mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
4f806387f71d3f0e8517650d5cd13b559003cdfd
ds4/ds4_cuda.cu
T
Salvatore Sanfilippo df6eedd972 Add optimized CUDA inference for Qwen3.8 Flash Next 
Share the Metal session graph with CUDA for text and vision, MTP and exact sampling, steering, checkpoints, and disk-only BF16 n-grams. Add bounded concurrent Linux n-gram reads, vectorized dense and expert kernels, register-held recurrent prefill state, and efficient small speculative batches.

Keep Qwen out of DeepSeek batch dispatch and exclude disk n-grams from resident memory accounting. Use normal file-backed weight loading in diagnostics and link n-gram support into the CUDA quality scorer.

Include independent kernel and memory-safety tests, mixed-prefill and recovery checks, long-context server and agent coverage, and user documentation with measured DGX Spark performance.
2026-09-15 15:00:11 +02:00

34040 lines
1.4 MiB
Plaintext
Raw Blame History

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <mma.h>
#include <cublas_v2.h>
#include <cub/block/block_radix_sort.cuh>
#include <stdint.h>
#include <errno.h>
#include <limits.h>
#include <float.h>
#include <math.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
#include <pthread.h>
#include <atomic>
#include "ds4_linux_memory.h"
#include <unordered_map>
#include <vector>
#include <algorithm>
#include "cuda/mmq/ds4_mmq.h"
#include "cuda/mmq/ds4_repack.h"
#include "ds4_image.h"
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
#define CUDA_QK_K 256
#define DS4_CUDA_UNUSED __attribute__((unused))
enum {
/* attention_decode_mixed_kernel stores raw-window scores plus visible
* compressed scores in shared memory. The host routes larger unmasked
* decode calls to the online attention kernel so this fixed buffer never
* becomes an out-of-bounds write at long context. */
DS4_CUDA_ATTENTION_SCORE_CAP = 8192u,
DS4_CUDA_ATTENTION_RAW_SCORE_CAP = 256u,
DS4_CUDA_TOPK_MERGE_GROUP = 8u,
/* perf-02 split-KV: fixed logical rows per chunk (shared scores = 2KB),
* grid S = ceil(n_score / CHUNK) clamped, so block count grows with ctx. */
DS4_CUDA_SPLITKV_CHUNK = 512u,
DS4_CUDA_SPLITKV_SCORE_CAP = 512u,
DS4_CUDA_SPLITKV_S_MAX = 16u,
DS4_CUDA_SPLITKV_S_FLOOR = 4u
};
/* struct ds4_gpu_tensor is defined in ds4_gpu.h (no longer opaque as of
* the device-aware CUDA PR). Field layout includes the new device_id
* tag and is read by the WITH_DEVICE-wrapped tensor APIs below. */
typedef struct {
uint8_t scales[CUDA_QK_K / 16];
uint8_t qs[CUDA_QK_K / 4];
uint16_t d;
uint16_t dmin;
} cuda_block_q2_K;
typedef struct {
uint16_t d;
uint16_t dmin;
uint8_t scales[12];
uint8_t qs[CUDA_QK_K / 2];
} cuda_block_q4_K;
typedef struct {
float d;
int8_t qs[CUDA_QK_K];
int16_t bsums[CUDA_QK_K / 16];
} cuda_block_q8_K;
typedef struct {
uint16_t d;
uint16_t qs[CUDA_QK_K / 8];
} cuda_block_iq2_xxs;
#include "ds4_gpu_mgpu.h"
#include "ds4_gpu_tp.h"
#include "ds4_iq2_tables_cuda.inc"
typedef struct {
ds4_gpu_attention_decode_row row[DS4_GPU_ATTENTION_DECODE_BATCH_MAX];
} cuda_attention_decode_row_table;
static_assert(sizeof(cuda_attention_decode_row_table) <= 3072u,
"attention row table must fit in CUDA kernel parameters");
static const void *g_model_host_base;
static const char *g_model_device_base;
static uint64_t g_model_registered_size;
static int g_model_registered;
static thread_local bool g_glm_mtp_verify_mode;
static int g_model_device_owned;
static int g_model_range_mapping_supported = 1;
static int g_model_hmm_direct;
static int g_model_fd = -1;
static const void *g_model_fd_host_base;
static int g_model_direct_fd = -1;
static uint64_t g_model_direct_align = 1;
static uint64_t g_model_file_size;
static int g_model_cache_full;
static cudaStream_t g_model_prefetch_stream;
static cudaStream_t g_model_upload_stream;
static int g_cublas_ready;
static int g_quality_mode;
static int g_decode_fast_attention;
static int g_decode_score_vec4;
static int g_xdev_sync_debug;
static int g_xdev_force_cuda_peer;
static int g_xdev_force_host_bounce;
static int g_cuda_disable_qkv_rms_fused;
static int g_cuda_no_window_attention;
static int g_cuda_decode_heads8_online;
static int g_cuda_decode_score4;
static int g_cuda_decode_score8;
static int g_cuda_no_decode_value512;
static int g_cuda_no_top1;
static int g_cuda_end_stream_sync;
static int g_cuda_no_setdevice_cache;
static int g_cuda_exact_score_split_graph;
static int g_cuda_exact_score_split_ldg;
static int g_cuda_exact_score_split_vec4;
static int g_cuda_exact_score_split_vec4_plain;
static int g_cuda_exact_score_split_dim2;
static int g_cuda_exact_score_split_fuse_inv_rope;
static int g_cuda_moe_decode_graph;
static int g_current_logical_tier = -1;
static int g_ssd_streaming_mode;
typedef struct ds4_gpu_stream_expert_table {
const void *model_map;
uint64_t model_size;
uint32_t layer, n_total_expert;
uint64_t gate_offset, up_offset, down_offset;
uint64_t gate_expert_bytes, down_expert_bytes;
} ds4_gpu_stream_expert_table;
extern "C" int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
const ds4_gpu_stream_expert_table *, const ds4_gpu_tensor *, uint32_t);
extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(uint64_t, uint64_t);
typedef struct {
int valid;
int logical_tier;
const void *model_map;
uint32_t layer;
uint32_t n_total_expert;
uint32_t slot_count;
uint32_t compact_count;
uint64_t gate_offset;
uint64_t up_offset;
uint64_t down_offset;
uint64_t gate_expert_bytes;
uint64_t down_expert_bytes;
char *gate_ptr;
char *up_ptr;
char *down_ptr;
uint64_t gate_capacity;
uint64_t up_capacity;
uint64_t down_capacity;
int32_t *slot_selected_ptr;
uint64_t slot_selected_capacity;
ds4_gpu_tensor slot_selected_tensor;
char *prefill_ptr;
uint64_t prefill_capacity;
} cuda_stream_selected_cache;
static cuda_stream_selected_cache g_stream_selected_cache;
static uint32_t g_stream_expert_budget;
static uint64_t g_stream_expert_bytes;
struct cuda_stream_expert_slot {
uint64_t gate, up, down, used;
};
static std::vector<cuda_stream_expert_slot> g_stream_expert_slots;
static std::unordered_map<uint64_t, uint32_t> g_stream_expert_by_gate;
/* Zero is empty; one is an unread look-ahead entry, older than any demand hit. */
static uint64_t g_stream_expert_clock = 1;
static std::vector<int32_t> g_stream_prefill_ids, g_stream_prefill_slots;
extern "C" void ds4_gpu_stream_expert_cache_prefetch_finish(bool cancel);
static void cuda_stream_prefetch_before_load(const ds4_gpu_stream_expert_table *table);
static bool cuda_stream_prefetch_protects(const cuda_stream_expert_slot &slot);
static int cuda_stream_compact_prefill(const char **gate, const char **up,
const char **down, const int32_t **ids,
uint32_t *experts, uint32_t in_dim = 0,
uint32_t mid_dim = 0, uint32_t out_dim = 0);
static void cuda_stream_selected_cache_invalidate(void) {
g_stream_selected_cache.valid = 0;
}
static void cuda_stream_selected_cache_release(void) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
const int tier = g_stream_selected_cache.logical_tier;
if (tier >= 0 && tier < g_n_gpus) {
(void)ds4_gpu_set_current_device(tier);
}
if (g_stream_selected_cache.gate_ptr) {
(void)cudaFree(g_stream_selected_cache.gate_ptr);
}
if (g_stream_selected_cache.up_ptr) {
(void)cudaFree(g_stream_selected_cache.up_ptr);
}
if (g_stream_selected_cache.down_ptr) {
(void)cudaFree(g_stream_selected_cache.down_ptr);
}
if (g_stream_selected_cache.slot_selected_ptr) {
(void)cudaFree(g_stream_selected_cache.slot_selected_ptr);
}
if (g_stream_selected_cache.prefill_ptr) {
(void)cudaFree(g_stream_selected_cache.prefill_ptr);
}
memset(&g_stream_selected_cache, 0, sizeof(g_stream_selected_cache));
g_stream_selected_cache.logical_tier = -1;
g_stream_expert_slots.clear();
g_stream_expert_by_gate.clear();
g_stream_expert_clock = 1;
g_stream_prefill_ids.clear();
g_stream_prefill_slots.clear();
}
typedef struct {
cudaGraph_t graph;
cudaGraphExec_t exec;
cudaGraphNode_t score_node;
cudaGraphNode_t final_node;
cudaGraphNode_t rope_node;
uint32_t n_head;
uint32_t head_dim;
uint32_t S;
uint32_t final_threads;
uint32_t n_rot;
int fuses_inv_rope;
int valid;
} cuda_score_split_graph_cache;
static cuda_score_split_graph_cache g_score_split_graph[DS4_MAX_GPUS];
static void attention_decode_score_split_graph_destroy_one(int logical_tier);
typedef struct {
cudaGraph_t graph;
cudaGraphExec_t exec;
cudaGraphNode_t xq_node;
cudaGraphNode_t gate_node;
cudaGraphNode_t midq_node;
cudaGraphNode_t down_node;
uint32_t n_expert;
uint32_t expert_in_dim;
uint32_t expert_mid_dim;
uint32_t out_dim;
int valid;
} cuda_moe_decode_graph_cache;
static cuda_moe_decode_graph_cache g_moe_decode_graph[DS4_MAX_GPUS];
static int cuda_q4_mma_ok(void);
static int cuda_q4_mma_tile16_shmem_ok(int which_down);
static void routed_moe_decode_graph_destroy_one(int logical_tier);
/* =========================================================================
* Multi-GPU plumbing (device-aware CUDA).
* ========================================================================= */
static_assert(DS4_MAX_GPUS == 16, "DS4_MAX_GPUS stack tables sized for 16");
ds4_gpu_ctx g_gpu[DS4_MAX_GPUS];
int g_n_gpus = 0;
int g_gpu_peer_ok[DS4_MAX_GPUS][DS4_MAX_GPUS];
static bool g_device_is_spark = false;
extern "C" int ds4_gpu_device_is_spark(void) {
return g_n_gpus == 1 && g_device_is_spark;
}
/* Per-pair pinned-host bounce buffers, indexed [src][dst]. Lazily grown
* to the largest copy seen for that pair. Each pair is its own allocation
* so concurrent fan-out copies from a single source GPU to multiple
* destinations cannot race for staging memory. */
static void *g_xdev_bounce[DS4_MAX_GPUS][DS4_MAX_GPUS];
static size_t g_xdev_bounce_bytes[DS4_MAX_GPUS][DS4_MAX_GPUS];
/* Internal helper: resolve a tensor's device index. -1 (untagged) is
* treated as device 0 for legacy callers. */
static inline int ds4_tensor_device_idx(const ds4_gpu_tensor *t) {
if (!t) return 0;
int d = t->device_id;
if (d < 0) return 0;
return d;
}
/* Debug/override flags are read once per CUDA init. The hot decode path calls
* the xdev helpers many times per token, so they must not re-enter getenv(). */
static void cuda_xdev_env_refresh(void) {
g_xdev_sync_debug = getenv("DS4_CUDA_SYNC_XDEV") != NULL;
g_xdev_force_cuda_peer = getenv("DS4_FORCE_CUDA_PEER") != NULL;
g_xdev_force_host_bounce = getenv("DS4_FORCE_HOST_BOUNCE") != NULL;
}
static void cuda_decode_dispatch_env_refresh(void) {
g_cuda_disable_qkv_rms_fused = getenv("DS4_CUDA_DISABLE_QKV_RMS_FUSED") != NULL;
g_cuda_no_window_attention = getenv("DS4_CUDA_NO_WINDOW_ATTENTION") != NULL;
g_cuda_decode_heads8_online = getenv("DS4_CUDA_DECODE_HEADS8_ONLINE") != NULL;
g_cuda_decode_score4 = getenv("DS4_CUDA_DECODE_SCORE4") != NULL;
g_cuda_decode_score8 = getenv("DS4_CUDA_DECODE_SCORE8") != NULL;
g_cuda_no_decode_value512 = getenv("DS4_CUDA_NO_DECODE_VALUE512") != NULL;
g_cuda_no_top1 = getenv("DS4_CUDA_NO_TOP1") != NULL;
g_cuda_end_stream_sync = getenv("DS4_CUDA_END_STREAM_SYNC") != NULL;
g_cuda_no_setdevice_cache = getenv("DS4_CUDA_NO_SETDEVICE_CACHE") != NULL;
g_cuda_exact_score_split_graph =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_GRAPH") != NULL;
g_cuda_exact_score_split_ldg =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_LDG") != NULL;
g_cuda_exact_score_split_vec4 =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4") != NULL;
g_cuda_exact_score_split_vec4_plain =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_VEC4_PLAIN") != NULL;
g_cuda_exact_score_split_dim2 =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_DIM2") != NULL &&
getenv("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DIM2") == NULL;
g_cuda_exact_score_split_fuse_inv_rope =
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_FUSE_INV_ROPE") != NULL;
g_cuda_moe_decode_graph = getenv("DS4_CUDA_MOE_DECODE_GRAPH") != NULL;
}
/* WITH_DEVICE(d) { ... } scope macro.
*
* Save the calling thread's current CUDA device, switch to device `d`,
* run the body exactly once, then restore the previous device. If the
* required CUDA calls fail, the body still runs (we don't have a clean
* way to early-exit a containing function from a macro), but the next
* CUDA call inside the body will surface the error naturally.
*
* Implementation: a for-loop with two synthetic variables. Iter 0 runs
* the body; on iter 1, the iteration step restores the previous device
* via cudaSetDevice and sets _wd_first = 0 so the loop exits. The
* single-statement-body restriction of for-loops is removed by the
* required `{ ... }` block in the call site.
*/
#define WITH_DEVICE(d) \
for (int _wd_prev = -1, _wd_first = 1; \
_wd_first; \
_wd_first = 0, \
(_wd_prev >= 0 ? (void)cudaSetDevice(_wd_prev) : (void)0)) \
if (cudaGetDevice(&_wd_prev) != cudaSuccess) { /* leave */ } else \
if (cudaSetDevice(d) != cudaSuccess) { /* leave */ } else
/* =========================================================================
* Per-device selective model cache (selective model cache).
*
* The public API in ds4_gpu.h declares ds4_tensor_range and the
* device_cache_tensors / lookup_cache entry points. ds4_cuda.cu does NOT
* include ds4_gpu.h historically (a pre-existing project convention), so
* we redeclare the struct here with the same layout the header uses.
* The implementation links by C linkage; struct compatibility is by
* field layout. */
typedef struct {
uint64_t source_offset;
uint64_t bytes;
int target_device;
} ds4_tensor_range;
struct cuda_device_cache {
void *base; /* device-side slab base */
size_t bytes;
int present;
};
static cuda_device_cache g_dev_cache[DS4_MAX_GPUS];
struct cache_range_entry {
uint64_t source_offset;
uint64_t bytes;
int device_id;
void *device_ptr;
};
static std::vector<cache_range_entry> g_cache_ranges;
struct cuda_model_range {
const void *host_base;
uint64_t offset;
uint64_t bytes;
char *device_ptr;
void *registered_base;
char *registered_device_base;
uint64_t registered_bytes;
int host_registered;
int arena_allocated;
int borrowed;
};
struct cuda_model_arena {
char *device_ptr;
uint64_t bytes;
uint64_t used;
};
struct cuda_q8_f16_range {
const void *host_base;
uint64_t offset;
uint64_t weight_bytes;
uint64_t in_dim;
uint64_t out_dim;
__half *device_ptr;
int device_id; /* physical CUDA device id; 0 in single-tier */
};
struct cuda_q8_f32_range {
const void *host_base;
uint64_t offset;
uint64_t weight_bytes;
uint64_t in_dim;
uint64_t out_dim;
float *device_ptr;
int device_id; /* physical CUDA device id; 0 in single-tier */
};
enum cuda_derived_kind {
CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE = 4,
CUDA_DERIVED_Q8_0_ALIGNED_DENSE = 5,
CUDA_DERIVED_Q2_K_ALIGNED_MOE = 6,
};
struct cuda_derived_range {
const void *host_base;
uint64_t source_offset;
uint64_t source_bytes;
uint32_t kind;
uint64_t in_dim;
uint64_t out_dim;
uint32_t group_count;
uint64_t bytes;
char *device_ptr;
};
static std::vector<cuda_model_range> g_model_ranges;
static std::vector<cuda_model_arena> g_model_arenas;
static std::unordered_map<uint64_t, size_t> g_model_range_by_offset;
static std::vector<cuda_q8_f16_range> g_q8_f16_ranges;
static std::unordered_map<uint64_t, size_t> g_q8_f16_by_offset;
static std::vector<cuda_q8_f32_range> g_q8_f32_ranges;
static std::unordered_map<uint64_t, size_t> g_q8_f32_by_offset;
static std::vector<cuda_derived_range> g_derived_ranges;
static const void *g_derived_replace_map;
static uint64_t g_derived_artifact_bytes;
static double g_derived_artifact_build_secs;
static int g_derived_replaces_complete;
static void *g_aligned_q81_scratch;
static uint64_t g_model_range_bytes;
static uint64_t g_q8_f16_bytes;
static uint64_t g_q8_f32_bytes;
static int g_q8_cache_suppressed;
static int g_q8_f16_disabled_after_oom;
static int g_q8_f16_budget_notice_printed;
static uint64_t g_model_load_progress_next;
static double g_model_load_progress_last;
static int g_model_load_progress_started;
static int g_model_load_progress_tty;
static void *g_cuda_tmp;
static uint64_t g_cuda_tmp_bytes;
static void *g_tt_scratch;
static uint64_t g_tt_scratch_bytes;
static int g_tt_scratch_device = -1;
#ifdef DS4_CUDA_HAVE_MXF4
static void *g_indexer_mxf4_scratch;
static uint64_t g_indexer_mxf4_scratch_bytes;
static int g_indexer_mxf4_scratch_device = -1;
#endif
static void *g_model_stage_raw[4];
static void *g_model_stage[4];
static cudaEvent_t g_model_stage_event[4];
static uint64_t g_model_stage_bytes;
static void *g_stream_selected_stage_raw[4];
static void *g_stream_selected_stage[4];
static cudaEvent_t g_stream_selected_stage_event[4];
static uint64_t g_stream_selected_stage_bytes;
static cudaStream_t g_stream_selected_upload_stream;
static int cuda_ok(cudaError_t err, const char *what);
extern "C" void ds4_gpu_decode_graphs_invalidate(void);
static const char *cuda_model_range_ptr_from_fd(
const void *model_map,
uint64_t offset,
uint64_t bytes,
const char *what);
/* Forward declaration: defined later in this file. The resolver wrapper
* below uses it for multi-tier dispatch. */
extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset,
uint64_t bytes,
int expected_device,
void **out_device_ptr);
__global__ static void dequant_q8_0_to_f16_kernel(
__half *out,
const unsigned char *w,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks);
__global__ static void dequant_q8_0_to_f32_kernel(
float *out,
const unsigned char *w,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks);
static int cuda_aligned_iq2_enabled(void) {
const char *s = getenv("DS4_CUDA_MOE_NO_IQ2_ALIGNED");
return !(s && s[0] && strcmp(s, "0") != 0);
}
static int cuda_aligned_q2k_enabled(void) {
const char *s = getenv("DS4_CUDA_MOE_NO_Q2K_ALIGNED");
return !(s && s[0] && strcmp(s, "0") != 0);
}
static int cuda_aligned_q8_enabled(void) {
const char *s = getenv("DS4_CUDA_Q8_NO_ALIGNED");
return !(s && s[0] && strcmp(s, "0") != 0);
}
static const char *cuda_derived_weight_ptr(
const void *model_map,
uint64_t source_offset,
uint64_t source_bytes,
uint32_t kind,
uint64_t in_dim,
uint64_t out_dim,
uint32_t group_count,
uint64_t bytes) {
if (getenv("DS4_CUDA_NO_DERIVED_WEIGHTS") != NULL) return NULL;
for (const cuda_derived_range &r : g_derived_ranges) {
if (r.host_base == model_map &&
r.source_offset == source_offset &&
r.source_bytes == source_bytes &&
r.kind == kind &&
r.in_dim == in_dim &&
r.out_dim == out_dim &&
r.group_count == group_count &&
bytes <= r.bytes) {
return r.device_ptr;
}
}
return NULL;
}
static int cuda_model_map_replaces_complete(const void *model_map) {
return g_derived_replaces_complete && model_map == g_derived_replace_map;
}
static int cuda_integrated_artifact_map(const void *model_map) {
if (!cuda_model_map_replaces_complete(model_map)) return 0;
int device = 0;
int integrated = 0;
return cudaGetDevice(&device) == cudaSuccess &&
cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated, device) == cudaSuccess &&
integrated;
}
/* Startup combines adjacent tensors into cache spans. A span is replaceable
* when aligned IQ2/Q2K artifacts cover every tensor payload in it; GGUF
* alignment gaps of at most 64 KiB do not need device residency. */
static int cuda_span_fully_replaced(
const void *model_map,
uint64_t offset,
uint64_t bytes) {
if (!cuda_model_map_replaces_complete(model_map) || bytes == 0) return 0;
const uint64_t end = offset + bytes;
if (end < offset) return 0;
uint64_t cursor = offset;
int found = 0;
while (cursor < end) {
uint64_t next = cursor;
for (const cuda_derived_range &r : g_derived_ranges) {
if (r.host_base != model_map ||
(r.kind != CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE &&
r.kind != CUDA_DERIVED_Q2_K_ALIGNED_MOE)) {
continue;
}
const uint64_t r_end = r.source_offset + r.source_bytes;
const uint64_t gap = found ? 65536u : 0u;
if (r_end >= r.source_offset &&
r.source_offset <= cursor + gap &&
r_end > next) {
next = r_end;
}
}
if (next <= cursor) return 0;
cursor = next;
found = 1;
}
return 1;
}
static void *cuda_tmp_alloc(uint64_t bytes, const char *what) {
if (bytes == 0) return NULL;
if (g_cuda_tmp_bytes >= bytes) return g_cuda_tmp;
if (g_cuda_tmp) {
if (!cuda_ok(cudaDeviceSynchronize(),
"synchronize CUDA scratch growth")) {
return NULL;
}
ds4_gpu_decode_graphs_invalidate();
(void)cudaFree(g_cuda_tmp);
g_cuda_tmp = NULL;
g_cuda_tmp_bytes = 0;
}
void *ptr = NULL;
cudaError_t err = cudaMalloc(&ptr, (size_t)bytes);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA temp alloc failed for %s (%.2f MiB): %s\n",
what ? what : "scratch", (double)bytes / 1048576.0, cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
g_cuda_tmp = ptr;
g_cuda_tmp_bytes = bytes;
return g_cuda_tmp;
}
static void *tt_scratch_ensure(uint64_t bytes, const char *what) {
if (bytes == 0) return NULL;
if (g_tt_scratch_bytes >= bytes) return g_tt_scratch;
int device = -1;
if (!cuda_ok(cudaGetDevice(&device), "get device for token-tile scratch")) {
return NULL;
}
if (g_tt_scratch) {
if (!cuda_ok(cudaDeviceSynchronize(),
"synchronize token-tile scratch growth")) {
return NULL;
}
(void)cudaFree(g_tt_scratch);
g_tt_scratch = NULL;
g_tt_scratch_bytes = 0;
g_tt_scratch_device = -1;
}
if (!cuda_ok(cudaMalloc(&g_tt_scratch, (size_t)bytes),
what ? what : "allocate token-tile scratch")) {
g_tt_scratch = NULL;
return NULL;
}
g_tt_scratch_bytes = bytes;
g_tt_scratch_device = device;
return g_tt_scratch;
}
static inline uint64_t tt_align256_u64(uint64_t x) {
return (x + 255ull) & ~255ull;
}
/* Per-tier scratch accessor.
*
* Behavior:
* - At g_n_gpus <= 1 (single-tier), delegates to cuda_tmp_alloc which
* manages the legacy g_cuda_tmp slab. This guarantees byte-identical
* behavior to pre-task code for the gpu_cfg == NULL case.
* - For multi-tier, grows the per-device g_gpu[logical_tier].scratch
* slab on the corresponding physical device. Cleanup is already
* handled by ds4_gpu_cleanup (which walks g_gpu[i].scratch).
*
* The legacy g_cuda_tmp slab is untouched: init-time / preload callers
* still use cuda_tmp_alloc directly. No aliasing between g_cuda_tmp
* and g_gpu[0].scratch — they are independently owned and freed.
*
* Added for multi-GPU execution (multi-GPU execution), step A3 of the
* spec (sub-area 2). */
/* Scalar Q8 shared-expert work may run beside the routed experts. Its
* activation quantization must not reuse the main stream's scratch. */
static struct {
cudaStream_t stream;
cudaEvent_t ready, done;
void *scratch;
bool active, pending, disabled;
} g_dsv41_shared;
static const uint64_t CUDA_DSV41_SHARED_SCRATCH = 65536u;
static void *cuda_tmp_alloc_on(int logical_tier, uint64_t bytes, const char *what) {
if (g_dsv41_shared.active)
return logical_tier == 0 && bytes <= CUDA_DSV41_SHARED_SCRATCH
? g_dsv41_shared.scratch : NULL;
if (bytes == 0) return NULL;
if (g_n_gpus <= 1) {
return cuda_tmp_alloc(bytes, what);
}
if (logical_tier < 0 || logical_tier >= g_n_gpus) {
fprintf(stderr, "ds4: cuda_tmp_alloc_on: bad tier %d (n_gpus=%d, what=%s)\n",
logical_tier, g_n_gpus, what ? what : "?");
return NULL;
}
ds4_gpu_ctx *ctx = &g_gpu[logical_tier];
if (ctx->scratch_bytes >= bytes) return ctx->scratch;
int prev = -1;
cudaError_t derr = cudaGetDevice(&prev);
if (derr != cudaSuccess) {
fprintf(stderr,
"ds4: cudaGetDevice failed before scratch alloc on tier %d (dev=%d, what=%s): %s\n",
logical_tier, ctx->device_id, what ? what : "scratch",
cudaGetErrorString(derr));
(void)cudaGetLastError();
return NULL;
}
derr = cudaSetDevice(ctx->device_id);
if (derr != cudaSuccess) {
fprintf(stderr,
"ds4: cudaSetDevice(%d) failed before scratch alloc on tier %d (what=%s): %s\n",
ctx->device_id, logical_tier, what ? what : "scratch",
cudaGetErrorString(derr));
(void)cudaGetLastError();
if (prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
if (ctx->scratch) {
(void)cudaFree(ctx->scratch);
ctx->scratch = NULL;
ctx->scratch_bytes = 0;
}
void *p = NULL;
cudaError_t err = cudaMalloc(&p, (size_t)bytes);
if (prev >= 0) (void)cudaSetDevice(prev);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA scratch alloc on tier %d (dev=%d) failed for %s (%.2f MiB): %s\n",
logical_tier, ctx->device_id, what ? what : "scratch",
(double)bytes / 1048576.0, cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
ctx->scratch = p;
ctx->scratch_bytes = (size_t)bytes;
return p;
}
static int cuda_attention_score_buffer_fits(uint32_t n_comp) {
return n_comp <= DS4_CUDA_ATTENTION_SCORE_CAP - DS4_CUDA_ATTENTION_RAW_SCORE_CAP;
}
static const char *cuda_model_ptr(const void *model_map, uint64_t offset) {
if (model_map == g_model_host_base && g_model_device_base) return g_model_device_base + offset;
return (const char *)model_map + offset;
}
static const char *cuda_model_range_ptr(const void *model_map, uint64_t offset, uint64_t bytes, const char *what) {
if (bytes == 0) return cuda_model_ptr(model_map, offset);
const uint64_t end = offset + bytes;
if (end < offset) return NULL;
auto exact = g_model_range_by_offset.find(offset);
if (exact != g_model_range_by_offset.end()) {
const cuda_model_range &r = g_model_ranges[exact->second];
if (r.host_base == model_map && bytes <= r.bytes) return r.device_ptr;
}
for (const cuda_model_range &r : g_model_ranges) {
if (r.host_base == model_map && offset >= r.offset &&
end <= r.offset + r.bytes) {
return r.device_ptr + (offset - r.offset);
}
if (r.host_base == model_map && r.host_registered &&
r.registered_base && r.registered_device_base) {
const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset);
const uintptr_t h1 = h0 + bytes;
const uintptr_t r0 = (uintptr_t)r.registered_base;
const uintptr_t r1 = r0 + r.registered_bytes;
if (h1 >= h0 && h0 >= r0 && h1 <= r1) {
return r.registered_device_base + (h0 - r0);
}
}
}
if (model_map == g_model_host_base &&
(g_model_device_owned || g_model_registered)) {
return cuda_model_ptr(model_map, offset);
}
if (model_map == g_model_host_base && g_model_hmm_direct &&
getenv("DS4_CUDA_WEIGHT_CACHE") == NULL &&
getenv("DS4_CUDA_WEIGHT_PRELOAD") == NULL) {
return cuda_model_ptr(model_map, offset);
}
const char *direct_env = getenv("DS4_CUDA_DIRECT_MODEL");
if (model_map == g_model_host_base && direct_env && direct_env[0]) {
return cuda_model_ptr(model_map, offset);
}
if (getenv("DS4_CUDA_NO_FD_CACHE") == NULL) {
const char *fd_ptr = cuda_model_range_ptr_from_fd(model_map, offset, bytes, what);
if (fd_ptr || (g_ssd_streaming_mode && g_model_fd >= 0 && model_map == g_model_fd_host_base))
return fd_ptr;
}
cudaError_t err = cudaSuccess;
if (g_model_range_mapping_supported) {
const long page_sz_l = sysconf(_SC_PAGESIZE);
const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u;
const uintptr_t host_addr = (uintptr_t)((const char *)model_map + offset);
const uintptr_t reg_addr = host_addr & ~(uintptr_t)(page_sz - 1u);
const uint64_t reg_delta = (uint64_t)(host_addr - reg_addr);
const uint64_t reg_bytes = (reg_delta + bytes + page_sz - 1u) & ~(page_sz - 1u);
void *reg_dev = NULL;
err = cudaHostRegister((void *)reg_addr,
(size_t)reg_bytes,
cudaHostRegisterMapped | cudaHostRegisterReadOnly);
if (err == cudaSuccess) {
err = cudaHostGetDevicePointer(&reg_dev, (void *)reg_addr, 0);
if (err == cudaSuccess && reg_dev) {
char *dev_ptr = (char *)reg_dev + reg_delta;
g_model_ranges.push_back({model_map, offset, bytes, dev_ptr, (void *)reg_addr, (char *)reg_dev, reg_bytes, 1, 0});
g_model_range_by_offset[offset] = g_model_ranges.size() - 1u;
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA mapped %s %.2f MiB\n",
what ? what : "weights",
(double)bytes / 1048576.0);
}
return dev_ptr;
}
fprintf(stderr, "ds4: CUDA model range map pointer failed for %s: %s\n",
what ? what : "weights", cudaGetErrorString(err));
(void)cudaHostUnregister((void *)reg_addr);
(void)cudaGetLastError();
} else {
if (err == cudaErrorNotSupported || err == cudaErrorInvalidValue) g_model_range_mapping_supported = 0;
(void)cudaGetLastError();
}
}
void *dev = NULL;
err = cudaMalloc(&dev, (size_t)bytes);
if (err != cudaSuccess) {
(void)cudaGetLastError();
fprintf(stderr, "ds4: CUDA model range alloc failed for %s (%.2f MiB): %s\n",
what ? what : "weights", (double)bytes / 1048576.0, cudaGetErrorString(err));
return NULL;
}
const char *src = (const char *)model_map + offset;
const uint64_t chunk = 64ull * 1024ull * 1024ull;
for (uint64_t done = 0; done < bytes; done += chunk) {
uint64_t n = bytes - done < chunk ? bytes - done : chunk;
err = cudaMemcpy((char *)dev + done, src + done, (size_t)n, cudaMemcpyHostToDevice);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f/%.2f MiB: %s\n",
what ? what : "weights",
(double)done / 1048576.0,
(double)bytes / 1048576.0,
cudaGetErrorString(err));
(void)cudaFree(dev);
(void)cudaGetLastError();
return NULL;
}
}
g_model_ranges.push_back({model_map, offset, bytes, (char *)dev, NULL, NULL, 0, 0, 0});
g_model_range_by_offset[offset] = g_model_ranges.size() - 1u;
g_model_range_bytes += bytes;
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA cached %s %.2f MiB (total %.2f GiB)\n",
what ? what : "weights",
(double)bytes / 1048576.0,
(double)g_model_range_bytes / 1073741824.0);
}
return (const char *)dev;
}
/* Per-tier cuBLAS handle. Used by kernel-dispatch wrappers; returns the
* cuBLAS handle for the logical tier. The wrapper is expected to have
* cudaSetDevice'd to that tier's physical device already (kernels and
* cuBLAS calls ride the default stream and are naturally serialized).
*
* Added for multi-GPU execution (multi-GPU execution), sub-area 1. */
static inline cublasHandle_t cuda_cublas_for_tier(int logical_tier) {
if (g_n_gpus <= 1) {
return (cublasHandle_t)g_gpu[0].cublas;
}
/* The executing device is authoritative: GLM per-layer switching runs
* generic launchers whose out tensors live on device 0 while the layer
* executes elsewhere. On DS4 paths the current device always equals the
* requested tier, so this is behavior-preserving there. */
int cur_dev = -1;
if (cudaGetDevice(&cur_dev) == cudaSuccess) {
for (int t = 0; t < g_n_gpus; t++) {
if (g_gpu[t].device_id == cur_dev) {
return (cublasHandle_t)g_gpu[t].cublas;
}
}
}
if (logical_tier < 0 || logical_tier >= g_n_gpus) {
return (cublasHandle_t)g_gpu[0].cublas;
}
return (cublasHandle_t)g_gpu[logical_tier].cublas;
}
/* ------------------------------------------------------------------------
* Decode-island CUDA graph capture.
*
* Design ported from the Entrpi/ds4 batched-serving fork's per-layer
* decode graph work (fork commits abe26575, 2200c675, 9caaa508; author
* Entrpi <entrpi@proton.me>), re-implemented against upstream's decode
* body. Phase 1 captures the two position-independent islands of each
* decode layer:
* island 0: layer top (hc pre-norm, mixes, QKV projections), ending
* before the position-dependent QKV rope.
* island 1: attention output projection through the FFN/MoE tail.
* Network V4.1 TP separates the output projection into island 2 because
* its attention all-reduce sits before the FFN in island 1.
* Kernel launches inside the islands ride cuda_decode_stream(): the
* legacy NULL stream in eager mode (bit-identical behavior to before),
* or a dedicated blocking stream while capturing/replaying. The legacy
* stream cannot be captured; a blocking (default-flags) stream keeps the
* implicit serialization against NULL-stream work, so the eager middle
* section of each layer orders correctly around replays.
*
* Per-entry state machine: first sight of a key runs eager (warms lazy
* allocators -- cudaMalloc is forbidden inside capture), second sight
* captures + instantiates + launches, later sights replay. Any failure
* marks the entry dead and the caller re-encodes the island eagerly, so
* a broken capture can never lose a token's work. Graphs replay the
* captured kernels byte-for-byte, so replayed output is bit-identical
* to the eager encode of the same island.
*
* DS4_CUDA_DECODE_GRAPHS=0 (or off/no/false) disables everything. */
#define CUDA_DECODE_GRAPH_LAYERS 64u
#define CUDA_DECODE_GRAPH_ISLANDS 3u
#define CUDA_DECODE_GRAPH_VARIANTS 4u
/* Mirrors the public `struct ds4_decode_graph_key` decl in ds4_gpu.h
* byte-for-byte (ds4_cuda.cu does not include that header; it carries
* its own extern "C" signatures inline). The static_assert catches
* accidental drift: adding a field means updating both sites. */
struct ds4_decode_graph_key {
uint32_t il;
uint32_t island;
uint32_t variant;
uint32_t _pad;
void *cur_hc;
void *after_attn_hc;
void *after_ffn_hc;
void *attn_norm;
};
typedef struct ds4_decode_graph_key ds4_decode_graph_key;
static_assert(sizeof(ds4_decode_graph_key) == 48u,
"ds4_decode_graph_key must match the ds4_gpu.h decl");
struct cuda_decode_graph_entry {
ds4_decode_graph_key key;
cudaGraphExec_t exec;
int state; /* 0 empty, 1 warmed, 2 ready, 3 dead */
uint64_t hits;
};
static cuda_decode_graph_entry
g_decode_graphs[CUDA_DECODE_GRAPH_LAYERS][CUDA_DECODE_GRAPH_ISLANDS]
[CUDA_DECODE_GRAPH_VARIANTS];
static cudaStream_t g_decode_graph_stream = NULL;
static int g_decode_graph_capturing = 0;
static uint64_t g_decode_graph_replays = 0;
static uint64_t g_decode_graph_captures = 0;
extern "C" int ds4_gpu_decode_graphs_supported(void) {
static int init = 0;
static int enabled = 0;
if (!init) {
init = 1;
const char *s = getenv("DS4_CUDA_DECODE_GRAPHS");
const int off = s && *s &&
(s[0] == '0' ||
strcmp(s, "off") == 0 || strcmp(s, "OFF") == 0 ||
strcmp(s, "no") == 0 || strcmp(s, "NO") == 0 ||
strcmp(s, "false") == 0 || strcmp(s, "FALSE") == 0);
if (off) {
fprintf(stderr, "ds4: DS4_CUDA_DECODE_GRAPHS=%s - decode graph capture disabled\n", s);
enabled = 0;
} else {
enabled = 1;
}
}
return enabled && g_n_gpus == 1;
}
/* Stream the decode-island kernels launch on. Legacy NULL stream in
* eager mode (unchanged behavior); the capture stream while a capture
* or replay is in flight. */
static inline cudaStream_t cuda_decode_stream(void) {
if (g_dsv41_shared.active) return g_dsv41_shared.stream;
return g_decode_graph_capturing ? g_decode_graph_stream : (cudaStream_t)0;
}
static void cuda_dsv41_shared_free(void) {
if (g_dsv41_shared.stream) (void)cudaStreamDestroy(g_dsv41_shared.stream);
if (g_dsv41_shared.ready) (void)cudaEventDestroy(g_dsv41_shared.ready);
if (g_dsv41_shared.done) (void)cudaEventDestroy(g_dsv41_shared.done);
if (g_dsv41_shared.scratch) (void)cudaFree(g_dsv41_shared.scratch);
g_dsv41_shared = {};
}
static bool cuda_dsv41_shared_prepare(void) {
if (g_dsv41_shared.stream) return true;
if (g_dsv41_shared.disabled || g_decode_graph_capturing) return false;
if (cuda_ok(cudaStreamCreateWithFlags(&g_dsv41_shared.stream, cudaStreamNonBlocking),
"shared expert stream") &&
cuda_ok(cudaEventCreateWithFlags(&g_dsv41_shared.ready, cudaEventDisableTiming),
"shared expert input event") &&
cuda_ok(cudaEventCreateWithFlags(&g_dsv41_shared.done, cudaEventDisableTiming),
"shared expert output event") &&
cuda_ok(cudaMalloc(&g_dsv41_shared.scratch, CUDA_DSV41_SHARED_SCRATCH),
"shared expert scratch")) return true;
cuda_dsv41_shared_free();
(void)cudaGetLastError();
g_dsv41_shared.disabled = true;
return false;
}
static void cuda_decode_graph_entry_kill(cuda_decode_graph_entry *e) {
if (e->exec) {
(void)cudaGraphExecDestroy(e->exec);
e->exec = NULL;
}
e->state = 3;
}
extern "C" void ds4_gpu_decode_graphs_invalidate(void) {
for (uint32_t il = 0; il < CUDA_DECODE_GRAPH_LAYERS; il++) {
for (uint32_t is = 0; is < CUDA_DECODE_GRAPH_ISLANDS; is++) {
for (uint32_t v = 0; v < CUDA_DECODE_GRAPH_VARIANTS; v++) {
cuda_decode_graph_entry *e = &g_decode_graphs[il][is][v];
if (e->exec) {
(void)cudaGraphExecDestroy(e->exec);
e->exec = NULL;
}
e->state = 0;
e->hits = 0;
memset(&e->key, 0, sizeof(e->key));
}
}
}
}
static cuda_decode_graph_entry *cuda_decode_graph_find(
const ds4_decode_graph_key *key) {
if (key->il >= CUDA_DECODE_GRAPH_LAYERS ||
key->island >= CUDA_DECODE_GRAPH_ISLANDS) return NULL;
cuda_decode_graph_entry *slot = NULL;
for (uint32_t v = 0; v < CUDA_DECODE_GRAPH_VARIANTS; v++) {
cuda_decode_graph_entry *e = &g_decode_graphs[key->il][key->island][v];
if (e->state != 0 &&
memcmp(&e->key, key, sizeof(*key)) == 0) return e;
if (e->state == 0 && !slot) slot = e;
}
if (slot) {
memcpy(&slot->key, key, sizeof(*key));
slot->state = 0; /* caller advances the state machine */
return slot;
}
return NULL; /* all variants busy with other keys: stay eager */
}
extern "C" int ds4_gpu_decode_graph_begin(const ds4_decode_graph_key *key) {
if (!key || !ds4_gpu_decode_graphs_supported()) return -1;
if (g_decode_graph_capturing) return -1; /* no nesting */
cuda_decode_graph_entry *e = cuda_decode_graph_find(key);
if (!e || e->state == 3) return -1;
if (e->state == 0) {
/* Warm pass: run eagerly once so lazy allocators (tmp scratch,
* cuBLAS workspaces) reach steady-state sizes before capture. */
e->state = 1;
return -1;
}
if (e->state == 2) {
cudaError_t err = cudaGraphLaunch(e->exec, g_decode_graph_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: decode graph replay failed (il=%u island=%u): %s\n",
key->il, key->island, cudaGetErrorString(err));
(void)cudaGetLastError();
cuda_decode_graph_entry_kill(e);
return -1; /* caller encodes eagerly; nothing was consumed */
}
e->hits++;
g_decode_graph_replays++;
return 1;
}
/* state == 1: capture this encode. */
if (!g_decode_graph_stream) {
if (!cuda_ok(cudaStreamCreate(&g_decode_graph_stream),
"decode graph stream create")) {
g_decode_graph_stream = NULL;
cuda_decode_graph_entry_kill(e);
return -1;
}
}
/* cuBLAS rides the legacy NULL stream, which cannot be captured:
* point the handle at the capture stream for the duration. */
(void)cublasSetStream(cuda_cublas_for_tier(0), g_decode_graph_stream);
if (!cuda_ok(cudaStreamBeginCapture(g_decode_graph_stream,
cudaStreamCaptureModeGlobal),
"decode graph begin capture")) {
(void)cublasSetStream(cuda_cublas_for_tier(0), NULL);
cuda_decode_graph_entry_kill(e);
return -1;
}
g_decode_graph_capturing = 1;
return 0;
}
extern "C" int ds4_gpu_decode_graph_end(const ds4_decode_graph_key *key) {
if (!key || !g_decode_graph_capturing) return -1;
g_decode_graph_capturing = 0;
cudaGraph_t graph = NULL;
cudaError_t err = cudaStreamEndCapture(g_decode_graph_stream, &graph);
(void)cublasSetStream(cuda_cublas_for_tier(0), NULL);
cuda_decode_graph_entry *e = cuda_decode_graph_find(key);
if (err != cudaSuccess || graph == NULL) {
fprintf(stderr, "ds4: decode graph capture failed (il=%u island=%u): %s\n",
key->il, key->island, cudaGetErrorString(err));
(void)cudaGetLastError();
if (graph) (void)cudaGraphDestroy(graph);
if (e) cuda_decode_graph_entry_kill(e);
return -1; /* caller re-encodes the island eagerly */
}
if (!e) { /* cannot happen: begin() found it */
(void)cudaGraphDestroy(graph);
return -1;
}
cudaGraphExec_t exec = NULL;
err = cudaGraphInstantiate(&exec, graph, NULL, NULL, 0);
(void)cudaGraphDestroy(graph);
if (err != cudaSuccess || exec == NULL) {
fprintf(stderr, "ds4: decode graph instantiate failed (il=%u island=%u): %s\n",
key->il, key->island, cudaGetErrorString(err));
(void)cudaGetLastError();
cuda_decode_graph_entry_kill(e);
return -1;
}
/* Capture recorded the work without executing it: launch now so this
* token's island actually runs. */
err = cudaGraphLaunch(exec, g_decode_graph_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: decode graph first launch failed (il=%u island=%u): %s\n",
key->il, key->island, cudaGetErrorString(err));
(void)cudaGetLastError();
(void)cudaGraphExecDestroy(exec);
cuda_decode_graph_entry_kill(e);
return -1;
}
e->exec = exec;
e->state = 2;
g_decode_graph_captures++;
if (getenv("DS4_CUDA_DECODE_GRAPH_LOG") != NULL) {
fprintf(stderr, "ds4: decode graph captured il=%u island=%u (total %llu)\n",
key->il, key->island,
(unsigned long long)g_decode_graph_captures);
}
return 0;
}
/* ------------------------------------------------------------------------
* Vendored llama.cpp mmq prefill tier (cuda/mmq/, see cuda/mmq/VENDOR.md).
*
* Ported from the Entrpi/ds4 batched-serving fork (fork commits 39d3877c,
* a56e07a5, 944482d5 and the Phase 5/6 MoE wiring; author Entrpi
* <entrpi@proton.me>). Phase 1 takes only the RAW-layout entries: dense
* Q8_0 GEMMs and the IQ2_XXS gate/up + Q2_K down routed-MoE pipeline for
* n_tok >= 2 (prefill and batched verify). The aligned-SoA / D2R /
* producer-quantized tiers and the weight server are deliberately left
* behind: they need the repack artifact pipeline for a further ~25% on
* top of the ~2.5x this tier gives. Decode (n_tok == 1) is untouched.
* mmq changes FP32 reduction order vs the cublas+dequant pipeline, so
* prefill logits drift at ULP scale: validated against the official
* continuation vectors rather than byte-diffs. DS4_CUDA_MMQ=0 restores
* the legacy dispatch. */
/* Producer-fold registry lookup the vendored ds4_mmq.cu entries probe for
* pre-quantized q8_1 activations (the fork's flat-pool M2-Inc2a fold).
* The flat-pool producers are not ported, so nothing is ever registered:
* always miss, and the entries run their own activation quantize. */
extern "C" int ds4_cuda_q8_fold_take_q81(const void *src, uint64_t in_dim,
const void **q81) {
(void)src; (void)in_dim; (void)q81;
return 0;
}
static int cuda_use_mmq(void) {
static int init = 0;
static int use = 0;
if (!init) {
init = 1;
const char *s = getenv("DS4_CUDA_MMQ");
const int off = (s && s[0] == '0') || g_quality_mode || g_n_gpus > 1;
if (off) {
if (s && s[0] == '0') {
fprintf(stderr, "ds4: DS4_CUDA_MMQ=0 - mmq prefill tier disabled\n");
}
} else if (ds4_mmq_init(0) == 0) {
use = 1;
} else {
fprintf(stderr, "ds4: ds4_mmq_init failed - mmq prefill tier disabled\n");
}
}
return use;
}
/* MXFP4 has no dequant+cublas fallback, so it must retain MMQ on multi-GPU
* placements where the optional Q8/IQ2 prefill tier stays disabled. MMQ
* resolves the active CUDA device on every call; initialization only warms
* its device-info singleton. The experimental global persistent scratch is
* intentionally rejected because one pointer cannot span CUDA devices. */
static int cuda_use_mxfp4_mmq(void) {
static int init = 0;
static int use = 0;
if (!init) {
init = 1;
const char *s = getenv("DS4_CUDA_MMQ");
if (s && s[0] == '0') {
fprintf(stderr,
"ds4: DS4_CUDA_MMQ=0 - MXFP4 MMQ disabled\n");
} else if (g_n_gpus > 1 &&
getenv("DS4_CUDA_MMQ_Q81_PERSISTENT") != NULL) {
fprintf(stderr,
"ds4: persistent MMQ Q8_1 scratch is unavailable with multi-GPU MXFP4\n");
} else {
int device = 0;
if (cudaGetDevice(&device) == cudaSuccess &&
ds4_mmq_init(device) == 0) {
use = 1;
} else {
fprintf(stderr,
"ds4: ds4_mmq_init failed - MXFP4 unavailable\n");
}
}
}
return use;
}
/* mmq pipeline helpers, ported from the fork: SwiGLU + clamp + router
* weight in the (token, slot, feature) layout ds4_mmq_*_moe writes, and
* the guarded per-token slot sum. */
__global__ static void moe_mmq_swiglu_weighted_clamp_kernel(
float *mid_out,
const float *gate_buf, const float *up_buf,
const float *weights,
uint32_t expert_mid_dim,
uint32_t n_tokens,
uint32_t n_expert_used,
float clamp) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_expert_used * expert_mid_dim;
if (gid >= n) return;
uint64_t slot_pair = gid / expert_mid_dim;
uint32_t tok = (uint32_t)(slot_pair / n_expert_used);
uint32_t slot = (uint32_t)(slot_pair - (uint64_t)tok * n_expert_used);
const float w = weights[(uint64_t)tok * n_expert_used + slot];
/* Unowned assignments have no gate/up output. Never read those slots. */
if (w == 0.0f) { mid_out[gid] = 0.0f; return; }
float g = gate_buf[gid];
float u = up_buf[gid];
if (!isfinite(g)) g = 0.0f;
if (!isfinite(u)) u = 0.0f;
if (clamp > 1.0e-6f) {
if (g > clamp) g = clamp;
if (u > clamp) u = clamp;
if (u < -clamp) u = -clamp;
}
const float s = g / (1.0f + expf(-g));
mid_out[gid] = s * u * w;
}
__global__ static void moe_mmq_sum_kernel(float *out, const float *down,
const int32_t *selected, uint32_t out_dim, uint32_t n_expert,
uint32_t n_tokens, uint32_t guard_nonfinite) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * out_dim;
if (gid >= n) return;
uint32_t tok = gid / out_dim;
uint32_t row = gid - (uint64_t)tok * out_dim;
float acc = 0.0f;
for (uint32_t e = 0; e < n_expert; e++) {
if (selected && selected[(uint64_t)tok * n_expert + e] < 0) continue;
const float v = down[((uint64_t)tok * n_expert + e) * out_dim + row];
if (!guard_nonfinite || isfinite(v)) acc += v;
}
out[gid] = acc;
}
/* Abort an in-flight capture after an encode error inside the island.
* Nothing was executed; the caller re-encodes eagerly. */
extern "C" void ds4_gpu_decode_graph_abort(const ds4_decode_graph_key *key) {
if (!g_decode_graph_capturing) return;
g_decode_graph_capturing = 0;
cudaGraph_t graph = NULL;
(void)cudaStreamEndCapture(g_decode_graph_stream, &graph);
if (graph) (void)cudaGraphDestroy(graph);
(void)cublasSetStream(cuda_cublas_for_tier(0), NULL);
(void)cudaGetLastError();
if (key) {
cuda_decode_graph_entry *e = cuda_decode_graph_find(key);
if (e) cuda_decode_graph_entry_kill(e);
}
}
/* Multi-tier-aware weight pointer resolver.
*
* Used by kernel-dispatch wrappers in the per-layer execution path to
* obtain a device pointer for a weight slice on the layer's logical
* tier. Behavior:
*
* - When g_n_gpus <= 1 (single-tier engine), delegates to the existing
* cuda_model_range_ptr path. This short-circuit guarantees byte-
* identical behavior to pre-multi-tier code for the gpu_cfg == NULL
* case.
*
* - When g_n_gpus >= 2 (multi-tier engine), translates the logical
* tier index to the corresponding physical CUDA device id via
* g_gpu[logical_tier].device_id and looks up the slice strictly
* in the per-device selective cache via ds4_gpu_lookup_cache_strict.
* On miss, logs a diagnostic and returns NULL (no host-pointer
* fallback — a miss here is a placement/install bug).
*
* The caller is responsible for cudaSetDevice'ing to the right physical
* device before launching the kernel that consumes the returned pointer.
* Wrappers in this file thread `int logical_tier` from the dispatch
* caller; single-tier callers pass 0, which hits the short-circuit.
*
* Added for multi-GPU execution (multi-GPU execution), sub-area 3 of the
* spec. */
/* Optional second (support) model map for speculative decoding. The strict
* multi-tier cache is keyed by source offset only, so support tensors are
* installed and resolved with a large disjoint offset bias. */
static const void *g_support_host_base = NULL;
static uint64_t g_support_host_size = 0;
static uint64_t g_support_offset_bias = 0;
extern "C" uint64_t ds4_gpu_tier_free_vram(int logical_tier) {
if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0;
int prev = -1;
if (cudaGetDevice(&prev) != cudaSuccess) prev = -1;
if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) return 0;
size_t free_b = 0, total_b = 0;
uint64_t out = 0;
if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) out = (uint64_t)free_b;
if (prev >= 0) (void)cudaSetDevice(prev);
return out;
}
extern "C" int ds4_gpu_register_support_map(const void *map, uint64_t size, uint64_t bias) {
if (!map || size == 0 || bias == 0) return 0;
g_support_host_base = map;
g_support_host_size = size;
g_support_offset_bias = bias;
return 1;
}
static const char *cuda_resolve_weight_ptr(const void *model_map,
uint64_t offset,
uint64_t bytes,
int logical_tier,
const char *label) {
if (g_n_gpus <= 1) {
return cuda_model_range_ptr(model_map, offset, bytes, label);
}
if (g_support_host_base && model_map == g_support_host_base) {
offset += g_support_offset_bias;
}
if (logical_tier < 0 || logical_tier >= g_n_gpus) {
fprintf(stderr,
"ds4: cuda_resolve_weight_ptr: bad tier %d (n_gpus=%d, label=%s)\n",
logical_tier, g_n_gpus, label ? label : "?");
return NULL;
}
const int physical_device = g_gpu[logical_tier].device_id;
void *dev_ptr = NULL;
if (ds4_gpu_lookup_cache_strict(offset, bytes, physical_device, &dev_ptr)
&& dev_ptr) {
return (const char *)dev_ptr;
}
/* GLM multi-tier: generic launchers resolve by the OUT tensor's tier,
* but the executing device (set per layer) is where the weights were
* cached. Retry with the current device before declaring a miss;
* DS4 paths never reach this (out tier == current device). */
int cur_dev = -1;
if (cudaGetDevice(&cur_dev) == cudaSuccess &&
cur_dev != physical_device &&
ds4_gpu_lookup_cache_strict(offset, bytes, cur_dev, &dev_ptr) &&
dev_ptr) {
return (const char *)dev_ptr;
}
fprintf(stderr,
"ds4: selective-cache miss for offset=%llu bytes=%llu on "
"logical_tier=%d (physical_device=%d, current_device=%d, "
"label=%s); this is a placement/cache-install bug\n",
(unsigned long long)offset, (unsigned long long)bytes,
logical_tier, physical_device, cur_dev, label ? label : "?");
return NULL;
}
static int cuda_model_range_is_cached(const void *model_map, uint64_t offset, uint64_t bytes) {
if (bytes == 0) return 1;
if (model_map == g_model_host_base &&
(g_model_device_owned || g_model_registered)) return 1;
const uint64_t end = offset + bytes;
if (end < offset) return 0;
for (const cuda_model_range &r : g_model_ranges) {
if (r.host_base == model_map &&
offset >= r.offset &&
end <= r.offset + r.bytes) {
return 1;
}
if (r.host_base == model_map &&
r.host_registered &&
r.registered_base &&
r.registered_device_base) {
const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset);
const uintptr_t h1 = h0 + bytes;
const uintptr_t r0 = (uintptr_t)r.registered_base;
const uintptr_t r1 = r0 + r.registered_bytes;
if (h1 >= h0 && h0 >= r0 && h1 <= r1) return 1;
}
}
return 0;
}
static void cuda_q8_f16_cache_release_all(void) {
for (const cuda_q8_f16_range &r : g_q8_f16_ranges) {
(void)cudaFree(r.device_ptr);
}
g_q8_f16_ranges.clear();
g_q8_f16_by_offset.clear();
g_q8_f16_bytes = 0;
}
static uint64_t cuda_parse_mib_env(const char *name, int *present) {
const char *env = getenv(name);
if (present) *present = 0;
if (!env || !env[0]) return 0;
char *end = NULL;
unsigned long long v = strtoull(env, &end, 10);
if (end == env || *end != '\0') return 0;
if (present) *present = 1;
if (v > UINT64_MAX / 1048576ull) return UINT64_MAX;
return (uint64_t)v * 1048576ull;
}
static uint32_t cuda_parse_u32_env_clamped(const char *name, uint32_t fallback,
uint32_t min_value, uint32_t max_value,
int *present) {
const char *env = getenv(name);
if (present) *present = 0;
if (!env || !env[0]) return fallback;
errno = 0;
char *end = NULL;
unsigned long v = strtoul(env, &end, 10);
if (errno != 0 || end == env || *end != '\0') return fallback;
if (present) *present = 1;
if (v < min_value) return min_value;
if (v > max_value) return max_value;
return (uint32_t)v;
}
static int cuda_env_flag_enabled(const char *name, int fallback) {
const char *env = getenv(name);
if (!env || !env[0]) return fallback;
return strcmp(env, "0") != 0;
}
extern "C" int ds4_gpu_set_decode_fast_attention(int enabled) {
const int old = g_decode_fast_attention;
g_decode_fast_attention = enabled != 0;
return old;
}
extern "C" int ds4_gpu_set_decode_score_vec4(int enabled) {
const int old = g_decode_score_vec4;
g_decode_score_vec4 = enabled != 0;
return old;
}
static bool cuda_splitkv_decode_requested(void) {
if (cuda_env_flag_enabled("DS4_CUDA_NO_SPLITKV_DECODE", 0)) return false;
return g_decode_fast_attention ||
cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0);
}
static uint64_t cuda_q8_f16_cache_limit_bytes(void) {
int present = 0;
const uint64_t limit = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_MB", &present);
return present ? limit : UINT64_MAX;
}
static uint64_t cuda_q8_f16_cache_reserve_bytes(uint64_t total_bytes) {
int present = 0;
const uint64_t reserve = cuda_parse_mib_env("DS4_CUDA_Q8_F16_CACHE_RESERVE_MB", &present);
if (present) return reserve;
if (total_bytes >= 112ull * 1024ull * 1024ull * 1024ull) {
return 512ull * 1048576ull;
}
/* High-VRAM cards (>= 40 GiB, e.g. 48 GiB RTX 6000 Ada): use a small
* reserve so the selective Q8->F16 cache can actually engage at tight
* budgets (e.g. --gpu-vram 47,47, where the 81 GB model leaves only ~1.3
* GiB free and the old 4 GiB floor rejected every cache allocation,
* forcing the scalar DP4A prefill kernel).
*
* NOTE: this 768 MiB value is a *bounded cache-growth guard*, not a hard
* guarantee that live free VRAM stays >= 768 MiB. cuda_q8_f16_cache_has_budget
* only blocks a *cache* allocation when free - request < reserve at that
* moment; allocations made outside cache accounting (cuda_tmp_alloc_on
* activation/prequant buffers, cuBLAS internal workspaces) can still dip
* below it. 768 MiB is chosen to leave headroom above the ~0.5 GiB
* memory-safety floor for those out-of-cache allocations; actual minimum
* free VRAM is verified by measurement, and the disable-after-failure path
* degrades gracefully if cuBLAS/alloc ever fails under pressure. Set
* DS4_CUDA_Q8_F16_CACHE_RESERVE_MB=4096 to restore the prior behavior. */
if (total_bytes >= 40ull * 1024ull * 1024ull * 1024ull) {
const uint64_t hi_min_reserve = 768ull * 1048576ull;
const uint64_t hi_pct_reserve = total_bytes / 100u; /* 1% */
return hi_pct_reserve > hi_min_reserve ? hi_pct_reserve : hi_min_reserve;
}
/* Smaller cards (< 40 GiB): keep the conservative reserve. The expanded
* Q8->F16 cache is only an acceleration path; on a small card a sub-GiB
* reserve would be a large fraction of total VRAM, so keep enough free for
* cuBLAS workspaces, transient graph buffers, and driver bookkeeping. */
const uint64_t min_reserve = 4096ull * 1048576ull;
const uint64_t pct_reserve = total_bytes / 20u; /* 5% */
return pct_reserve > min_reserve ? pct_reserve : min_reserve;
}
static void cuda_q8_f16_cache_budget_notice(
const char *reason,
uint64_t request_bytes,
uint64_t free_bytes,
uint64_t total_bytes,
uint64_t reserve_bytes,
uint64_t limit_bytes) {
if (g_q8_f16_budget_notice_printed && getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") == NULL) return;
g_q8_f16_budget_notice_printed = 1;
if (limit_bytes != UINT64_MAX && free_bytes == 0 && total_bytes == 0 && reserve_bytes == 0) {
fprintf(stderr,
"ds4: CUDA q8 fp16 cache %s; using q8 kernels "
"(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB)\n",
reason,
(double)request_bytes / 1048576.0,
(double)g_q8_f16_bytes / 1073741824.0,
(double)limit_bytes / 1073741824.0);
} else if (limit_bytes == UINT64_MAX) {
fprintf(stderr,
"ds4: CUDA q8 fp16 cache %s; using q8 kernels "
"(request=%.2f MiB cached=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n",
reason,
(double)request_bytes / 1048576.0,
(double)g_q8_f16_bytes / 1073741824.0,
(double)free_bytes / 1073741824.0,
(double)reserve_bytes / 1073741824.0,
(double)total_bytes / 1073741824.0);
} else {
fprintf(stderr,
"ds4: CUDA q8 fp16 cache %s; using q8 kernels "
"(request=%.2f MiB cached=%.2f GiB limit=%.2f GiB free=%.2f GiB reserve=%.2f GiB total=%.2f GiB)\n",
reason,
(double)request_bytes / 1048576.0,
(double)g_q8_f16_bytes / 1073741824.0,
(double)limit_bytes / 1073741824.0,
(double)free_bytes / 1073741824.0,
(double)reserve_bytes / 1073741824.0,
(double)total_bytes / 1073741824.0);
}
}
static int cuda_q8_f16_cache_has_budget(uint64_t request_bytes, const char *label) {
(void)label;
const uint64_t limit = cuda_q8_f16_cache_limit_bytes();
if (limit == 0) return 0;
if (g_q8_f16_bytes > limit || request_bytes > limit - g_q8_f16_bytes) {
cuda_q8_f16_cache_budget_notice("limit reached", request_bytes, 0, 0, 0, limit);
return 0;
}
size_t free_b = 0;
size_t total_b = 0;
cudaError_t err = cudaMemGetInfo(&free_b, &total_b);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA q8 fp16 cache memory query failed: %s; using q8 kernels\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
const uint64_t free_bytes = (uint64_t)free_b;
const uint64_t total_bytes = (uint64_t)total_b;
const uint64_t reserve_bytes = cuda_q8_f16_cache_reserve_bytes(total_bytes);
if (request_bytes > free_bytes ||
free_bytes - request_bytes < reserve_bytes) {
cuda_q8_f16_cache_budget_notice("budget exhausted", request_bytes,
free_bytes, total_bytes,
reserve_bytes, limit);
return 0;
}
return 1;
}
static void cuda_q8_f16_cache_disable_after_failure(const char *what, uint64_t request_bytes) {
if (!g_q8_f16_disabled_after_oom) {
fprintf(stderr,
"ds4: CUDA q8 fp16 cache disabled after %s "
"(request=%.2f MiB cached=%.2f GiB); using q8 kernels\n",
what ? what : "allocation failure",
(double)request_bytes / 1048576.0,
(double)g_q8_f16_bytes / 1073741824.0);
}
g_q8_f16_disabled_after_oom = 1;
if (!g_q8_f16_ranges.empty()) {
(void)cudaDeviceSynchronize();
cuda_q8_f16_cache_release_all();
}
(void)cudaGetLastError();
}
static int cuda_q8_f16_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) {
if (g_quality_mode) return 0;
if (g_q8_cache_suppressed) return 0;
if (g_q8_f16_disabled_after_oom) return 0;
if (getenv("DS4_CUDA_NO_Q8_F16_CACHE") != NULL) return 0;
if (cuda_q8_f16_cache_limit_bytes() == 0) return 0;
if (getenv("DS4_CUDA_Q8_F16_ALL") != NULL) return 1;
if (!label) return 0;
if (strstr(label, "attn_output_a") != NULL ||
strstr(label, "attn_output_b") != NULL ||
strstr(label, "attention_output_a") != NULL ||
strstr(label, "attention_output_b") != NULL) {
return getenv("DS4_CUDA_NO_ATTENTION_OUTPUT_F16_CACHE") == NULL;
}
if (strstr(label, "attn_q_b") != NULL) {
return getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL;
}
if (strstr(label, "ffn_gate_shexp") != NULL ||
strstr(label, "ffn_up_shexp") != NULL ||
strstr(label, "ffn_down_shexp") != NULL) {
return 1;
}
return (in_dim == 4096u && out_dim == 2048u) ||
(in_dim == 2048u && out_dim == 4096u) ||
(in_dim == 4096u && out_dim == 1024u) ||
(in_dim == 4096u && out_dim == 512u) ||
(getenv("DS4_CUDA_NO_ATTN_Q_B_F16_CACHE") == NULL &&
in_dim == 1024u && out_dim == 32768u);
}
static int cuda_q8_label_is_attention_output(const char *label) {
return label &&
(strstr(label, "attn_output_a") != NULL ||
strstr(label, "attn_output_b") != NULL ||
strstr(label, "attention_output_a") != NULL ||
strstr(label, "attention_output_b") != NULL);
}
static int cuda_q8_use_dp4a(void) {
return getenv("DS4_CUDA_NO_Q8_DP4A") == NULL;
}
static unsigned cuda_q8_exact_threads(uint64_t blocks) {
if (blocks <= 64u) return 64u;
if (blocks <= 128u) return 128u;
return 256u;
}
static int cuda_q8_f16_preload_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) {
if (cuda_q8_label_is_attention_output(label) &&
getenv("DS4_CUDA_ATTENTION_OUTPUT_PRELOAD") == NULL &&
getenv("DS4_CUDA_Q8_F16_ALL") == NULL) {
return 0;
}
return cuda_q8_f16_cache_allowed(label, in_dim, out_dim);
}
static int cuda_q8_f32_cache_allowed(const char *label, uint64_t in_dim, uint64_t out_dim) {
if (g_q8_cache_suppressed) return 0;
if (getenv("DS4_CUDA_NO_Q8_F32_CACHE") != NULL) return 0;
if (getenv("DS4_CUDA_Q8_F32_ALL") != NULL) return 1;
if (label && strstr(label, "attn_q_b") != NULL) {
return getenv("DS4_CUDA_ATTN_Q_B_F32_CACHE") != NULL;
}
return getenv("DS4_CUDA_Q8_F32_LARGE") != NULL &&
in_dim == 1024u && out_dim == 32768u;
}
/* Look up a per-device dequantized fp16 slice of the Q8_0 weight at
* (model_map, offset, weight_bytes, in_dim, out_dim). expected_device is a
* PHYSICAL CUDA device id (0 in single-tier; g_gpu[logical_tier].device_id in
* multi-tier). On hit returns the cached pointer for that device. On miss
* cudaSetDevice's to expected_device, allocates + dequants there, stamps the
* new entry with device_id == expected_device, restores the previous device,
* and returns the new pointer.
*
* Single-tier (g_n_gpus <= 1) uses the offset-keyed map for a fast path —
* legacy entries were stamped device_id=0, so the map remains authoritative.
* Multi-tier linear-scans the ranges vector filtering on device_id (the same
* offset may now legitimately map to multiple entries, one per device).
*/
static const __half *cuda_q8_f16_ptr(
const void *model_map,
uint64_t offset,
uint64_t weight_bytes,
uint64_t in_dim,
uint64_t out_dim,
int expected_device,
const char *label) {
if (g_n_gpus <= 1) {
auto exact = g_q8_f16_by_offset.find(offset);
if (exact != g_q8_f16_by_offset.end()) {
const cuda_q8_f16_range &r = g_q8_f16_ranges[exact->second];
if (r.host_base == model_map && r.weight_bytes == weight_bytes &&
r.in_dim == in_dim && r.out_dim == out_dim) {
return r.device_ptr;
}
}
} else {
for (const cuda_q8_f16_range &r : g_q8_f16_ranges) {
if (r.host_base == model_map &&
r.offset == offset &&
r.weight_bytes == weight_bytes &&
r.in_dim == in_dim &&
r.out_dim == out_dim &&
r.device_id == expected_device) {
return r.device_ptr;
}
}
}
if (!cuda_q8_f16_cache_allowed(label, in_dim, out_dim)) return NULL;
/* Source Q8 bytes:
* - Single-tier (g_n_gpus <= 1): cuda_model_range_ptr — preserves the
* legacy behavior (FD cache, host-register, or cudaMalloc-and-copy).
* - Multi-tier (g_n_gpus > 1): the per-device selective cache must
* already contain the weight on expected_device. Use the strict
* lookup; on miss this is a placement bug and we hard-fail.
*/
const char *q8;
if (g_n_gpus <= 1) {
q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, "q8_0");
} else {
void *strict_ptr = NULL;
if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) ||
!strict_ptr) {
fprintf(stderr,
"ds4: q8 fp16 cache miss: source bytes not in selective cache for "
"offset=%llu bytes=%llu device=%d (label=%s); placement bug\n",
(unsigned long long)offset, (unsigned long long)weight_bytes,
expected_device, label ? label : "?");
return NULL;
}
q8 = (const char *)strict_ptr;
}
if (!q8) return NULL;
if (in_dim != 0 && out_dim > UINT64_MAX / in_dim / sizeof(__half)) return NULL;
const uint64_t out_bytes = in_dim * out_dim * sizeof(__half);
if (!cuda_q8_f16_cache_has_budget(out_bytes, label)) return NULL;
int prev = -1;
if (g_n_gpus > 1) {
cudaError_t derr = cudaGetDevice(&prev);
if (derr != cudaSuccess) {
fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp16 alloc on device %d: %s\n",
expected_device, cudaGetErrorString(derr));
(void)cudaGetLastError();
return NULL;
}
derr = cudaSetDevice(expected_device);
if (derr != cudaSuccess) {
fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp16 alloc: %s\n",
expected_device, cudaGetErrorString(derr));
(void)cudaGetLastError();
if (prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
}
__half *dev = NULL;
cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA q8 fp16 cache alloc failed on device %d (%.2f MiB): %s\n",
expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err));
cuda_q8_f16_cache_disable_after_failure("allocation failure", out_bytes);
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
const uint64_t blocks = (in_dim + 31) / 32;
const uint64_t n = in_dim * out_dim;
dequant_q8_0_to_f16_kernel<<<(n + 255) / 256, 256>>>(dev,
(const unsigned char *)q8,
in_dim,
out_dim,
blocks);
if (!cuda_ok(cudaGetLastError(), "q8 fp16 dequant launch")) {
(void)cudaFree(dev);
cuda_q8_f16_cache_disable_after_failure("dequant launch failure", out_bytes);
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
g_q8_f16_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device});
if (g_n_gpus <= 1) {
g_q8_f16_by_offset[offset] = g_q8_f16_ranges.size() - 1u;
}
g_q8_f16_bytes += out_bytes;
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA cached q8 fp16 %.2f MiB on device %d (total %.2f GiB)\n",
(double)out_bytes / 1048576.0, expected_device,
(double)g_q8_f16_bytes / 1073741824.0);
}
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return dev;
}
/* Per-device dequantized fp32 cache. Same conventions as cuda_q8_f16_ptr. */
static float *cuda_q8_f32_ptr(
const void *model_map,
uint64_t offset,
uint64_t weight_bytes,
uint64_t in_dim,
uint64_t out_dim,
int expected_device,
const char *label) {
if (g_n_gpus <= 1) {
auto exact = g_q8_f32_by_offset.find(offset);
if (exact != g_q8_f32_by_offset.end()) {
const cuda_q8_f32_range &r = g_q8_f32_ranges[exact->second];
if (r.host_base == model_map && r.weight_bytes == weight_bytes &&
r.in_dim == in_dim && r.out_dim == out_dim) {
return r.device_ptr;
}
}
} else {
for (const cuda_q8_f32_range &r : g_q8_f32_ranges) {
if (r.host_base == model_map &&
r.offset == offset &&
r.weight_bytes == weight_bytes &&
r.in_dim == in_dim &&
r.out_dim == out_dim &&
r.device_id == expected_device) {
return r.device_ptr;
}
}
}
if (!cuda_q8_f32_cache_allowed(label, in_dim, out_dim)) return NULL;
/* Source Q8 bytes: legacy path in single-tier; strict per-device lookup
* in multi-tier (same rationale as cuda_q8_f16_ptr). */
const char *q8;
if (g_n_gpus <= 1) {
q8 = cuda_model_range_ptr(model_map, offset, weight_bytes, label ? label : "q8_0");
} else {
void *strict_ptr = NULL;
if (!ds4_gpu_lookup_cache_strict(offset, weight_bytes, expected_device, &strict_ptr) ||
!strict_ptr) {
fprintf(stderr,
"ds4: q8 fp32 cache miss: source bytes not in selective cache for "
"offset=%llu bytes=%llu device=%d (label=%s); placement bug\n",
(unsigned long long)offset, (unsigned long long)weight_bytes,
expected_device, label ? label : "?");
return NULL;
}
q8 = (const char *)strict_ptr;
}
if (!q8) return NULL;
const uint64_t out_bytes = in_dim * out_dim * sizeof(float);
int prev = -1;
if (g_n_gpus > 1) {
cudaError_t derr = cudaGetDevice(&prev);
if (derr != cudaSuccess) {
fprintf(stderr, "ds4: cudaGetDevice failed before q8 fp32 alloc on device %d: %s\n",
expected_device, cudaGetErrorString(derr));
(void)cudaGetLastError();
return NULL;
}
derr = cudaSetDevice(expected_device);
if (derr != cudaSuccess) {
fprintf(stderr, "ds4: cudaSetDevice(%d) failed before q8 fp32 alloc: %s\n",
expected_device, cudaGetErrorString(derr));
(void)cudaGetLastError();
if (prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
}
float *dev = NULL;
cudaError_t err = cudaMalloc(&dev, (size_t)out_bytes);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA q8 fp32 cache alloc failed on device %d (%.2f MiB): %s\n",
expected_device, (double)out_bytes / 1048576.0, cudaGetErrorString(err));
(void)cudaGetLastError();
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
const uint64_t blocks = (in_dim + 31) / 32;
const uint64_t n = in_dim * out_dim;
dequant_q8_0_to_f32_kernel<<<(n + 255) / 256, 256>>>(dev,
(const unsigned char *)q8,
in_dim,
out_dim,
blocks);
if (!cuda_ok(cudaGetLastError(), "q8 fp32 dequant launch")) {
(void)cudaFree(dev);
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return NULL;
}
g_q8_f32_ranges.push_back({model_map, offset, weight_bytes, in_dim, out_dim, dev, expected_device});
if (g_n_gpus <= 1) {
g_q8_f32_by_offset[offset] = g_q8_f32_ranges.size() - 1u;
}
g_q8_f32_bytes += out_bytes;
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA cached q8 fp32 %.2f MiB on device %d (total %.2f GiB)\n",
(double)out_bytes / 1048576.0, expected_device,
(double)g_q8_f32_bytes / 1073741824.0);
}
if (g_n_gpus > 1 && prev >= 0) (void)cudaSetDevice(prev);
return dev;
}
static int cuda_ok(cudaError_t err, const char *what) {
if (err == cudaSuccess) return 1;
fprintf(stderr, "ds4: CUDA %s failed: %s\n", what, cudaGetErrorString(err));
return 0;
}
static double cuda_wall_sec(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec * 1.0e-9;
}
static int cuda_model_load_progress_enabled(void) {
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE") != NULL) return 0;
return 1;
}
static void cuda_model_load_progress_reset(void) {
g_model_load_progress_next = 0;
g_model_load_progress_last = 0.0;
g_model_load_progress_started = 0;
g_model_load_progress_tty = 0;
}
static void cuda_model_load_progress_note(uint64_t cached_bytes) {
if (!cuda_model_load_progress_enabled()) return;
const double now = cuda_wall_sec();
if (!g_model_load_progress_started) {
g_model_load_progress_started = 1;
g_model_load_progress_tty = isatty(STDERR_FILENO) != 0;
g_model_load_progress_next = (g_model_load_progress_tty ? 2ull : 16ull) *
1024ull * 1024ull * 1024ull;
g_model_load_progress_last = now;
if (g_model_load_progress_tty) {
fprintf(stderr, "ds4: CUDA loading model tensors into device cache: 0.00 GiB");
} else {
fprintf(stderr, "ds4: CUDA loading model tensors into device cache\n");
}
}
if (cached_bytes < g_model_load_progress_next &&
now - g_model_load_progress_last < (g_model_load_progress_tty ? 2.0 : 10.0)) {
return;
}
if (g_model_load_progress_tty) {
fprintf(stderr, "\rds4: CUDA loading model tensors into device cache: %.2f GiB",
(double)cached_bytes / 1073741824.0);
} else {
fprintf(stderr, "ds4: CUDA loading model tensors %.2f GiB cached\n",
(double)cached_bytes / 1073741824.0);
}
fflush(stderr);
g_model_load_progress_last = now;
const uint64_t step = (g_model_load_progress_tty ? 2ull : 16ull) *
1024ull * 1024ull * 1024ull;
while (g_model_load_progress_next <= cached_bytes) {
g_model_load_progress_next += step;
}
}
static int cuda_model_prefetch_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) {
if (!model_map || map_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0;
if (getenv("DS4_CUDA_NO_MODEL_PREFETCH") != NULL ||
getenv("DS4_CUDA_COPY_MODEL") != NULL ||
getenv("DS4_CUDA_WEIGHT_CACHE") != NULL ||
getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) {
return 0;
}
int device = 0;
if (cudaGetDevice(&device) != cudaSuccess) {
(void)cudaGetLastError();
return 0;
}
int pageable = 0;
cudaError_t err = cudaDeviceGetAttribute(&pageable, cudaDevAttrPageableMemoryAccess, device);
if (err != cudaSuccess || !pageable) {
(void)cudaGetLastError();
return 0;
}
#if CUDART_VERSION >= 13000
cudaMemLocation loc;
memset(&loc, 0, sizeof(loc));
loc.type = cudaMemLocationTypeDevice;
loc.id = device;
#else
int loc = device;
#endif
const long page_sz_l = sysconf(_SC_PAGESIZE);
const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u;
const uintptr_t host_addr = (uintptr_t)((const char *)model_map + map_offset);
const uintptr_t pre_addr = host_addr & ~(uintptr_t)(page_sz - 1u);
const uint64_t pre_delta = (uint64_t)(host_addr - pre_addr);
const uint64_t pre_bytes = (pre_delta + map_size + page_sz - 1u) & ~(page_sz - 1u);
void *pre_ptr = (void *)pre_addr;
const double t0 = cuda_wall_sec();
err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetReadMostly, loc);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model read-mostly advise skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
err = cudaMemAdvise(pre_ptr, (size_t)pre_bytes, cudaMemAdviseSetPreferredLocation, loc);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model preferred-location advise skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
if (!g_model_prefetch_stream) {
err = cudaStreamCreateWithFlags(&g_model_prefetch_stream, cudaStreamNonBlocking);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model prefetch stream creation skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
}
#if CUDART_VERSION >= 13000
err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, 0, g_model_prefetch_stream);
#else
err = cudaMemPrefetchAsync(pre_ptr, (size_t)pre_bytes, loc, g_model_prefetch_stream);
#endif
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model prefetch skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
if (getenv("DS4_CUDA_MODEL_PREFETCH_SYNC") != NULL) {
err = cudaStreamSynchronize(g_model_prefetch_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model prefetch sync failed: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
}
const double t1 = cuda_wall_sec();
fprintf(stderr,
"ds4: CUDA ATS/HMM prefetch queued %.2f GiB of model tensors in %.3fs\n",
(double)map_size / 1073741824.0,
t1 - t0);
g_model_hmm_direct = 1;
return 1;
}
static uint64_t cuda_model_copy_chunk_bytes(void) {
uint64_t mb = 64;
const char *env = getenv("DS4_CUDA_MODEL_COPY_CHUNK_MB");
if (env && env[0]) {
char *end = NULL;
unsigned long long v = strtoull(env, &end, 10);
if (end != env && v > 0) mb = (uint64_t)v;
}
if (mb < 16) mb = 16;
if (mb > 4096) mb = 4096;
return mb * 1048576ull;
}
static void cuda_model_discard_source_pages(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes) {
#if defined(POSIX_MADV_DONTNEED)
if (getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || !model_map || bytes == 0 || offset > model_size) return;
if (bytes > model_size - offset) bytes = model_size - offset;
const long page_sz_l = sysconf(_SC_PAGESIZE);
const uint64_t page_sz = page_sz_l > 0 ? (uint64_t)page_sz_l : 4096u;
const uintptr_t h0 = (uintptr_t)((const char *)model_map + offset);
const uintptr_t h1 = h0 + bytes;
const uintptr_t p0 = h0 & ~(uintptr_t)(page_sz - 1u);
const uintptr_t p1 = (h1 + page_sz - 1u) & ~(uintptr_t)(page_sz - 1u);
if (p1 > p0) (void)posix_madvise((void *)p0, (size_t)(p1 - p0), POSIX_MADV_DONTNEED);
#else
(void)model_map;
(void)model_size;
(void)offset;
(void)bytes;
#endif
}
static void cuda_model_drop_file_pages(uint64_t offset, uint64_t bytes) {
#if defined(POSIX_FADV_DONTNEED)
if (g_model_fd < 0 || getenv("DS4_CUDA_KEEP_MODEL_PAGES") != NULL || bytes == 0) return;
(void)posix_fadvise(g_model_fd, (off_t)offset, (off_t)bytes, POSIX_FADV_DONTNEED);
#else
(void)offset;
(void)bytes;
#endif
}
static uint64_t cuda_round_down(uint64_t v, uint64_t align) {
if (align <= 1) return v;
return (v / align) * align;
}
static uint64_t cuda_round_up(uint64_t v, uint64_t align) {
if (align <= 1) return v;
const uint64_t rem = v % align;
return rem == 0 ? v : v + (align - rem);
}
static void *cuda_align_ptr(void *ptr, uint64_t align) {
if (align <= 1) return ptr;
uintptr_t p = (uintptr_t)ptr;
uintptr_t a = (uintptr_t)align;
return (void *)(((p + a - 1u) / a) * a);
}
static int cuda_model_stage_pool_alloc(uint64_t bytes) {
if (g_model_stage_bytes >= bytes) return 1;
for (size_t i = 0; i < 4; i++) {
if (g_model_stage_event[i]) {
(void)cudaEventDestroy(g_model_stage_event[i]);
g_model_stage_event[i] = NULL;
}
if (g_model_stage_raw[i]) {
(void)cudaFreeHost(g_model_stage_raw[i]);
g_model_stage_raw[i] = NULL;
g_model_stage[i] = NULL;
}
}
g_model_stage_bytes = 0;
if (!g_model_upload_stream) {
cudaError_t err = cudaStreamCreateWithFlags(&g_model_upload_stream, cudaStreamNonBlocking);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model upload stream creation failed: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
}
for (size_t i = 0; i < 4; i++) {
cudaError_t err = cudaMallocHost(&g_model_stage_raw[i], (size_t)bytes);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
g_model_stage[i] = cuda_align_ptr(g_model_stage_raw[i], g_model_direct_align);
err = cudaEventCreateWithFlags(&g_model_stage_event[i], cudaEventDisableTiming);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model staging event creation failed: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
}
g_model_stage_bytes = bytes;
return 1;
}
static int cuda_pread_full(int fd, void *buf, uint64_t bytes, uint64_t offset) {
uint64_t done = 0;
while (done < bytes) {
const size_t n_req = (bytes - done > (uint64_t)SSIZE_MAX) ? (size_t)SSIZE_MAX : (size_t)(bytes - done);
ssize_t n = pread(fd, (char *)buf + done, n_req, (off_t)(offset + done));
if (n < 0) {
if (errno == EINTR) continue;
return 0;
}
if (n == 0) { errno = EIO; return 0; }
done += (uint64_t)n;
}
return 1;
}
static int cuda_model_stage_read_from(int fd, int *direct_fd, uint64_t align,
uint64_t file_size, void *stage, uint64_t stage_bytes,
uint64_t offset, uint64_t bytes, const char **payload) {
*payload = (const char *)stage;
#if defined(__linux__) && defined(O_DIRECT)
if (*direct_fd >= 0 && align > 1 && file_size != 0) {
const uint64_t aligned_off = cuda_round_down(offset, align);
const uint64_t delta = offset - aligned_off;
uint64_t read_size = cuda_round_up(delta + bytes, align);
if (aligned_off <= file_size &&
read_size <= stage_bytes &&
read_size <= file_size - aligned_off) {
const int saved_errno = errno;
errno = 0;
if (cuda_pread_full(*direct_fd, stage, read_size, aligned_off)) {
*payload = (const char *)stage + delta;
errno = saved_errno;
return 1;
}
const int direct_errno = errno;
if (direct_errno == EINVAL || direct_errno == EFAULT || direct_errno == ENOTSUP || direct_errno == EOPNOTSUPP) {
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA direct model read disabled: %s\n", strerror(direct_errno));
}
(void)close(*direct_fd);
*direct_fd = -1;
}
errno = direct_errno;
}
}
#else
(void)stage_bytes;
(void)direct_fd; (void)align; (void)file_size;
#endif
return cuda_pread_full(fd, stage, bytes, offset);
}
static int cuda_model_stage_read(void *stage, uint64_t stage_bytes,
uint64_t offset, uint64_t bytes,
const char **payload) {
return cuda_model_stage_read_from(g_model_fd, &g_model_direct_fd,
g_model_direct_align, g_model_file_size, stage, stage_bytes, offset, bytes, payload);
}
static void cuda_stream_selected_stage_release(void) {
for (size_t i = 0; i < 4; i++) {
if (g_stream_selected_stage_event[i]) {
(void)cudaEventDestroy(g_stream_selected_stage_event[i]);
g_stream_selected_stage_event[i] = NULL;
}
if (g_stream_selected_stage_raw[i]) {
(void)cudaFreeHost(g_stream_selected_stage_raw[i]);
g_stream_selected_stage_raw[i] = NULL;
g_stream_selected_stage[i] = NULL;
}
}
g_stream_selected_stage_bytes = 0;
if (g_stream_selected_upload_stream) {
(void)cudaStreamDestroy(g_stream_selected_upload_stream);
g_stream_selected_upload_stream = NULL;
}
}
static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) {
if (g_stream_selected_stage_bytes >= bytes) return 1;
cuda_stream_selected_stage_release();
cudaError_t err = cudaStreamCreateWithFlags(
&g_stream_selected_upload_stream, cudaStreamNonBlocking);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected upload stream creation failed: %s\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
for (size_t i = 0; i < 4; i++) {
err = cudaMallocHost(&g_stream_selected_stage_raw[i], (size_t)bytes);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected staging allocation failed: %s\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
cuda_stream_selected_stage_release();
return 0;
}
g_stream_selected_stage[i] = cuda_align_ptr(
g_stream_selected_stage_raw[i], g_model_direct_align);
err = cudaEventCreateWithFlags(&g_stream_selected_stage_event[i],
cudaEventDisableTiming);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected staging event creation failed: %s\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
cuda_stream_selected_stage_release();
return 0;
}
}
g_stream_selected_stage_bytes = bytes;
return 1;
}
static int cuda_model_copy_to_device_streamed(
char *dst,
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t bytes,
const char *what,
uint64_t &chunk_idx) {
if (!dst || !model_map || offset > model_size ||
bytes > model_size - offset) {
return 0;
}
if (bytes == 0) return 1;
if (g_model_fd < 0 ||
(g_model_fd_host_base != NULL && model_map != g_model_fd_host_base)) {
return cuda_ok(cudaMemcpy(dst,
(const char *)model_map + offset,
(size_t)bytes,
cudaMemcpyHostToDevice),
what ? what : "stream selected expert copy");
}
const uint64_t chunk = cuda_model_copy_chunk_bytes();
const uint64_t stage_bytes =
chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1);
if (!cuda_stream_selected_stage_pool_alloc(stage_bytes)) return 0;
uint64_t copied = 0;
while (copied < bytes) {
const uint64_t n = bytes - copied < chunk ? bytes - copied : chunk;
const uint64_t bi = chunk_idx % 4u;
cudaError_t err;
if (chunk_idx >= 4u) {
err = cudaEventSynchronize(g_stream_selected_stage_event[bi]);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected staging wait failed for %s: %s\n",
what ? what : "expert", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
}
const char *payload = NULL;
if (!cuda_model_stage_read(g_stream_selected_stage[bi],
g_stream_selected_stage_bytes,
offset + copied, n, &payload)) {
fprintf(stderr,
"ds4: CUDA streaming selected read failed for %s at %.2f MiB: %s\n",
what ? what : "expert", (double)copied / 1048576.0,
strerror(errno));
return 0;
}
err = cudaMemcpyAsync(dst + copied, payload, (size_t)n,
cudaMemcpyHostToDevice,
g_stream_selected_upload_stream);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected copy failed for %s at %.2f MiB: %s\n",
what ? what : "expert", (double)copied / 1048576.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
err = cudaEventRecord(g_stream_selected_stage_event[bi],
g_stream_selected_upload_stream);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA streaming selected staging record failed for %s: %s\n",
what ? what : "expert", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
cuda_model_drop_file_pages(offset + copied, n);
cuda_model_discard_source_pages(model_map, model_size,
offset + copied, n);
copied += n;
chunk_idx++;
}
return 1;
}
static uint64_t cuda_model_cache_limit_bytes(void) {
uint64_t gb = 0;
const char *env = getenv("DS4_CUDA_WEIGHT_CACHE_LIMIT_GB");
if (env && env[0]) {
char *end = NULL;
unsigned long long v = strtoull(env, &end, 10);
if (end != env) gb = (uint64_t)v;
}
if (gb == 0) return UINT64_MAX;
return gb * 1073741824ull;
}
static uint64_t cuda_model_arena_chunk_bytes(uint64_t need) {
/* Large tensor/shard ranges rarely fit each other's remaining arena tail.
* Allocate them directly so cached-byte budgets reflect resident memory. */
if (need >= 256ull * 1048576ull) return need;
uint64_t mb = 1792;
const char *env = getenv("DS4_CUDA_WEIGHT_ARENA_CHUNK_MB");
if (env && env[0]) {
char *end = NULL;
unsigned long long v = strtoull(env, &end, 10);
if (end != env && v > 0) mb = (uint64_t)v;
}
if (mb < 256) mb = 256;
if (mb > 8192) mb = 8192;
return mb * 1048576ull;
}
static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) {
if (bytes == 0) return NULL;
if (g_model_cache_full) return NULL;
const uint64_t align = 256u;
const uint64_t aligned = (bytes + align - 1u) & ~(align - 1u);
for (cuda_model_arena &a : g_model_arenas) {
const uint64_t used = (a.used + align - 1u) & ~(align - 1u);
if (used <= a.bytes && aligned <= a.bytes - used) {
char *ptr = a.device_ptr + used;
a.used = used + aligned;
return ptr;
}
}
const uint64_t limit = cuda_model_cache_limit_bytes();
if (g_model_range_bytes > limit || aligned > limit - g_model_range_bytes) return NULL;
const uint64_t chunk = cuda_model_arena_chunk_bytes(aligned);
void *dev = NULL;
cudaError_t err = cudaMalloc(&dev, (size_t)chunk);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model arena alloc failed for %s (%.2f MiB chunk): %s\n",
what ? what : "weights",
(double)chunk / 1048576.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
g_model_cache_full = 1;
return NULL;
}
g_model_arenas.push_back({(char *)dev, chunk, aligned});
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
uint64_t arena_bytes = 0;
for (const cuda_model_arena &a : g_model_arenas) arena_bytes += a.bytes;
fprintf(stderr, "ds4: CUDA model arena allocated %.2f MiB (arenas %.2f GiB)\n",
(double)chunk / 1048576.0,
(double)arena_bytes / 1073741824.0);
}
return (char *)dev;
}
static const char *cuda_model_range_ptr_from_fd(
const void *model_map,
uint64_t offset,
uint64_t bytes,
const char *what) {
if (g_model_fd < 0 || bytes == 0) return NULL;
if (g_model_fd_host_base != NULL && model_map != g_model_fd_host_base) return NULL;
const uint64_t limit = cuda_model_cache_limit_bytes();
if (g_model_range_bytes > limit || bytes > limit - g_model_range_bytes) {
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA direct %s %.2f MiB (cache budget %.2f GiB exhausted)\n",
what ? what : "weights",
(double)bytes / 1048576.0,
(double)limit / 1073741824.0);
}
return g_ssd_streaming_mode ? NULL : cuda_model_ptr(model_map, offset);
}
char *dev = cuda_model_arena_alloc(bytes, what);
if (!dev) {
if (g_ssd_streaming_mode || getenv("DS4_CUDA_STRICT_WEIGHT_CACHE") != NULL) return NULL;
return cuda_model_ptr(model_map, offset);
}
cudaError_t err = cudaSuccess;
const uint64_t chunk = cuda_model_copy_chunk_bytes();
const uint64_t stage_bytes = chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1);
if (!cuda_model_stage_pool_alloc(stage_bytes)) return NULL;
uint64_t copied = 0;
uint64_t chunk_idx = 0;
while (copied < bytes) {
const uint64_t n = (bytes - copied < chunk) ? (bytes - copied) : chunk;
const uint64_t bi = chunk_idx % 4u;
if (chunk_idx >= 4u) {
err = cudaEventSynchronize(g_model_stage_event[bi]);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model staging wait failed for %s: %s\n",
what ? what : "weights", cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
}
const char *payload = NULL;
if (!cuda_model_stage_read(g_model_stage[bi], g_model_stage_bytes,
offset + copied, n, &payload)) {
fprintf(stderr, "ds4: CUDA model range read failed for %s at %.2f MiB: %s\n",
what ? what : "weights",
(double)copied / 1048576.0,
strerror(errno));
return NULL;
}
err = cudaMemcpyAsync(dev + copied, payload, (size_t)n,
cudaMemcpyHostToDevice, g_model_upload_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model range copy failed for %s at %.2f MiB: %s\n",
what ? what : "weights",
(double)copied / 1048576.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
err = cudaEventRecord(g_model_stage_event[bi], g_model_upload_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model staging record failed for %s: %s\n",
what ? what : "weights", cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
cuda_model_drop_file_pages(offset + copied, n);
cuda_model_discard_source_pages(model_map, g_model_registered_size, offset + copied, n);
copied += n;
cuda_model_load_progress_note(g_model_range_bytes + copied);
chunk_idx++;
}
err = cudaStreamSynchronize(g_model_upload_stream);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model range upload sync failed for %s: %s\n",
what ? what : "weights", cudaGetErrorString(err));
(void)cudaGetLastError();
return NULL;
}
g_model_ranges.push_back({model_map, offset, bytes, dev, NULL, NULL, 0, 0, 1});
g_model_range_by_offset[offset] = g_model_ranges.size() - 1u;
g_model_range_bytes += bytes;
cuda_model_load_progress_note(g_model_range_bytes);
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA fd-cached %s %.2f MiB (total %.2f GiB)\n",
what ? what : "weights",
(double)bytes / 1048576.0,
(double)g_model_range_bytes / 1073741824.0);
}
return (const char *)dev;
}
static int cuda_model_copy_chunked(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size) {
if (!model_map || model_size == 0 || map_offset > model_size || map_size > model_size - map_offset) return 0;
if (getenv("DS4_CUDA_NO_MODEL_COPY") != NULL ||
getenv("DS4_CUDA_DIRECT_MODEL") != NULL ||
getenv("DS4_CUDA_WEIGHT_CACHE") != NULL ||
getenv("DS4_CUDA_WEIGHT_PRELOAD") != NULL) {
return 0;
}
if (g_model_device_owned || g_model_registered) return 1;
void *dev = NULL;
const double t0 = cuda_wall_sec();
cudaError_t err = cudaMalloc(&dev, (size_t)model_size);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
fprintf(stderr, "ds4: CUDA chunk-copying %.2f GiB model image\n",
(double)model_size / 1073741824.0);
const uint64_t chunk = cuda_model_copy_chunk_bytes();
void *stage = NULL;
err = cudaMallocHost(&stage, (size_t)chunk);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA pinned model staging allocation failed: %s\n", cudaGetErrorString(err));
(void)cudaFree(dev);
(void)cudaGetLastError();
return 0;
}
if (map_offset > 0) {
uint64_t copied_header = 0;
while (copied_header < map_offset) {
const uint64_t n = (map_offset - copied_header < chunk) ? (map_offset - copied_header) : chunk;
memcpy(stage, (const char *)model_map + copied_header, (size_t)n);
err = cudaMemcpy((char *)dev + copied_header, stage, (size_t)n, cudaMemcpyHostToDevice);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model header copy failed: %s\n", cudaGetErrorString(err));
(void)cudaFreeHost(stage);
(void)cudaFree(dev);
(void)cudaGetLastError();
return 0;
}
copied_header += n;
}
}
uint64_t copied = 0;
double last_report = t0;
while (copied < map_size) {
const uint64_t n = (map_size - copied < chunk) ? (map_size - copied) : chunk;
const uint64_t off = map_offset + copied;
memcpy(stage, (const char *)model_map + off, (size_t)n);
err = cudaMemcpy((char *)dev + off, stage, (size_t)n, cudaMemcpyHostToDevice);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA model chunk copy failed at %.2f GiB: %s\n",
(double)copied / 1073741824.0, cudaGetErrorString(err));
(void)cudaFreeHost(stage);
(void)cudaFree(dev);
(void)cudaGetLastError();
return 0;
}
cuda_model_discard_source_pages(model_map, model_size, off, n);
copied += n;
const double now = cuda_wall_sec();
if (getenv("DS4_CUDA_MODEL_COPY_VERBOSE") != NULL && now - last_report >= 2.0) {
fprintf(stderr, "ds4: CUDA model chunk copy %.2f/%.2f GiB\n",
(double)copied / 1073741824.0,
(double)map_size / 1073741824.0);
last_report = now;
}
}
(void)cudaFreeHost(stage);
g_model_device_base = (const char *)dev;
g_model_device_owned = 1;
g_model_hmm_direct = 0;
const double t1 = cuda_wall_sec();
fprintf(stderr,
"ds4: CUDA model chunk copy complete in %.3fs (%.2f GiB tensors)\n",
t1 - t0,
(double)map_size / 1073741824.0);
return 1;
}
static void cuda_model_range_release_all(void) {
for (const cuda_model_range &r : g_model_ranges) {
if (r.host_registered && r.registered_base) {
(void)cudaHostUnregister(r.registered_base);
} else if (r.device_ptr && !r.arena_allocated && !r.borrowed) {
(void)cudaFree(r.device_ptr);
}
}
for (const cuda_model_arena &a : g_model_arenas) {
if (a.device_ptr) (void)cudaFree(a.device_ptr);
}
g_model_arenas.clear();
g_model_ranges.clear();
g_model_range_by_offset.clear();
g_model_range_bytes = 0;
cuda_model_load_progress_reset();
}
static void cuda_derived_range_release_all(void) {
ds4_mmq_set_aligned_q81_scratch(NULL, 0);
if (g_aligned_q81_scratch) {
(void)cudaFree(g_aligned_q81_scratch);
g_aligned_q81_scratch = NULL;
}
for (const cuda_derived_range &r : g_derived_ranges) {
if (r.device_ptr) (void)cudaFree(r.device_ptr);
}
g_derived_ranges.clear();
g_derived_replace_map = NULL;
g_derived_artifact_bytes = 0;
g_derived_artifact_build_secs = 0.0;
g_derived_replaces_complete = 0;
}
static int cublas_ok(cublasStatus_t st, const char *what) {
if (st == CUBLAS_STATUS_SUCCESS) return 1;
fprintf(stderr, "ds4: cuBLAS %s failed: status %d\n", what, (int)st);
return 0;
}
extern "C" int ds4_gpu_init_multi(const ds4_gpu_config *cfg) {
if (!cfg || cfg->n_gpus < 1 || cfg->n_gpus > DS4_MAX_GPUS) return 0;
cuda_xdev_env_refresh();
cuda_decode_dispatch_env_refresh();
g_current_logical_tier = -1;
/* g_n_gpus is published incrementally so ds4_gpu_cleanup() can unwind
* partial state on failure. We publish `i + 1` BEFORE allocating any
* resources for context `i`, so even if (e.g.) stream creation
* succeeds but event creation fails, cleanup still walks device `i`
* and destroys the stream. ds4_gpu_cleanup is null-safe per field —
* partial state is OK. */
for (int i = 0; i < cfg->n_gpus; i++) {
ds4_gpu_ctx *c = &g_gpu[i];
c->device_id = cfg->device_indices[i];
if (c->device_id < 0) return 0;
/* Publish the in-progress device id so cleanup can target it on
* any later failure. cudaSetDevice is also required before
* cleanup's cudaEventDestroy / cudaStreamDestroy / cublasDestroy
* calls hit the right context. */
g_n_gpus = i + 1;
if (!cuda_ok(cudaSetDevice(c->device_id), "init set device")) return 0;
cudaDeviceProp prop;
if (cudaGetDeviceProperties(&prop, c->device_id) == cudaSuccess) {
if (i == 0) g_device_is_spark =
prop.integrated && prop.major == 12 && prop.minor == 1;
fprintf(stderr, "ds4: CUDA backend initialized on %s (sm_%d%d) dev=%d\n",
prop.name, prop.major, prop.minor, c->device_id);
}
/* Per-device stream. */
cudaStream_t s = NULL;
if (!cuda_ok(cudaStreamCreate(&s), "init stream")) return 0;
c->stream = (void *)s;
/* Per-device boundary event (reusable, no timing). */
cudaEvent_t ev = NULL;
if (!cuda_ok(cudaEventCreateWithFlags(&ev, cudaEventDisableTiming),
"init event")) return 0;
c->boundary_event = (void *)ev;
/* Per-device cuBLAS handle. */
cublasHandle_t h = NULL;
if (!cublas_ok(cublasCreate(&h), "init cublas")) return 0;
c->cublas = (void *)h;
const cublasMath_t math_mode =
(g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL)
? CUBLAS_DEFAULT_MATH
: CUBLAS_TF32_TENSOR_OP_MATH;
(void)cublasSetMathMode(h, math_mode);
c->cublas_ready = 1;
c->budget_bytes = cfg->vram_bytes[i];
c->used_bytes = 0;
c->scratch = NULL;
c->scratch_bytes = 0;
}
/* NxN peer-access matrix.
*
* Driver semantics: cudaDeviceCanAccessPeer + cudaDeviceEnablePeerAccess
* can both succeed even on hardware/drivers where cudaMemcpyPeerAsync
* silently delivers wrong data (notably RTX 6000 Ada under recent
* NVIDIA drivers, per the v0 design doc). The corruption is
* non-deterministic and can affect either or both directions of a
* pair. To guard against this, we run a multi-size, multi-iteration
* validation at init (see the loop below): write distinct known
* patterns, peer-copy them to the destination, read back, and only
* set peer_ok[i][j] if every probe round-trips byte-perfect. A
* single small probe is not sufficient — it can pass while realistic
* activation-sized copies still corrupt. On any failure the entry
* stays at 0 and cross-device copies fall back to the pinned-host
* bounce path automatically. */
for (int i = 0; i < g_n_gpus; i++) {
for (int j = 0; j < g_n_gpus; j++) {
if (i == j) { g_gpu_peer_ok[i][j] = 1; continue; }
int can = 0;
(void)cudaDeviceCanAccessPeer(&can, g_gpu[i].device_id,
g_gpu[j].device_id);
if (!can) { g_gpu_peer_ok[i][j] = 0; continue; }
(void)cudaSetDevice(g_gpu[i].device_id);
cudaError_t e = cudaDeviceEnablePeerAccess(g_gpu[j].device_id, 0);
int enabled = (e == cudaSuccess ||
e == cudaErrorPeerAccessAlreadyEnabled);
(void)cudaGetLastError();
if (!enabled) { g_gpu_peer_ok[i][j] = 0; continue; }
/* Runtime validation: peer copies on RTX 6000 Ada under recent
* NVIDIA drivers silently corrupt at realistic sizes even though
* the API returns success. cudaDeviceCanAccessPeer and
* cudaDeviceEnablePeerAccess can both report success while
* cudaMemcpyPeer delivers wrong data non-deterministically.
* Probe with multiple sizes and iterations; ALL must round-trip
* byte-perfect or we disable peer for this pair and silently fall
* back to the pinned-host bounce path. */
static const size_t kValidateSizes[] = {
4u * 1024u,
256u * 1024u,
1u * 1024u * 1024u,
16u * 1024u * 1024u,
};
const int kValidateIters = 4;
const int kNValidateSizes = (int)(sizeof(kValidateSizes) /
sizeof(kValidateSizes[0]));
const size_t kMaxValidate = kValidateSizes[kNValidateSizes - 1];
unsigned char *vh_src = (unsigned char *)malloc(kMaxValidate);
unsigned char *vh_dst = (unsigned char *)malloc(kMaxValidate);
if (!vh_src || !vh_dst) {
free(vh_src); free(vh_dst);
g_gpu_peer_ok[i][j] = 0; continue;
}
void *src_dev = NULL; void *dst_dev = NULL;
(void)cudaSetDevice(g_gpu[i].device_id);
if (cudaMalloc(&src_dev, kMaxValidate) != cudaSuccess) {
(void)cudaGetLastError();
free(vh_src); free(vh_dst);
g_gpu_peer_ok[i][j] = 0; continue;
}
(void)cudaSetDevice(g_gpu[j].device_id);
if (cudaMalloc(&dst_dev, kMaxValidate) != cudaSuccess) {
(void)cudaGetLastError();
(void)cudaSetDevice(g_gpu[i].device_id);
(void)cudaFree(src_dev);
free(vh_src); free(vh_dst);
g_gpu_peer_ok[i][j] = 0; continue;
}
int peer_validated = 1;
size_t failed_bytes = 0;
int failed_iter = -1;
for (int s_idx = 0;
s_idx < kNValidateSizes && peer_validated;
s_idx++) {
size_t n = kValidateSizes[s_idx];
for (int it = 0; it < kValidateIters && peer_validated; it++) {
for (size_t k = 0; k < n; k++) {
vh_src[k] = (unsigned char)
((k * 31u + (size_t)it * 17u +
(size_t)s_idx * 53u + 11u) & 0xffu);
}
(void)cudaSetDevice(g_gpu[i].device_id);
if (cudaMemcpy(src_dev, vh_src, n,
cudaMemcpyHostToDevice) != cudaSuccess) {
peer_validated = 0; failed_bytes = n; failed_iter = it;
break;
}
cudaError_t pc = cudaMemcpyPeer(
dst_dev, g_gpu[j].device_id,
src_dev, g_gpu[i].device_id, n);
if (pc != cudaSuccess) {
peer_validated = 0; failed_bytes = n; failed_iter = it;
break;
}
(void)cudaSetDevice(g_gpu[j].device_id);
if (cudaMemcpy(vh_dst, dst_dev, n,
cudaMemcpyDeviceToHost) != cudaSuccess) {
peer_validated = 0; failed_bytes = n; failed_iter = it;
break;
}
if (memcmp(vh_src, vh_dst, n) != 0) {
peer_validated = 0; failed_bytes = n; failed_iter = it;
break;
}
}
}
(void)cudaSetDevice(g_gpu[j].device_id);
(void)cudaFree(dst_dev);
(void)cudaSetDevice(g_gpu[i].device_id);
(void)cudaFree(src_dev);
free(vh_src);
free(vh_dst);
g_gpu_peer_ok[i][j] = peer_validated;
if (peer_validated) {
fprintf(stderr,
"ds4: peer access %d->%d validated across %d sizes x %d"
" iterations (max %zu MiB)\n",
g_gpu[i].device_id, g_gpu[j].device_id,
kNValidateSizes, kValidateIters,
kMaxValidate / (1024u * 1024u));
} else {
fprintf(stderr,
"ds4: peer access %d->%d FAILED validation at"
" size=%zu iter=%d; falling back to pinned-host bounce\n",
g_gpu[i].device_id, g_gpu[j].device_id,
failed_bytes, failed_iter);
}
}
}
g_cublas_ready = 1;
return 1;
}
extern "C" int ds4_gpu_init(void) {
ds4_gpu_config cfg;
memset(&cfg, 0, sizeof(cfg));
cfg.device_indices[0] = 0;
cfg.n_gpus = 1;
return ds4_gpu_init_multi(&cfg);
}
extern "C" void ds4_gpu_cleanup(void) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
ds4_gpu_tp_shutdown();
(void)cudaDeviceSynchronize();
ds4_gpu_decode_graphs_invalidate();
g_current_logical_tier = -1;
/* Multi-GPU teardown: events, streams, cublas handles, scratch
* slabs, per-pair bounce buffers. */
for (int i = 0; i < g_n_gpus; i++) {
ds4_gpu_ctx *c = &g_gpu[i];
(void)cudaSetDevice(c->device_id);
if (i == 0) cuda_dsv41_shared_free();
attention_decode_score_split_graph_destroy_one(i);
routed_moe_decode_graph_destroy_one(i);
if (c->boundary_event) {
(void)cudaEventDestroy((cudaEvent_t)c->boundary_event);
c->boundary_event = NULL;
}
if (c->stream) {
(void)cudaStreamDestroy((cudaStream_t)c->stream);
c->stream = NULL;
}
if (c->cublas) {
(void)cublasDestroy((cublasHandle_t)c->cublas);
c->cublas = NULL;
c->cublas_ready = 0;
}
if (c->scratch) {
(void)cudaFree(c->scratch);
c->scratch = NULL;
c->scratch_bytes = 0;
}
if (g_tt_scratch && g_tt_scratch_device == c->device_id) {
(void)cudaFree(g_tt_scratch);
g_tt_scratch = NULL;
g_tt_scratch_bytes = 0;
g_tt_scratch_device = -1;
}
#ifdef DS4_CUDA_HAVE_MXF4
if (g_indexer_mxf4_scratch &&
g_indexer_mxf4_scratch_device == c->device_id) {
(void)cudaFree(g_indexer_mxf4_scratch);
g_indexer_mxf4_scratch = NULL;
g_indexer_mxf4_scratch_bytes = 0;
g_indexer_mxf4_scratch_device = -1;
}
#endif
}
for (int i = 0; i < DS4_MAX_GPUS; i++) {
for (int j = 0; j < DS4_MAX_GPUS; j++) {
if (g_xdev_bounce[i][j]) {
(void)cudaFreeHost(g_xdev_bounce[i][j]);
g_xdev_bounce[i][j] = NULL;
g_xdev_bounce_bytes[i][j] = 0;
}
}
}
cuda_stream_selected_cache_release();
cuda_stream_selected_stage_release();
g_n_gpus = 0;
g_device_is_spark = false;
g_cublas_ready = 0;
/* Per-device selective cache teardown (selective model cache). */
for (int d = 0; d < DS4_MAX_GPUS; d++) {
if (!g_dev_cache[d].present) continue;
int prev = -1;
(void)cudaGetDevice(&prev);
(void)cudaSetDevice(d);
if (g_dev_cache[d].base) (void)cudaFree(g_dev_cache[d].base);
g_dev_cache[d].base = NULL;
g_dev_cache[d].bytes = 0;
g_dev_cache[d].present = 0;
if (prev >= 0) (void)cudaSetDevice(prev);
}
g_cache_ranges.clear();
/* Continue with legacy global teardown below. */
cuda_model_range_release_all();
cuda_derived_range_release_all();
cuda_q8_f16_cache_release_all();
g_q8_f16_disabled_after_oom = 0;
g_q8_f16_budget_notice_printed = 0;
for (const cuda_q8_f32_range &r : g_q8_f32_ranges) {
(void)cudaFree(r.device_ptr);
}
g_q8_f32_ranges.clear();
g_q8_f32_by_offset.clear();
g_q8_f32_bytes = 0;
if (g_cuda_tmp) {
(void)cudaFree(g_cuda_tmp);
g_cuda_tmp = NULL;
g_cuda_tmp_bytes = 0;
}
for (size_t i = 0; i < 4; i++) {
if (g_model_stage_event[i]) {
(void)cudaEventDestroy(g_model_stage_event[i]);
g_model_stage_event[i] = NULL;
}
if (g_model_stage_raw[i]) {
(void)cudaFreeHost(g_model_stage_raw[i]);
g_model_stage_raw[i] = NULL;
g_model_stage[i] = NULL;
}
}
g_model_stage_bytes = 0;
if (g_model_upload_stream) {
(void)cudaStreamDestroy(g_model_upload_stream);
g_model_upload_stream = NULL;
}
if (g_model_device_owned && g_model_device_base) {
(void)cudaFree((void *)g_model_device_base);
}
if (g_model_registered && g_model_host_base) {
(void)cudaHostUnregister((void *)g_model_host_base);
}
g_model_host_base = NULL;
g_model_device_base = NULL;
g_model_registered_size = 0;
g_model_registered = 0;
g_model_device_owned = 0;
g_model_range_mapping_supported = 1;
g_model_hmm_direct = 0;
g_model_fd = -1;
if (g_model_direct_fd >= 0) {
(void)close(g_model_direct_fd);
g_model_direct_fd = -1;
}
g_model_direct_align = 1;
g_model_file_size = 0;
g_model_cache_full = 0;
if (g_model_prefetch_stream) {
(void)cudaStreamDestroy(g_model_prefetch_stream);
g_model_prefetch_stream = NULL;
}
}
__global__ static void fill_f32_kernel(float *x, uint64_t n, float v);
extern "C" int ds4_gpu_tensor_alloc_on(ds4_gpu_tensor *t, int device_id,
uint64_t bytes) {
if (!t) return 1;
if (device_id < 0 || device_id >= g_n_gpus) return 2;
if (bytes == 0) bytes = 1;
int ok = 0;
WITH_DEVICE(g_gpu[device_id].device_id) {
ok = cuda_ok(cudaMalloc(&t->ptr, (size_t)bytes), "tensor alloc");
}
if (!ok) return 3;
t->bytes = bytes;
t->owner = 1;
t->device_id = device_id;
g_gpu[device_id].used_bytes += bytes;
return 0;
}
/* Async D2D copy queued on the destination device's default stream —
* ordering against the producer comes from the caller's fence; the CPU
* does not block (unlike ds4_gpu_tensor_copy's sync cudaMemcpy). */
extern "C" int ds4_gpu_tensor_copy_async(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes) {
if (!dst || !src || bytes > dst->bytes || bytes > src->bytes) return 0;
if (bytes == 0) return 1;
return cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes,
cudaMemcpyDeviceToDevice, 0),
"tensor copy async");
}
extern "C" void ds4_gpu_tensor_free_in_place(ds4_gpu_tensor *t) {
if (!t) return;
int d = ds4_tensor_device_idx(t);
if (t->owner && t->ptr) {
WITH_DEVICE(g_gpu[d].device_id) {
(void)cudaFree(t->ptr);
}
}
t->ptr = NULL;
t->bytes = 0;
t->owner = 0;
}
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
if (ds4_gpu_tensor_alloc_on(t, 0, bytes) != 0) {
free(t);
return NULL;
}
return t;
}
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed(uint64_t bytes) {
if (bytes == 0) bytes = 1;
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
int ok = 0;
/* Managed memory is not device-bound, but we record device 0 so that
* subsequent ds4_gpu_tensor_free pairs with WITH_DEVICE(0) safely. */
WITH_DEVICE(g_gpu[0].device_id) {
ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes),
"managed tensor alloc");
}
if (!ok) { free(t); return NULL; }
t->bytes = bytes;
t->owner = 1;
t->device_id = 0;
return t;
}
/* Heap-allocated tensor on a specific logical tier.
*
* Mirrors the legacy ds4_gpu_tensor_alloc ABI (returns ds4_gpu_tensor *)
* with an explicit tier argument. Internally calls
* ds4_gpu_tensor_alloc_on on a freshly malloc'd struct.
*
* The legacy ds4_gpu_tensor_alloc(bytes) (above) calls
* ds4_gpu_tensor_alloc_on(t, 0, bytes); ds4_gpu_tensor_alloc_ptr_on(0,
* bytes) is byte-equivalent. Single-tier callers MAY remain on the
* legacy 1-arg helper; new multi-tier callers in ds4.c use _ptr_on. */
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_ptr_on(int tier, uint64_t bytes) {
if (tier < 0 || tier >= g_n_gpus) {
fprintf(stderr,
"ds4: ds4_gpu_tensor_alloc_ptr_on: bad tier %d (n_gpus=%d)\n",
tier, g_n_gpus);
return NULL;
}
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
if (ds4_gpu_tensor_alloc_on(t, tier, bytes) != 0) {
free(t);
return NULL;
}
return t;
}
/* Heap-allocated managed-memory tensor on a specific logical tier.
* Differs from ds4_gpu_tensor_alloc_managed only in stamping
* tier instead of 0. Used by the per-layer KV cache when tier !=0.
*
* Managed-memory paging behavior: cudaMallocManaged pages between
* devices on first-touch. In a single-tier pipeline the page lives on
* tier 0; in a multi-tier pipeline the layer's kernels run on the
* layer's tier so the page lives there after first-touch and stays
* unless another device touches it. Stamping tier matches the home
* device for free-time accounting. */
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_alloc_managed_on(int tier, uint64_t bytes) {
if (tier < 0 || tier >= g_n_gpus) {
fprintf(stderr,
"ds4: ds4_gpu_tensor_alloc_managed_on: bad tier %d (n_gpus=%d)\n",
tier, g_n_gpus);
return NULL;
}
if (bytes == 0) bytes = 1;
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
int ok = 0;
/* Run the cudaMallocManaged call under the home tier's device so the
* first-touch home matches the stamped device_id; the page itself can
* migrate freely under managed-memory semantics. */
WITH_DEVICE(g_gpu[tier].device_id) {
ok = cuda_ok(cudaMallocManaged(&t->ptr, (size_t)bytes),
"managed tensor alloc (tier)");
}
if (!ok) { free(t); return NULL; }
t->bytes = bytes;
t->owner = 1;
t->device_id = tier;
return t;
}
extern "C" int ds4_gpu_tensor_device(const ds4_gpu_tensor *t) {
return t ? t->device_id : -1;
}
static uint64_t cuda_managed_kv_reserve_bytes(uint64_t total_bytes) {
const uint64_t min_reserve = 8ull * 1073741824ull;
const uint64_t max_reserve = 40ull * 1073741824ull;
uint64_t reserve = total_bytes / 4u;
if (reserve < min_reserve) reserve = min_reserve;
if (reserve > max_reserve) reserve = max_reserve;
return reserve;
}
extern "C" int ds4_gpu_should_use_managed_kv_cache(uint64_t kv_cache_bytes, uint64_t context_bytes) {
if (kv_cache_bytes == 0) return 0;
/* Very large KV caches are where device-only cudaMalloc() can make a
* unified-memory machine unresponsive. Managed memory restores the old
* demand-paged behavior for this one long-lived allocation class only. */
const uint64_t huge_kv = 8ull * 1073741824ull;
if (kv_cache_bytes >= huge_kv) return 1;
const uint64_t large_context = 8ull * 1073741824ull;
if (context_bytes < large_context) return 0;
size_t free_b = 0;
size_t total_b = 0;
cudaError_t err = cudaMemGetInfo(&free_b, &total_b);
if (err != cudaSuccess) {
(void)cudaGetLastError();
return 0;
}
const uint64_t free_bytes = (uint64_t)free_b;
const uint64_t total_bytes = (uint64_t)total_b;
const uint64_t reserve_bytes = cuda_managed_kv_reserve_bytes(total_bytes);
if (context_bytes > free_bytes) return 1;
return free_bytes - context_bytes < reserve_bytes;
}
extern "C" ds4_gpu_tensor *ds4_gpu_tensor_view(const ds4_gpu_tensor *base, uint64_t offset, uint64_t bytes) {
if (!base || offset > base->bytes || bytes > base->bytes - offset) return NULL;
ds4_gpu_tensor *t = (ds4_gpu_tensor *)calloc(1, sizeof(*t));
if (!t) return NULL;
t->ptr = (char *)base->ptr + offset;
t->bytes = bytes;
t->owner = 0;
t->device_id = base->device_id; /* inherit owning device */
return t;
}
extern "C" void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) {
if (!tensor) return;
int d = ds4_tensor_device_idx(tensor);
if (tensor->owner && tensor->ptr) {
WITH_DEVICE(g_gpu[d].device_id) {
(void)cudaFree(tensor->ptr);
}
}
free(tensor);
}
extern "C" uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) {
return tensor ? tensor->bytes : 0;
}
extern "C" void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) {
if (!tensor) return NULL;
/* Full-device sync preserves legacy semantics. */
(void)cudaDeviceSynchronize();
return tensor->ptr;
}
extern "C" int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) {
if (!tensor || count > tensor->bytes / sizeof(float)) return 0;
if (count == 0) return 1;
int d = ds4_tensor_device_idx(tensor);
int ok = 0;
WITH_DEVICE(g_gpu[d].device_id) {
fill_f32_kernel<<<(count + 255u) / 256u, 256>>>((float *)tensor->ptr, count, value);
ok = cuda_ok(cudaGetLastError(), "tensor fill f32 launch");
}
return ok;
}
extern "C" int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) {
if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0;
int d = ds4_tensor_device_idx(tensor);
int ok = 0;
WITH_DEVICE(g_gpu[d].device_id) {
ok = cuda_ok(cudaMemcpy((char *)tensor->ptr + offset, data, (size_t)bytes,
cudaMemcpyHostToDevice),
"tensor write");
}
return ok;
}
extern "C" int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) {
if (!tensor || !data || offset > tensor->bytes || bytes > tensor->bytes - offset) return 0;
int d = ds4_tensor_device_idx(tensor);
int ok = 0;
WITH_DEVICE(g_gpu[d].device_id) {
ok = cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset, (size_t)bytes,
cudaMemcpyDeviceToHost),
"tensor read");
}
return ok;
}
extern "C" int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t bytes) {
if (!dst || !src || dst_offset > dst->bytes || src_offset > src->bytes ||
bytes > dst->bytes - dst_offset || bytes > src->bytes - src_offset) {
return 0;
}
if (bytes == 0) return 1;
/* Same-device fast path; for cross-device, callers should use
* ds4_gpu_tensor_copy_xdev. We still tolerate cross-device callers
* here by routing to D2D copy on the destination's device. */
int d = ds4_tensor_device_idx(dst);
int ok = 0;
WITH_DEVICE(g_gpu[d].device_id) {
if (g_decode_graph_capturing) {
ok = cuda_ok(cudaMemcpyAsync((char *)dst->ptr + dst_offset,
(const char *)src->ptr + src_offset,
(size_t)bytes,
cudaMemcpyDeviceToDevice,
cuda_decode_stream()),
"tensor copy");
} else {
ok = cuda_ok(cudaMemcpy((char *)dst->ptr + dst_offset,
(const char *)src->ptr + src_offset,
(size_t)bytes,
cudaMemcpyDeviceToDevice),
"tensor copy");
}
}
return ok;
}
__global__ static void moe_handoff_pack_kernel(
unsigned char *packed,
const float *ffn_norm,
const int32_t *selected,
const float *weights,
uint32_t n_embd,
uint32_t n_expert) {
const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
float *packed_norm = (float *)packed;
int32_t *packed_selected = (int32_t *)(packed + (uint64_t)n_embd * sizeof(float));
float *packed_weights = (float *)(packed + (uint64_t)n_embd * sizeof(float) +
(uint64_t)n_expert * sizeof(int32_t));
if (i < n_embd) packed_norm[i] = ffn_norm[i];
if (i < n_expert) {
packed_selected[i] = selected[i];
packed_weights[i] = weights[i];
}
}
extern "C" int ds4_gpu_moe_handoff_pack_tensor(
ds4_gpu_tensor *packed,
const ds4_gpu_tensor *ffn_norm,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_embd,
uint32_t n_expert) {
if (!packed || !ffn_norm || !selected || !weights ||
n_embd == 0 || n_expert == 0) {
return 0;
}
const uint64_t bytes = (uint64_t)n_embd * sizeof(float) +
(uint64_t)n_expert * sizeof(int32_t) +
(uint64_t)n_expert * sizeof(float);
if (packed->bytes < bytes ||
ffn_norm->bytes < (uint64_t)n_embd * sizeof(float) ||
selected->bytes < (uint64_t)n_expert * sizeof(int32_t) ||
weights->bytes < (uint64_t)n_expert * sizeof(float)) {
return 0;
}
const uint32_t n = n_embd > n_expert ? n_embd : n_expert;
moe_handoff_pack_kernel<<<(n + 255u) / 256u, 256>>>(
(unsigned char *)packed->ptr,
(const float *)ffn_norm->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
n_embd,
n_expert);
return cuda_ok(cudaGetLastError(), "moe handoff pack launch");
}
/* Cross-device copy primitive. Path selection (highest priority first):
* DS4_FORCE_HOST_BOUNCE=1 -> always pinned-host bounce
* DS4_FORCE_CUDA_PEER=1 -> always cudaMemcpyPeerAsync (manual-testing
* override; bypasses g_gpu_peer_ok)
* otherwise -> peer if validation passed at init, else
* pinned-host bounce. */
static int ds4_gpu_tensor_copy_xdev_impl(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes,
bool order_dst_before_write) {
if (!dst || !src) return 0;
if (bytes == 0) return 1;
if (bytes > dst->bytes || bytes > src->bytes) return 0;
int sd = ds4_tensor_device_idx(src);
int dd = ds4_tensor_device_idx(dst);
/* Same-device fast path. */
if (sd == dd) {
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, bytes,
cudaMemcpyDeviceToDevice, s),
"xdev same-device copy");
if (ok && g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s), "xdev same-device sync");
}
}
return ok;
}
int peer = g_gpu_peer_ok[sd][dd];
if (g_xdev_force_cuda_peer) peer = 1;
if (g_xdev_force_host_bounce) peer = 0;
if (peer) {
int ok = 0;
if (order_dst_before_write) {
WITH_DEVICE(g_gpu[dd].device_id) {
cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream;
cudaEvent_t e2 = (cudaEvent_t)g_gpu[dd].boundary_event;
ok = cuda_ok(cudaEventRecord(e2, s2), "peer dst-ready event record");
}
if (!ok) return 0;
} else {
ok = 1;
}
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event;
if (order_dst_before_write) {
ok = cuda_ok(cudaStreamWaitEvent(s, (cudaEvent_t)g_gpu[dd].boundary_event, 0),
"peer src wait dst-ready");
}
if (ok) ok = cuda_ok(cudaMemcpyPeerAsync(
dst->ptr, g_gpu[dd].device_id,
src->ptr, g_gpu[sd].device_id,
bytes, s),
"peer copy");
if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer event record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream;
(void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0);
if (g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s2), "peer dst sync");
}
}
return ok;
}
/* Per-pair pinned-host bounce buffer. */
if (g_xdev_bounce_bytes[sd][dd] < bytes) {
if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]);
if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes),
"bounce alloc")) return 0;
g_xdev_bounce_bytes[sd][dd] = bytes;
}
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event;
ok = cuda_ok(cudaMemcpyAsync(g_xdev_bounce[sd][dd], src->ptr, bytes,
cudaMemcpyDeviceToHost, s),
"bounce d2h");
if (ok) ok = cuda_ok(cudaEventRecord(e, s), "bounce event record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream;
(void)cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0);
ok = cuda_ok(cudaMemcpyAsync(dst->ptr, g_xdev_bounce[sd][dd], bytes,
cudaMemcpyHostToDevice, s2),
"bounce h2d");
if (ok) ok = cuda_ok(cudaStreamSynchronize(s2), "bounce dst sync");
}
return ok;
}
extern "C" int ds4_gpu_tensor_copy_xdev(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes) {
return ds4_gpu_tensor_copy_xdev_impl(dst, src, bytes, false);
}
static int ds4_gpu_tensor_copy_xdev_default_impl(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes) {
if (!dst || !src || bytes > dst->bytes || bytes > src->bytes) return 0;
if (bytes == 0u) return 1;
const int sd = ds4_tensor_device_idx(src);
const int dd = ds4_tensor_device_idx(dst);
if (sd == dd) {
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
ok = cuda_ok(cudaMemcpyAsync(dst->ptr, src->ptr, (size_t)bytes,
cudaMemcpyDeviceToDevice, 0),
"default-stream same-device copy");
}
return ok;
}
int peer = g_gpu_peer_ok[sd][dd];
if (g_xdev_force_cuda_peer) peer = 1;
if (g_xdev_force_host_bounce) peer = 0;
if (peer) {
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
ok = cuda_ok(cudaMemcpyPeerAsync(
dst->ptr, g_gpu[dd].device_id,
src->ptr, g_gpu[sd].device_id,
(size_t)bytes, 0),
"default-stream peer copy");
if (ok) {
ok = cuda_ok(cudaEventRecord(
(cudaEvent_t)g_gpu[sd].boundary_event, 0),
"default-stream peer event record");
}
}
if (ok) {
WITH_DEVICE(g_gpu[dd].device_id) {
ok = cuda_ok(cudaStreamWaitEvent(
0,
(cudaEvent_t)g_gpu[sd].boundary_event,
0),
"default-stream peer destination wait");
}
}
return ok;
}
if (g_xdev_bounce_bytes[sd][dd] < bytes) {
if (g_xdev_bounce[sd][dd]) (void)cudaFreeHost(g_xdev_bounce[sd][dd]);
if (!cuda_ok(cudaMallocHost(&g_xdev_bounce[sd][dd], (size_t)bytes),
"default-stream bounce alloc")) return 0;
g_xdev_bounce_bytes[sd][dd] = bytes;
}
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
ok = cuda_ok(cudaMemcpy(g_xdev_bounce[sd][dd], src->ptr, (size_t)bytes,
cudaMemcpyDeviceToHost),
"default-stream bounce d2h");
}
if (ok) {
WITH_DEVICE(g_gpu[dd].device_id) {
ok = cuda_ok(cudaMemcpy(dst->ptr, g_xdev_bounce[sd][dd],
(size_t)bytes, cudaMemcpyHostToDevice),
"default-stream bounce h2d");
}
}
return ok;
}
extern "C" int ds4_gpu_tensor_copy_xdev_default(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes) {
return ds4_gpu_tensor_copy_xdev_default_impl(dst, src, bytes);
}
extern "C" int ds4_gpu_tensor_copy_xdev3_default_dst(
ds4_gpu_tensor *dst0,
const ds4_gpu_tensor *src0,
uint64_t bytes0,
ds4_gpu_tensor *dst1,
const ds4_gpu_tensor *src1,
uint64_t bytes1,
ds4_gpu_tensor *dst2,
const ds4_gpu_tensor *src2,
uint64_t bytes2) {
ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2};
const ds4_gpu_tensor *srcs[3] = {src0, src1, src2};
const uint64_t sizes[3] = {bytes0, bytes1, bytes2};
int sd = -1;
int dd = -1;
for (int i = 0; i < 3; i++) {
if (sizes[i] == 0u) continue;
if (!dsts[i] || !srcs[i] || sizes[i] > dsts[i]->bytes ||
sizes[i] > srcs[i]->bytes) {
return 0;
}
const int this_sd = ds4_tensor_device_idx(srcs[i]);
const int this_dd = ds4_tensor_device_idx(dsts[i]);
if (sd < 0) {
sd = this_sd;
dd = this_dd;
} else if (sd != this_sd || dd != this_dd) {
return 0;
}
}
if (sd < 0) return 1;
if (sd == dd) {
int ok = 1;
WITH_DEVICE(g_gpu[sd].device_id) {
for (int i = 0; ok && i < 3; i++) {
if (sizes[i] == 0u) continue;
ok = cuda_ok(cudaMemcpyAsync(
dsts[i]->ptr, srcs[i]->ptr, (size_t)sizes[i],
cudaMemcpyDeviceToDevice, 0),
"grouped default same-device copy");
}
}
return ok;
}
int peer = g_gpu_peer_ok[dd][sd];
if (g_xdev_force_cuda_peer) peer = 1;
if (g_xdev_force_host_bounce) peer = 0;
if (!peer) {
for (int i = 0; i < 3; i++) {
if (sizes[i] != 0u &&
!ds4_gpu_tensor_copy_xdev_default_impl(
dsts[i], srcs[i], sizes[i])) {
return 0;
}
}
return 1;
}
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
ok = cuda_ok(cudaEventRecord(
(cudaEvent_t)g_gpu[sd].boundary_event, 0),
"grouped default source-ready record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
ok = cuda_ok(cudaStreamWaitEvent(
0, (cudaEvent_t)g_gpu[sd].boundary_event, 0),
"grouped default destination wait");
for (int i = 0; ok && i < 3; i++) {
if (sizes[i] == 0u) continue;
ok = cuda_ok(cudaMemcpyPeerAsync(
dsts[i]->ptr, g_gpu[dd].device_id,
srcs[i]->ptr, g_gpu[sd].device_id,
(size_t)sizes[i], 0),
"grouped destination-stream peer copy");
}
}
return ok;
}
extern "C" int ds4_gpu_tensor_copy_xdev3(ds4_gpu_tensor *dst0,
const ds4_gpu_tensor *src0,
uint64_t bytes0,
ds4_gpu_tensor *dst1,
const ds4_gpu_tensor *src1,
uint64_t bytes1,
ds4_gpu_tensor *dst2,
const ds4_gpu_tensor *src2,
uint64_t bytes2) {
ds4_gpu_tensor *dsts[3] = {dst0, dst1, dst2};
const ds4_gpu_tensor *srcs[3] = {src0, src1, src2};
uint64_t bytes[3] = {bytes0, bytes1, bytes2};
int first = -1;
for (int i = 0; i < 3; i++) {
if (bytes[i] == 0) continue;
if (!dsts[i] || !srcs[i] ||
bytes[i] > dsts[i]->bytes || bytes[i] > srcs[i]->bytes) {
return 0;
}
if (first < 0) first = i;
}
if (first < 0) return 1;
const int sd = ds4_tensor_device_idx(srcs[first]);
const int dd = ds4_tensor_device_idx(dsts[first]);
for (int i = first + 1; i < 3; i++) {
if (bytes[i] == 0) continue;
if (ds4_tensor_device_idx(srcs[i]) != sd ||
ds4_tensor_device_idx(dsts[i]) != dd) {
int ok = 1;
for (int j = 0; ok && j < 3; j++) {
if (bytes[j] == 0) continue;
ok = ds4_gpu_tensor_copy_xdev(dsts[j], srcs[j], bytes[j]);
}
return ok;
}
}
if (sd == dd) {
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
ok = 1;
for (int i = 0; ok && i < 3; i++) {
if (bytes[i] == 0) continue;
ok = cuda_ok(cudaMemcpyAsync(dsts[i]->ptr, srcs[i]->ptr, bytes[i],
cudaMemcpyDeviceToDevice, s),
"xdev3 same-device copy");
}
if (ok && g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s), "xdev3 same-device sync");
}
}
return ok;
}
int peer = g_gpu_peer_ok[sd][dd];
if (g_xdev_force_cuda_peer) peer = 1;
if (g_xdev_force_host_bounce) peer = 0;
if (!peer) {
int ok = 1;
for (int i = 0; ok && i < 3; i++) {
if (bytes[i] == 0) continue;
ok = ds4_gpu_tensor_copy_xdev(dsts[i], srcs[i], bytes[i]);
}
return ok;
}
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event;
ok = 1;
for (int i = 0; ok && i < 3; i++) {
if (bytes[i] == 0) continue;
ok = cuda_ok(cudaMemcpyPeerAsync(
dsts[i]->ptr, g_gpu[dd].device_id,
srcs[i]->ptr, g_gpu[sd].device_id,
bytes[i], s),
"peer copy3");
}
if (ok) ok = cuda_ok(cudaEventRecord(e, s), "peer copy3 event record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream;
ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0),
"peer copy3 dst wait");
if (ok && g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s2), "peer copy3 dst sync");
}
}
return ok;
}
extern "C" int ds4_gpu_tensor_copy_xdev_ordered(ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint64_t bytes) {
return ds4_gpu_tensor_copy_xdev_impl(dst, src, bytes, true);
}
extern "C" int ds4_gpu_tensor_wait_xdev(const ds4_gpu_tensor *src, int dst_tier) {
if (!src) return 0;
if (dst_tier < 0 || dst_tier >= g_n_gpus) return 0;
int sd = ds4_tensor_device_idx(src);
int dd = dst_tier;
if (sd == dd) return 1;
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[sd].stream;
cudaEvent_t e = (cudaEvent_t)g_gpu[sd].boundary_event;
ok = cuda_ok(cudaEventRecord(e, s), "xdev wait source event record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
cudaStream_t s2 = (cudaStream_t)g_gpu[dd].stream;
ok = cuda_ok(cudaStreamWaitEvent(s2, (cudaEvent_t)g_gpu[sd].boundary_event, 0),
"xdev wait destination wait");
if (ok && g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s2), "xdev wait dst sync");
}
}
return ok;
}
extern "C" int ds4_gpu_tensor_wait_xdev_default(
const ds4_gpu_tensor *src,
int dst_tier) {
if (!src || dst_tier < 0 || dst_tier >= g_n_gpus) return 0;
const int sd = ds4_tensor_device_idx(src);
const int dd = dst_tier;
if (sd == dd) return 1;
int ok = 0;
WITH_DEVICE(g_gpu[sd].device_id) {
ok = cuda_ok(cudaEventRecord(
(cudaEvent_t)g_gpu[sd].boundary_event, 0),
"default xdev wait source event record");
}
if (!ok) return 0;
WITH_DEVICE(g_gpu[dd].device_id) {
ok = cuda_ok(cudaStreamWaitEvent(
0,
(cudaEvent_t)g_gpu[sd].boundary_event,
0),
"default xdev wait destination wait");
}
return ok;
}
extern "C" int ds4_gpu_q8_cache_suppressed(void) {
return g_q8_cache_suppressed;
}
extern "C" void ds4_gpu_set_q8_cache_suppressed(int suppressed) {
g_q8_cache_suppressed = suppressed ? 1 : 0;
}
__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap);
extern "C" int ds4_gpu_pack_slot_rows_f32_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *slots,
uint32_t n_rows,
uint32_t width,
uint32_t n_slots,
uint32_t slot_cap) {
uint64_t slot_rows = 0;
uint64_t slot_elems = 0;
uint64_t out_rows = 0;
uint64_t out_elems = 0;
if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 ||
slot_cap == 0 || n_rows > slot_cap ||
(uint64_t)n_slots > UINT64_MAX / slot_cap ||
(slot_rows = (uint64_t)n_slots * slot_cap) > UINT64_MAX / width ||
(slot_elems = slot_rows * width) > UINT64_MAX / sizeof(float) ||
(uint64_t)n_rows > UINT64_MAX / n_slots ||
(out_rows = (uint64_t)n_rows * n_slots) > UINT64_MAX / width ||
(out_elems = out_rows * width) > UINT64_MAX / sizeof(float) ||
slots->bytes < slot_elems * sizeof(float) ||
out->bytes < out_elems * sizeof(float)) {
return 0;
}
const uint64_t blocks = (out_elems + 255u) / 256u;
if (blocks > UINT32_MAX) return 0;
pack_slot_rows_f32_kernel<<<(unsigned)blocks, 256>>>(
(float *)out->ptr,
(const float *)slots->ptr,
n_rows,
width,
n_slots,
slot_cap);
return cuda_ok(cudaGetLastError(), "pack_slot_rows_f32 launch");
}
extern "C" int ds4_gpu_begin_commands(void) { return 1; }
extern "C" int ds4_gpu_flush_commands(void) { return cuda_ok(cudaDeviceSynchronize(), "flush"); }
extern "C" int ds4_gpu_end_commands(void) {
if (g_cuda_end_stream_sync) {
return cuda_ok(cudaStreamSynchronize(0), "end commands stream");
}
return cuda_ok(cudaDeviceSynchronize(), "end commands");
}
extern "C" int ds4_gpu_synchronize(void) { return cuda_ok(cudaDeviceSynchronize(), "synchronize"); }
extern "C" int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) {
if (!model_map || model_size == 0) return 0;
if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1;
cuda_stream_selected_cache_release();
cuda_model_range_release_all();
cuda_q8_f16_cache_release_all();
g_q8_f16_disabled_after_oom = 0;
g_q8_f16_budget_notice_printed = 0;
for (const cuda_q8_f32_range &r : g_q8_f32_ranges) {
(void)cudaFree(r.device_ptr);
}
g_q8_f32_ranges.clear();
g_q8_f32_by_offset.clear();
g_q8_f32_bytes = 0;
if (g_model_device_owned && g_model_device_base) {
(void)cudaFree((void *)g_model_device_base);
g_model_device_owned = 0;
}
if (g_model_registered && g_model_host_base) {
(void)cudaHostUnregister((void *)g_model_host_base);
g_model_registered = 0;
}
g_model_host_base = model_map;
g_model_device_base = (const char *)model_map;
g_model_registered_size = model_size;
g_model_range_mapping_supported = 1;
g_model_hmm_direct = 0;
g_model_cache_full = 0;
if (g_model_fd >= 0 && g_model_fd_host_base == NULL) {
g_model_fd_host_base = model_map;
}
/* Never pin or copy the complete model in SSD mode, even on integrated GPUs. */
if (g_ssd_streaming_mode) return 1;
const char *copy_env = getenv("DS4_CUDA_COPY_MODEL");
if (copy_env && copy_env[0]) {
void *dev = NULL;
const double t0 = clock() / (double)CLOCKS_PER_SEC;
cudaError_t err = cudaMalloc(&dev, (size_t)model_size);
if (err == cudaSuccess) {
fprintf(stderr, "ds4: CUDA copying %.2f GiB model to device memory\n",
(double)model_size / 1073741824.0);
err = cudaMemcpy(dev, model_map, (size_t)model_size, cudaMemcpyHostToDevice);
if (err == cudaSuccess) {
g_model_device_base = (const char *)dev;
g_model_device_owned = 1;
const double t1 = clock() / (double)CLOCKS_PER_SEC;
fprintf(stderr, "ds4: CUDA model copy complete in %.3fs\n", t1 - t0);
return 1;
}
fprintf(stderr, "ds4: CUDA model copy failed: %s\n", cudaGetErrorString(err));
(void)cudaFree(dev);
(void)cudaGetLastError();
} else {
fprintf(stderr, "ds4: CUDA model allocation skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
}
}
if (cuda_integrated_artifact_map(model_map)) {
fprintf(stderr,
"ds4: CUDA aligned artifacts replace expert residency; "
"leaving the %.2f GiB model mmap unpinned\n",
(double)model_size / 1073741824.0);
return 1;
}
cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size,
cudaHostRegisterMapped | cudaHostRegisterReadOnly);
if (err == cudaSuccess) {
void *dev = NULL;
err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0);
if (err == cudaSuccess && dev) {
g_model_device_base = (const char *)dev;
g_model_registered = 1;
fprintf(stderr, "ds4: CUDA registered %.2f GiB model mapping for device access\n",
(double)model_size / 1073741824.0);
} else {
fprintf(stderr, "ds4: CUDA host registration pointer lookup failed: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
}
} else {
fprintf(stderr, "ds4: CUDA host registration skipped: %s\n", cudaGetErrorString(err));
(void)cudaGetLastError();
}
return 1;
}
extern "C" int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes) {
(void)max_tensor_bytes;
if (!ds4_gpu_register_model_map_no_copy(model_map, model_size)) return 0;
if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL &&
!cuda_model_copy_chunked(model_map, model_size, map_offset, map_size)) {
(void)cuda_model_prefetch_range(model_map, model_size, map_offset, map_size);
}
return 1;
}
extern "C" int ds4_gpu_set_aux_model_map_range(
const void *model_map,
uint64_t model_size,
uint64_t map_offset,
uint64_t map_size) {
if (!model_map || model_size == 0 || map_size == 0 ||
map_offset > model_size || map_size > model_size - map_offset) {
return 0;
}
if (cuda_model_range_is_cached(model_map, map_offset, map_size)) return 1;
int current_device = 0;
int integrated = 0;
int pageable = 0;
if (getenv("DS4_CUDA_AUX_FORCE_COPY") == NULL &&
cudaGetDevice(&current_device) == cudaSuccess &&
cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated,
current_device) ==
cudaSuccess &&
cudaDeviceGetAttribute(&pageable, cudaDevAttrPageableMemoryAccess,
current_device) == cudaSuccess &&
integrated && pageable) {
g_model_ranges.push_back({
model_map, map_offset, map_size,
(char *)model_map + map_offset,
NULL, NULL, 0, 0, 0, 1});
fprintf(stderr,
"ds4: CUDA directly mapped %.2f GiB auxiliary model\n",
(double)map_size / 1073741824.0);
return 1;
}
void *device = NULL;
cudaError_t err = cudaMalloc(&device, (size_t)map_size);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA vision map allocation failed (%.2f GiB): %s\n",
(double)map_size / 1073741824.0,
cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
const uint64_t chunk = 64ull * 1024ull * 1024ull;
const char *source = (const char *)model_map + map_offset;
for (uint64_t copied = 0; copied < map_size; copied += chunk) {
const uint64_t bytes = map_size - copied < chunk ?
map_size - copied : chunk;
err = cudaMemcpy((char *)device + copied, source + copied,
(size_t)bytes, cudaMemcpyHostToDevice);
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: CUDA vision map copy failed at %.2f/%.2f GiB: %s\n",
(double)copied / 1073741824.0,
(double)map_size / 1073741824.0,
cudaGetErrorString(err));
(void)cudaFree(device);
(void)cudaGetLastError();
return 0;
}
}
g_model_ranges.push_back({model_map, map_offset, map_size,
(char *)device, NULL, NULL, 0, 0, 0});
g_model_range_bytes += map_size;
fprintf(stderr, "ds4: CUDA mapped %.2f GiB auxiliary model\n",
(double)map_size / 1073741824.0);
return 1;
}
/* Register the mmap'd host model pointer for selective-cache lookups WITHOUT
* triggering any device-side copy. Used by multi-GPU placement scaffolding's
* multi-tier path so DS4_CUDA_COPY_MODEL cannot reintroduce a full-model
* copy that defeats the per-device selective cache.
*
* This is the no-copy subset of ds4_gpu_set_model_map: same bookkeeping
* for the host pointer plus cudaHostRegister, but skipping the
* DS4_CUDA_COPY_MODEL branch that allocates and copies the entire model. */
static int cuda_register_model_map(const void *model_map, uint64_t model_size,
bool register_all) {
if (!model_map || model_size == 0) return 0;
if (g_model_host_base == model_map && g_model_registered_size == model_size) return 1;
cuda_stream_selected_cache_release();
cuda_model_range_release_all();
cuda_q8_f16_cache_release_all();
g_q8_f16_disabled_after_oom = 0;
g_q8_f16_budget_notice_printed = 0;
for (const cuda_q8_f32_range &r : g_q8_f32_ranges) {
(void)cudaFree(r.device_ptr);
}
g_q8_f32_ranges.clear();
g_q8_f32_by_offset.clear();
g_q8_f32_bytes = 0;
if (g_model_device_owned && g_model_device_base) {
(void)cudaFree((void *)g_model_device_base);
g_model_device_owned = 0;
}
if (g_model_registered && g_model_host_base) {
(void)cudaHostUnregister((void *)g_model_host_base);
g_model_registered = 0;
}
g_model_host_base = model_map;
g_model_device_base = (const char *)model_map;
g_model_registered_size = model_size;
g_model_range_mapping_supported = 1;
g_model_hmm_direct = 0;
g_model_cache_full = 0;
if (g_model_fd >= 0 && g_model_fd_host_base == NULL) {
g_model_fd_host_base = model_map;
}
/* A sharded mapping must never pin the unowned expert ranges. */
if (!register_all) return 1;
if (cuda_integrated_artifact_map(model_map)) {
fprintf(stderr,
"ds4: CUDA aligned artifacts replace expert residency; "
"leaving the %.2f GiB model mmap unpinned\n",
(double)model_size / 1073741824.0);
return 1;
}
cudaError_t err = cudaHostRegister((void *)model_map, (size_t)model_size,
cudaHostRegisterMapped | cudaHostRegisterReadOnly);
if (err == cudaSuccess) {
void *dev = NULL;
err = cudaHostGetDevicePointer(&dev, (void *)model_map, 0);
if (err == cudaSuccess && dev) {
g_model_device_base = (const char *)dev;
g_model_registered = 1;
fprintf(stderr,
"ds4: CUDA (no-copy) registered %.2f GiB model mapping for multi-tier selective cache\n",
(double)model_size / 1073741824.0);
} else {
fprintf(stderr,
"ds4: CUDA (no-copy) host registration pointer lookup failed: %s\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
}
} else {
fprintf(stderr,
"ds4: CUDA (no-copy) host registration skipped: %s\n",
cudaGetErrorString(err));
(void)cudaGetLastError();
}
return 1;
}
extern "C" int ds4_gpu_register_model_map_no_copy(const void *model_map, uint64_t model_size) {
return cuda_register_model_map(model_map, model_size, true);
}
/* Set the current CUDA device by LOGICAL tier index (0..g_n_gpus-1).
* Maps to the physical CUDA device id stored in g_gpu[].device_id.
* Added for multi-GPU placement scaffolding (multi-GPU CLI); first executed by
* multi-GPU execution (follow-up). */
extern "C" int ds4_gpu_set_current_device(int logical_tier) {
if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1;
if (!g_cuda_no_setdevice_cache && g_current_logical_tier == logical_tier) {
return 0;
}
if (cudaSetDevice(g_gpu[logical_tier].device_id) == cudaSuccess) {
g_current_logical_tier = logical_tier;
return 0;
}
g_current_logical_tier = -1;
return -1;
}
/* Fenced device switch for sequential cross-device pipelines (GLM
* per-layer placement): work queued on the next device's default stream
* waits for everything queued so far on the previous device's default
* stream. Async — no host sync. Falls back to a plain switch when the
* device does not change. */
extern "C" int ds4_gpu_set_current_device_fenced(int logical_tier) {
if (logical_tier < 0 || logical_tier >= g_n_gpus) return -1;
static cudaEvent_t fence_ev[DS4_MAX_GPUS];
/* Resolve the ACTUAL current device: WITH_DEVICE blocks and direct
* cudaSetDevice calls can leave g_current_logical_tier stale, and a
* false "already there" here strands work on the wrong device. */
int cur_dev = -1;
(void)cudaGetDevice(&cur_dev);
int prev = -1;
for (int t = 0; t < g_n_gpus; t++) {
if (g_gpu[t].device_id == cur_dev) { prev = t; break; }
}
if (getenv("DS4_GLM_FENCE_TRACE")) {
fprintf(stderr, "ds4: fenced switch %d -> %d\n", prev, logical_tier);
}
if (prev == logical_tier) {
g_current_logical_tier = logical_tier;
return 0;
}
if (prev >= 0 && prev < g_n_gpus && prev != logical_tier) {
if (cudaSetDevice(g_gpu[prev].device_id) != cudaSuccess) return -1;
if (!fence_ev[prev] &&
cudaEventCreateWithFlags(&fence_ev[prev],
cudaEventDisableTiming) != cudaSuccess) {
fence_ev[prev] = NULL;
}
if (fence_ev[prev]) {
(void)cudaEventRecord(fence_ev[prev], 0);
}
if (cudaSetDevice(g_gpu[logical_tier].device_id) != cudaSuccess) {
g_current_logical_tier = -1;
return -1;
}
g_current_logical_tier = logical_tier;
if (fence_ev[prev]) {
(void)cudaStreamWaitEvent(0, fence_ev[prev], 0);
}
return 0;
}
return ds4_gpu_set_current_device(logical_tier);
}
/* =========================================================================
* Per-device selective model cache (selective model cache).
*
* ds4_gpu_device_cache_tensors copies the listed source ranges from the
* host mmap onto device_id's selective slab and appends sorted lookup
* entries. The legacy chunked-copy machinery (cuda_model_range_*) is
* NOT disturbed — it continues to drive all existing callers. New
* lookups fall back to it when no selective entry covers the range.
*
* Caller-context preference for overlap: when the same source range is
* cached on multiple devices, ds4_gpu_lookup_cache returns the entry
* whose device matches cudaGetDevice().
* ========================================================================= */
extern "C" int ds4_gpu_device_cache_tensors(int device_id,
const ds4_tensor_range *ranges,
int n_ranges) {
if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1;
if (n_ranges < 0 || (!ranges && n_ranges > 0)) return 2;
if (n_ranges == 0) return 0;
if (!g_model_host_base || g_model_registered_size == 0) return 3;
/* Validate ranges against the mmap'd model bounds; reject ranges
* that overflow or extend past the mapped region. Done in a
* separate pass before any allocation so a bad input doesn't
* partially grow the slab. */
uint64_t want_bytes = 0;
for (int i = 0; i < n_ranges; i++) {
if (ranges[i].target_device != device_id) continue;
const uint64_t off = ranges[i].source_offset;
const uint64_t nb = ranges[i].bytes;
/* Overflow-safe upper bound: off + nb must not exceed model
* size, and the sum must not wrap. */
if (nb == 0) continue;
if (off > g_model_registered_size) return 8;
if (nb > g_model_registered_size - off) return 9;
/* Accumulate into want_bytes with overflow check. */
if (want_bytes > UINT64_MAX - nb) return 10;
want_bytes += nb;
}
if (want_bytes == 0) return 0;
cuda_device_cache &c = g_dev_cache[device_id];
int prev_device = -1;
if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1;
if (cudaSetDevice(device_id) != cudaSuccess) return 4;
/* Allocate or grow the slab via cudaMalloc + d2d copy. */
void *new_base = NULL;
size_t new_bytes = c.bytes + want_bytes;
/* Refuse cleanly before cudaMalloc if the device clearly cannot hold
* the slab. The multi-tier packer reserves per-tier runtime scratch
* before placing tensors, but it cannot predict the cudaMalloc
* allocator's overhead (alignment, fragmentation after CUDA context
* init, default driver-side reservations). On a borderline budget
* that overhead pushes a "fits-by-packer-math" layout past the actual
* free pool and the cudaMalloc below OOMs after the engine already
* committed to the layout — same silent-late-OOM failure mode the
* upfront refusal path was added to eliminate. Catch it here too. */
{
size_t free_b = 0, total_b = 0;
if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess) {
/* free_b already excludes the existing slab (it's still
* allocated), so the additional cudaMalloc only needs
* new_bytes free — not new_bytes + c.bytes. The old slab is
* freed AFTER the d2d copy succeeds. 2 GiB safety covers what
* the engine will allocate AFTER the cache slab in the same
* session_create: per-tier graph scratch (the planner can't
* predict its cumulative cudaMalloc alignment overhead),
* cuBLAS workspace beyond the 64 MiB the packer already
* reserves, and driver-side allocator slack. Without this
* headroom a borderline budget that fits the slab itself can
* still OOM at the per-tier tensor allocations a few moments
* later — same silent-late-OOM failure mode, one layer up. */
const size_t safety = (size_t)2ull * 1024ull * 1024ull * 1024ull;
const size_t need = new_bytes + safety;
if (need > free_b) {
fprintf(stderr,
"ds4: device cache slab needs %.2f GiB on device %d "
"but only %.2f GiB free (slab=%.2f GiB + %.2f GiB safety). "
"Lower --gpu-vram / --ctx-max, or use --gpu-vram auto on "
"a host with more free VRAM. Refusing upfront to avoid "
"late OOM at cudaMalloc.\n",
(double)need / 1073741824.0,
device_id,
(double)free_b / 1073741824.0,
(double)new_bytes / 1073741824.0,
(double)safety / 1073741824.0);
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 5;
}
}
/* If cudaMemGetInfo itself failed, fall through; cudaMalloc's own
* error path still catches the late case, just with a less helpful
* message. */
}
if (!cuda_ok(cudaMalloc(&new_base, new_bytes), "device cache alloc")) {
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 5;
}
if (c.present && c.bytes > 0) {
cudaError_t e = cudaMemcpy(new_base, c.base, c.bytes,
cudaMemcpyDeviceToDevice);
if (e != cudaSuccess) {
cuda_ok(e, "device cache grow d2d");
(void)cudaFree(new_base);
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 6;
}
/* Re-base existing entries on this device. */
char *old_base = (char *)c.base;
char *grown = (char *)new_base;
for (size_t k = 0; k < g_cache_ranges.size(); k++) {
if (g_cache_ranges[k].device_id == device_id) {
g_cache_ranges[k].device_ptr =
grown + ((char *)g_cache_ranges[k].device_ptr - old_base);
}
}
(void)cudaFree(c.base);
}
c.base = new_base;
c.bytes = new_bytes;
c.present = 1;
/* Copy ranges and append entries. */
const char *host_base = (const char *)g_model_host_base;
size_t write_off = c.bytes - want_bytes;
for (int i = 0; i < n_ranges; i++) {
if (ranges[i].target_device != device_id) continue;
char *dev_ptr = (char *)c.base + write_off;
cudaError_t e = cudaMemcpy(dev_ptr,
host_base + ranges[i].source_offset,
(size_t)ranges[i].bytes,
cudaMemcpyHostToDevice);
if (e != cudaSuccess) {
cuda_ok(e, "device cache range h2d");
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 7;
}
cache_range_entry ent;
ent.source_offset = ranges[i].source_offset;
ent.bytes = ranges[i].bytes;
ent.device_id = device_id;
ent.device_ptr = dev_ptr;
g_cache_ranges.push_back(ent);
write_off += ranges[i].bytes;
}
/* Keep sorted by source_offset for binary-search lookup. */
std::sort(g_cache_ranges.begin(), g_cache_ranges.end(),
[](const cache_range_entry &a, const cache_range_entry &b) {
if (a.source_offset != b.source_offset)
return a.source_offset < b.source_offset;
return a.device_id < b.device_id;
});
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 0;
}
/* Install support-model tensor ranges into device_id's strict cache,
* copying from the registered support map and keying entries at
* source_offset + bias. Standalone slab (does not touch the main cache
* slab growth path). */
extern "C" int ds4_gpu_device_cache_support_tensors(int device_id,
int entry_device_id,
const ds4_tensor_range *ranges,
int n_ranges,
int from_main_map) {
if (device_id < 0 || device_id >= DS4_MAX_GPUS) return 1;
if (entry_device_id < 0 || entry_device_id >= DS4_MAX_GPUS) return 1;
if (n_ranges <= 0 || !ranges) return 2;
const char *src_base;
uint64_t src_size;
uint64_t key_bias;
if (from_main_map) {
/* Auxiliary main-model ranges (e.g. the embedding bucket for the
* DSpark executor tier): standalone slab, unbiased offsets. */
src_base = (const char *)g_model_host_base;
src_size = g_model_registered_size;
key_bias = 0;
} else {
src_base = (const char *)g_support_host_base;
src_size = g_support_host_size;
key_bias = g_support_offset_bias;
if (key_bias == 0) return 3;
}
if (!src_base || src_size == 0) return 3;
uint64_t want = 0;
for (int i = 0; i < n_ranges; i++) {
const uint64_t off = ranges[i].source_offset;
const uint64_t nb = ranges[i].bytes;
if (nb == 0) continue;
if (off > src_size || nb > src_size - off) return 8;
if (want > UINT64_MAX - nb) return 9;
want += nb;
}
if (want == 0) return 0;
int prev_device = -1;
if (cudaGetDevice(&prev_device) != cudaSuccess) prev_device = -1;
if (cudaSetDevice(device_id) != cudaSuccess) return 4;
void *base = NULL;
if (!cuda_ok(cudaMalloc(&base, (size_t)want), "support cache alloc")) {
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 5;
}
const char *host_base = src_base;
size_t write_off = 0;
for (int i = 0; i < n_ranges; i++) {
if (ranges[i].bytes == 0) continue;
char *dev_ptr = (char *)base + write_off;
cudaError_t e = cudaMemcpy(dev_ptr,
host_base + ranges[i].source_offset,
(size_t)ranges[i].bytes,
cudaMemcpyHostToDevice);
if (e != cudaSuccess) {
cuda_ok(e, "support cache range h2d");
(void)cudaFree(base);
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 7;
}
cache_range_entry ent;
ent.source_offset = ranges[i].source_offset + key_bias;
ent.bytes = ranges[i].bytes;
/* Entries can claim a different (executor) device than the one the
* slab physically lives on: strict lookups filter by entry device,
* and peer access lets the executor's kernels dereference the
* spilled pointer directly. */
ent.device_id = entry_device_id;
ent.device_ptr = dev_ptr;
g_cache_ranges.push_back(ent);
write_off += ranges[i].bytes;
}
std::sort(g_cache_ranges.begin(), g_cache_ranges.end(),
[](const cache_range_entry &a, const cache_range_entry &b) {
if (a.source_offset != b.source_offset)
return a.source_offset < b.source_offset;
return a.device_id < b.device_id;
});
if (getenv("DS4_DSPARK_VERIFY_CACHE") != NULL) {
/* Read back every installed range and compare with the host copy. */
int bad = 0;
write_off = 0;
for (int i = 0; i < n_ranges; i++) {
if (ranges[i].bytes == 0) continue;
char *dev_ptr = (char *)base + write_off;
std::vector<char> tmp((size_t)ranges[i].bytes);
if (cudaMemcpy(tmp.data(), dev_ptr, (size_t)ranges[i].bytes,
cudaMemcpyDeviceToHost) != cudaSuccess ||
memcmp(tmp.data(), host_base + ranges[i].source_offset,
(size_t)ranges[i].bytes) != 0) {
fprintf(stderr,
"ds4: support cache VERIFY MISMATCH offset=%llu bytes=%llu dev=%d\n",
(unsigned long long)ranges[i].source_offset,
(unsigned long long)ranges[i].bytes, device_id);
bad++;
}
write_off += ranges[i].bytes;
}
fprintf(stderr, "ds4: support cache verify dev=%d ranges=%d bad=%d\n",
device_id, n_ranges, bad);
}
if (prev_device >= 0) (void)cudaSetDevice(prev_device);
return 0;
}
extern "C" int ds4_gpu_lookup_cache(uint64_t source_offset, uint64_t bytes,
int *out_device_id, void **out_device_ptr) {
int active_device = -1;
(void)cudaGetDevice(&active_device);
if (!g_cache_ranges.empty()) {
/* upper_bound: first entry with source_offset > query.
* Candidates are at strictly earlier positions; scan all of
* them rather than breaking on the first non-covering entry,
* because the table allows overlap across devices. */
auto it = std::upper_bound(
g_cache_ranges.begin(), g_cache_ranges.end(),
source_offset,
[](uint64_t off, const cache_range_entry &e) {
return off < e.source_offset;
});
const cache_range_entry *match_any = NULL;
const cache_range_entry *match_pref = NULL;
while (it != g_cache_ranges.begin()) {
--it;
/* Overflow-safe coverage check:
* 1. source_offset >= it->source_offset
* 2. bytes <= it->bytes - (source_offset - it->source_offset)
* The second form computes only the remaining capacity inside
* the entry, so neither side can overflow even with bytes ==
* UINT64_MAX. */
if (source_offset >= it->source_offset) {
uint64_t into = source_offset - it->source_offset;
if (into <= it->bytes && bytes <= it->bytes - into) {
if (it->device_id == active_device) {
match_pref = &*it;
break;
}
if (!match_any) match_any = &*it;
}
}
/* Do NOT break on non-covering: an earlier entry may still
* cover if its bytes extend far enough. */
}
const cache_range_entry *m = match_pref ? match_pref : match_any;
if (m) {
if (out_device_id) *out_device_id = m->device_id;
if (out_device_ptr) {
*out_device_ptr =
(char *)m->device_ptr + (source_offset - m->source_offset);
}
return 1;
}
}
/* Legacy chunk-aware fallback (device 0 only). */
const char *p = cuda_model_range_ptr_from_fd(g_model_host_base,
source_offset, bytes,
"lookup_cache");
if (p) {
if (out_device_id) *out_device_id = 0;
if (out_device_ptr) *out_device_ptr = (void *)p;
return 1;
}
return 0;
}
extern "C" int ds4_gpu_lookup_cache_device(uint64_t source_offset, uint64_t bytes) {
int d = -1;
if (!ds4_gpu_lookup_cache(source_offset, bytes, &d, NULL)) return -1;
return d;
}
/* Strict per-device selective-cache lookup.
*
* Returns 1 only if a covering entry exists whose device_id matches the
* caller-supplied expected_device. Otherwise returns 0 with *out_device_ptr
* untouched. Unlike ds4_gpu_lookup_cache, this variant performs NO host-
* pointer fallback (no FD-cache, no model_range_ptr_from_fd) and NO
* different-device fallback. It is the canonical lookup for multi-tier
* dispatch where consuming a different device's pointer would be a
* correctness bug. expected_device is a PHYSICAL CUDA device id (the
* value stored in g_gpu[logical_tier].device_id, not the logical tier
* index). The caller is expected to have cudaSetDevice'd to
* expected_device before invoking; the returned pointer is valid to
* consume from that device's kernel. Added for
* multi-GPU execution (multi-GPU execution). */
extern "C" int ds4_gpu_lookup_cache_strict(uint64_t source_offset,
uint64_t bytes,
int expected_device,
void **out_device_ptr) {
if (g_cache_ranges.empty()) return 0;
auto it = std::upper_bound(
g_cache_ranges.begin(), g_cache_ranges.end(),
source_offset,
[](uint64_t off, const cache_range_entry &e) {
return off < e.source_offset;
});
while (it != g_cache_ranges.begin()) {
--it;
if (source_offset < it->source_offset) {
/* Should not happen given upper_bound semantics, but defensive. */
continue;
}
uint64_t into = source_offset - it->source_offset;
if (into > it->bytes) continue;
if (bytes > it->bytes - into) continue;
if (it->device_id != expected_device) continue;
if (out_device_ptr) {
*out_device_ptr =
(char *)it->device_ptr + (source_offset - it->source_offset);
}
return 1;
}
return 0;
}
extern "C" int ds4_gpu_set_model_fd(int fd) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
g_model_fd = fd;
g_model_fd_host_base = g_model_host_base;
g_model_file_size = 0;
if (g_model_direct_fd >= 0) {
(void)close(g_model_direct_fd);
g_model_direct_fd = -1;
}
g_model_direct_align = 1;
if (fd >= 0) {
struct stat st;
if (fstat(fd, &st) == 0 && st.st_size > 0) {
g_model_file_size = (uint64_t)st.st_size;
if (st.st_blksize > 1) g_model_direct_align = (uint64_t)st.st_blksize;
}
#if defined(__linux__) && defined(O_DIRECT)
if (getenv("DS4_CUDA_NO_DIRECT_IO") == NULL) {
char proc_path[64];
snprintf(proc_path, sizeof(proc_path), "/proc/self/fd/%d", fd);
int direct_fd = open(proc_path, O_RDONLY | O_DIRECT);
if (direct_fd >= 0) {
g_model_direct_fd = direct_fd;
if (g_model_direct_align < 512) g_model_direct_align = 512;
if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA model direct I/O enabled (align=%llu)\n",
(unsigned long long)g_model_direct_align);
}
} else if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) {
fprintf(stderr, "ds4: CUDA model direct I/O unavailable: %s\n", strerror(errno));
}
}
#endif
}
return 1;
}
static int cuda_build_derived_artifacts(
const void *model_map,
uint64_t model_size, uint64_t file_size,
const char *model_path, uint32_t rank, uint32_t world) {
if ((world != 1u && world != 2u) || rank >= world) return 0;
if (!model_map || model_size == 0 || !model_path || !model_path[0]) return 0;
if (!g_derived_ranges.empty()) return (int)g_derived_ranges.size();
if (getenv("DS4_CUDA_NO_DERIVED_WEIGHTS") != NULL) return 0;
const char *build = getenv("DS4_CUDA_BUILD_ARTIFACTS");
if (build && strcmp(build, "0") == 0) return 0;
int device = 0;
int integrated = 0;
if (cudaGetDevice(&device) != cudaSuccess ||
cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated, device) != cudaSuccess ||
!integrated) {
return 0;
}
ds4_repack_file mapped;
if (!ds4_repack_map_file("ds4", model_path, mapped)) return 0;
if (mapped.size != file_size || model_size > file_size) {
ds4_repack_unmap_file(mapped);
fprintf(stderr, "ds4: aligned artifact build skipped: model size changed\n");
return 0;
}
std::vector<ds4_repack_tensor> records;
const bool catalog_ok =
ds4_repack_collect_catalog("ds4", mapped, nullptr, &records);
ds4_repack_unmap_file(mapped);
if (!catalog_ok) return 0;
if (world == 2u) {
/* Repack only the owned half. Dense projections have different slice
* layouts and keep their existing raw caches in network TP. */
std::vector<ds4_repack_tensor> owned;
for (ds4_repack_tensor t : records) {
if (!ds4_repack_iq2_candidate(t) && !ds4_repack_q2k_candidate(t)) continue;
if (t.dims[2] % world || t.bytes % world || t.elements % world) return 0;
t.dims[2] /= world;
t.bytes /= world;
t.elements /= world;
t.off += rank * t.bytes;
if (t.off > model_size || t.bytes > model_size - t.off) return 0;
owned.push_back(t);
}
records.swap(owned);
}
const bool build_moe =
cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled();
const bool build_q8 = world == 1u && cuda_aligned_q8_enabled();
if (!build_moe && !build_q8) return 0;
ds4_repack_build_args args;
args.log_prefix = "ds4";
args.model_id = "base";
args.path = model_path;
args.records = &records;
args.device = device;
args.copy_chunk_bytes = 256ull * 1048576ull;
const double t0 = cuda_wall_sec();
std::vector<ds4_repack_artifact> artifacts;
uint64_t built_bytes = 0;
uint64_t part_bytes = 0;
bool ok = true;
if (build_moe) {
ok = ds4_repack_build_q2k_aligned(args, artifacts, &part_bytes);
built_bytes += part_bytes;
}
if (ok && build_moe) {
ok = ds4_repack_build_iq2_aligned(args, artifacts, &part_bytes);
built_bytes += part_bytes;
}
if (ok && build_q8) {
ok = ds4_repack_build_q8_aligned(args, artifacts, &part_bytes);
built_bytes += part_bytes;
}
if (!ok || artifacts.empty()) {
for (ds4_repack_artifact &artifact : artifacts) {
if (artifact.dev) (void)cudaFree(artifact.dev);
}
fprintf(stderr, "ds4: aligned artifact build failed; using raw weights\n");
return 0;
}
for (const ds4_repack_artifact &artifact : artifacts) {
g_derived_ranges.push_back({
model_map,
artifact.t->off,
artifact.t->bytes,
artifact.kind,
artifact.in_dim,
artifact.out_dim,
artifact.group_count,
artifact.bytes,
(char *)artifact.dev,
});
}
int replaces_complete = build_moe ? 1 : 0;
uint64_t replace_candidates = 0;
for (const ds4_repack_tensor &tensor : records) {
if (!ds4_repack_iq2_candidate(tensor) &&
!ds4_repack_q2k_candidate(tensor)) {
continue;
}
replace_candidates++;
bool found = false;
for (const ds4_repack_artifact &artifact : artifacts) {
if (artifact.t == &tensor) {
found = true;
break;
}
}
if (!found) replaces_complete = 0;
}
if (replace_candidates == 0) replaces_complete = 0;
g_derived_replace_map = model_map;
g_derived_replaces_complete = replaces_complete;
g_derived_artifact_bytes = built_bytes;
g_derived_artifact_build_secs = cuda_wall_sec() - t0;
if (!g_aligned_q81_scratch) {
const size_t scratch_bytes = 96u * 1024u * 1024u;
cudaError_t scratch_err = cudaMalloc(&g_aligned_q81_scratch,
scratch_bytes);
if (scratch_err == cudaSuccess) {
ds4_mmq_set_aligned_q81_scratch(g_aligned_q81_scratch,
scratch_bytes);
} else {
g_aligned_q81_scratch = NULL;
(void)cudaGetLastError();
fprintf(stderr,
"ds4: aligned decode scratch allocation failed: %s\n",
cudaGetErrorString(scratch_err));
}
}
fprintf(stderr,
"ds4: built %llu aligned CUDA artifacts (%.2f GiB) in %.1fs%s\n",
(unsigned long long)g_derived_ranges.size(),
(double)g_derived_artifact_bytes / 1073741824.0,
g_derived_artifact_build_secs,
replaces_complete ? "; expert raw residency replaced" : "");
return (int)g_derived_ranges.size();
}
extern "C" int ds4_gpu_build_derived_artifacts(
const void *model_map, uint64_t model_size, const char *model_path) {
return cuda_build_derived_artifacts(model_map, model_size, model_size, model_path, 0, 1);
}
extern "C" int ds4_gpu_build_derived_artifacts_shard(
const void *model_map, uint64_t model_size, uint64_t file_size,
const char *model_path, uint32_t rank) {
return cuda_build_derived_artifacts(model_map, model_size, file_size, model_path, rank, 2);
}
extern "C" int ds4_gpu_model_range_replaced(
const void *model_map,
uint64_t offset,
uint64_t bytes) {
if (!cuda_model_map_replaces_complete(model_map) || bytes == 0) return 0;
for (const cuda_derived_range &range : g_derived_ranges) {
if (range.host_base == model_map &&
range.source_offset == offset &&
range.source_bytes == bytes &&
(range.kind == CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE ||
range.kind == CUDA_DERIVED_Q2_K_ALIGNED_MOE)) {
return 1;
}
}
return 0;
}
extern "C" int ds4_gpu_cache_model_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, const char *label) {
if (!model_map || bytes == 0) return 1;
if (offset > model_size || bytes > model_size - offset) return 0;
if (cuda_span_fully_replaced(model_map, offset, bytes)) return 1;
if (!cuda_model_range_ptr(model_map, offset, bytes, label ? label : "model_tensor")) return 0;
return cuda_model_range_is_cached(model_map, offset, bytes);
}
extern "C" int ds4_gpu_cache_q8_f16_range(const void *model_map, uint64_t model_size, uint64_t offset, uint64_t bytes, uint64_t in_dim, uint64_t out_dim, const char *label) {
if (!model_map || bytes == 0) return 1;
if (offset > model_size || bytes > model_size - offset) return 0;
static int optional_q8_preload_disabled = 0;
if (optional_q8_preload_disabled) return 1;
const char *cache_label = label ? label : "q8_0";
/* Preload runs before any multi-tier dispatch. The cache entries it creates
* are device-0 by construction; multi-tier callers in kernel wrappers will
* miss the linear scan (device_id filter) and allocate fresh per-device
* copies the first time they're consulted. */
if (getenv("DS4_CUDA_Q8_F32_PRELOAD") != NULL &&
cuda_q8_f32_cache_allowed(cache_label, in_dim, out_dim)) {
if (cuda_q8_f32_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1;
optional_q8_preload_disabled = 1;
return 1;
}
if (!cuda_q8_f16_preload_allowed(cache_label, in_dim, out_dim)) return 1;
if (cuda_q8_f16_ptr(model_map, offset, bytes, in_dim, out_dim, 0, cache_label)) return 1;
optional_q8_preload_disabled = 1;
return 1;
}
extern "C" void ds4_gpu_print_memory_report(const char *label) {
size_t free_b = 0, total_b = 0;
(void)cudaMemGetInfo(&free_b, &total_b);
fprintf(stderr, "ds4: CUDA memory report %s: free %.2f MiB total %.2f MiB\n",
label ? label : "", (double)free_b / 1048576.0, (double)total_b / 1048576.0);
}
extern "C" void ds4_gpu_set_quality(bool quality) {
g_quality_mode = quality ? 1 : 0;
const cublasMath_t math_mode =
(g_quality_mode || getenv("DS4_CUDA_NO_TF32") != NULL)
? CUBLAS_DEFAULT_MATH
: CUBLAS_TF32_TENSOR_OP_MATH;
/* Walk every initialized per-tier handle. Single-tier (g_n_gpus == 1)
* walks exactly one entry. On any device-switch failure,
* skip the tier and continue — the function is void and the math-mode
* setting is advisory, but log so misconfiguration is visible. */
for (int i = 0; i < g_n_gpus; i++) {
if (!g_gpu[i].cublas_ready || !g_gpu[i].cublas) continue;
int prev = -1;
cudaError_t derr = cudaGetDevice(&prev);
if (derr != cudaSuccess) {
fprintf(stderr,
"ds4: ds4_gpu_set_quality: cudaGetDevice failed before tier %d "
"(dev=%d): %s; skipping\n",
i, g_gpu[i].device_id, cudaGetErrorString(derr));
(void)cudaGetLastError();
continue;
}
derr = cudaSetDevice(g_gpu[i].device_id);
if (derr != cudaSuccess) {
fprintf(stderr,
"ds4: ds4_gpu_set_quality: cudaSetDevice(%d) failed for tier %d: "
"%s; skipping\n",
g_gpu[i].device_id, i, cudaGetErrorString(derr));
(void)cudaGetLastError();
if (prev >= 0) (void)cudaSetDevice(prev);
continue;
}
cublasStatus_t st = cublasSetMathMode((cublasHandle_t)g_gpu[i].cublas, math_mode);
if (st != CUBLAS_STATUS_SUCCESS) {
fprintf(stderr,
"ds4: ds4_gpu_set_quality: cublasSetMathMode failed on tier %d "
"(dev=%d): status %d\n",
i, g_gpu[i].device_id, (int)st);
}
if (prev >= 0) (void)cudaSetDevice(prev);
}
}
__global__ static void embed_token_hc_kernel(float *out, const unsigned short *w, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t n = n_embd * n_hc;
if (i >= n) return;
uint32_t e = i % n_embd;
out[i] = __half2float(reinterpret_cast<const __half *>(w)[(uint64_t)token * n_embd + e]);
}
__global__ static void embed_tokens_hc_kernel(
float *out,
const int32_t *tokens,
const __half *w,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
if (gid >= n) return;
uint32_t d = gid % n_embd;
uint64_t tmp = gid / n_embd;
uint32_t t = tmp / n_hc;
int32_t tok_i = tokens[t];
uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i;
if (tok >= n_vocab) tok = 0;
out[gid] = __half2float(w[(uint64_t)tok * n_embd + d]);
}
__global__ static void matmul_f16_kernel(
float *out,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok) {
uint64_t row = (uint64_t)blockIdx.x;
uint64_t tok = (uint64_t)blockIdx.y;
if (row >= out_dim || tok >= n_tok) return;
float sum = 0.0f;
const __half *wr = w + row * in_dim;
const float *xr = x + tok * in_dim;
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
sum += __half2float(wr[i]) * xr[i];
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
}
__global__ static void matmul_f16_serial_kernel(
float *out,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok) {
uint64_t row = (uint64_t)blockIdx.x;
uint64_t tok = (uint64_t)blockIdx.y;
if (row >= out_dim || tok >= n_tok || threadIdx.x != 0) return;
float sum = 0.0f;
const __half *wr = w + row * in_dim;
const float *xr = x + tok * in_dim;
for (uint64_t i = 0; i < in_dim; i++) {
sum += __half2float(wr[i]) * xr[i];
}
out[tok * out_dim + row] = sum;
}
__global__ static void matmul_f16_ordered_chunks_kernel(
float *out,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok) {
uint64_t row = (uint64_t)blockIdx.x;
uint64_t tok = (uint64_t)blockIdx.y;
if (row >= out_dim || tok >= n_tok) return;
__shared__ float partial[32];
const uint32_t tid = threadIdx.x;
float sum = 0.0f;
const uint64_t chunk = (in_dim + 31u) / 32u;
const uint64_t k0 = (uint64_t)tid * chunk;
uint64_t k1 = k0 + chunk;
if (k1 > in_dim) k1 = in_dim;
const __half *wr = w + row * in_dim;
const float *xr = x + tok * in_dim;
for (uint64_t i = k0; i < k1; i++) {
sum += __half2float(wr[i]) * xr[i];
}
partial[tid] = sum;
__syncthreads();
if (tid == 0) {
float total = 0.0f;
for (uint32_t i = 0; i < 32u; i++) total += partial[i];
out[tok * out_dim + row] = total;
}
}
__global__ static void matmul_f16_small_out_hx_ordered_chunks_kernel(
float *out,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim) {
uint64_t row = (uint64_t)blockIdx.x;
if (row >= out_dim) return;
__shared__ float partial[32];
const uint32_t tid = threadIdx.x;
float sum = 0.0f;
const uint64_t chunk = (in_dim + 31u) / 32u;
const uint64_t k0 = (uint64_t)tid * chunk;
uint64_t k1 = k0 + chunk;
if (k1 > in_dim) k1 = in_dim;
const __half *wr = w + row * in_dim;
for (uint64_t i = k0; i < k1; i++) {
const float xv = __half2float(__float2half(x[i]));
sum += __half2float(wr[i]) * xv;
}
partial[tid] = sum;
__syncthreads();
if (tid == 0) {
float total = 0.0f;
for (uint32_t i = 0; i < 32u; i++) total += partial[i];
out[row] = total;
}
}
__global__ static void matmul_f16_small_out_batch_kernel(
float *out,
const __half *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok) {
const uint64_t tok = (uint64_t)blockIdx.x;
const uint32_t tid = threadIdx.x;
if (tok >= n_tok || out_dim > 32u || blockDim.x != 256u) return;
float acc[32];
#pragma unroll
for (uint32_t r = 0; r < 32u; r++) acc[r] = 0.0f;
const float *xr = x + tok * in_dim;
for (uint64_t i = tid; i < in_dim; i += 256u) {
const float xv = xr[i];
#pragma unroll
for (uint32_t r = 0; r < 32u; r++) {
if (r < out_dim) {
acc[r] += __half2float(w[(uint64_t)r * in_dim + i]) * xv;
}
}
}
__shared__ float partial[32 * 256];
#pragma unroll
for (uint32_t r = 0; r < 32u; r++) {
if (r < out_dim) partial[r * 256u + tid] = acc[r];
}
__syncthreads();
for (uint32_t stride = 128u; stride > 0u; stride >>= 1u) {
if (tid < stride) {
#pragma unroll
for (uint32_t r = 0; r < 32u; r++) {
if (r < out_dim) {
partial[r * 256u + tid] += partial[r * 256u + tid + stride];
}
}
}
__syncthreads();
}
if (tid == 0) {
#pragma unroll
for (uint32_t r = 0; r < 32u; r++) {
if (r < out_dim) out[tok * out_dim + r] = partial[r * 256u];
}
}
}
__global__ static void matmul_f16_pair_ordered_chunks_kernel(
float *out0,
float *out1,
const __half *w0,
const __half *w1,
const float *x,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim) {
uint64_t row = (uint64_t)blockIdx.x;
if (row >= out0_dim && row >= out1_dim) return;
__shared__ float partial0[32];
__shared__ float partial1[32];
const uint32_t tid = threadIdx.x;
float sum0 = 0.0f;
float sum1 = 0.0f;
const uint64_t chunk = (in_dim + 31u) / 32u;
const uint64_t k0 = (uint64_t)tid * chunk;
uint64_t k1 = k0 + chunk;
if (k1 > in_dim) k1 = in_dim;
const __half *wr0 = row < out0_dim ? w0 + row * in_dim : w0;
const __half *wr1 = row < out1_dim ? w1 + row * in_dim : w1;
for (uint64_t i = k0; i < k1; i++) {
const float xv = x[i];
if (row < out0_dim) sum0 += __half2float(wr0[i]) * xv;
if (row < out1_dim) sum1 += __half2float(wr1[i]) * xv;
}
partial0[tid] = sum0;
partial1[tid] = sum1;
__syncthreads();
if (tid == 0) {
float total0 = 0.0f;
float total1 = 0.0f;
for (uint32_t i = 0; i < 32u; i++) {
total0 += partial0[i];
total1 += partial1[i];
}
if (row < out0_dim) out0[row] = total0;
if (row < out1_dim) out1[row] = total1;
}
}
/* Preserve the 256-thread halving tree: combine the eight warp partials
* before reducing within the first warp. Only two block barriers are needed. */
__device__ __forceinline__ static float block_sum_f32_256(float v, float *partial) {
const uint32_t t = threadIdx.x;
partial[t] = v;
__syncthreads();
if (t < 32u) {
v = ((partial[t] + partial[t + 128u]) +
(partial[t + 64u] + partial[t + 192u])) +
((partial[t + 32u] + partial[t + 160u]) +
(partial[t + 96u] + partial[t + 224u]));
for (uint32_t s = 16u; s; s >>= 1u)
v += __shfl_down_sync(0xffffffffu, v, s);
if (t == 0u) partial[0] = v;
}
__syncthreads();
return partial[0];
}
__global__ static void matmul_f32_kernel(
float *out,
const float *w,
const float *x,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok) {
uint64_t row = (uint64_t)blockIdx.x;
uint64_t tok = (uint64_t)blockIdx.y;
if (row >= out_dim || tok >= n_tok) return;
float sum = 0.0f;
const float *wr = w + row * in_dim;
const float *xr = x + tok * in_dim;
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
sum += wr[i] * xr[i];
}
__shared__ float partial[256];
sum = block_sum_f32_256(sum, partial);
if (threadIdx.x == 0) out[tok * out_dim + row] = sum;
}
__global__ static void repeat_hc_kernel(float *out, const float *row, uint32_t n_embd, uint32_t n_hc) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_embd * n_hc;
if (i >= n) return;
out[i] = row[i % n_embd];
}
__global__ static void repeat_hc_rows_kernel(float *out, const float *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
if (i >= n) return;
uint64_t hc_row = (uint64_t)n_hc * n_embd;
uint64_t tok = i / hc_row;
uint64_t embd = i % n_embd;
out[i] = rows[tok * n_embd + embd];
}
__global__ static void pack_slot_rows_f32_kernel(float *out, const float *slots, uint32_t n_rows, uint32_t width, uint32_t n_slots, uint32_t slot_cap) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_rows * n_slots * width;
if (i >= n) return;
uint64_t col = i % width;
uint64_t slot = (i / width) % n_slots;
uint64_t row = i / ((uint64_t)n_slots * width);
out[i] = slots[((slot * slot_cap) + row) * width + col];
}
__global__ static void f32_to_f16_kernel(__half *out, const float *x, uint64_t n) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = __float2half(x[i]);
}
__device__ static float warp_sum_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v += __shfl_down_sync(0xffffffffu, v, offset);
}
return v;
}
__device__ static float warp_max_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, offset));
}
return v;
}
__device__ static float dot4_f32(float4 a, float4 b) {
return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
}
__device__ __forceinline__ static int32_t load_i8x4_i32_aligned(const int8_t *p) {
return *(const int32_t *)p;
}
__device__ __forceinline__ static int32_t load_i8x4_i32_unaligned(const int8_t *p) {
/* Raw Q8_0 blocks have a 34-byte stride: two-byte, not four-byte,
* alignment. Both loads stay inside the requested four weight bytes. */
const uint16_t *u = (const uint16_t *)p;
return (int32_t)((uint32_t)u[0] | ((uint32_t)u[1] << 16));
}
__device__ __forceinline__ static int32_t dot_i8x32_dp4a(const int8_t *a, const int8_t *b) {
int32_t dot = 0;
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
dot = __dp4a(load_i8x4_i32_unaligned(a + i), load_i8x4_i32_aligned(b + i), dot);
}
return dot;
}
__device__ __forceinline__ static int32_t dot_i8_block(const int8_t *a, const int8_t *b, uint64_t n, int use_dp4a) {
if (use_dp4a && n == 32u) return dot_i8x32_dp4a(a, b);
int32_t dot = 0;
for (uint64_t i = 0; i < n; i++) dot += (int32_t)a[i] * (int32_t)b[i];
return dot;
}
__global__ static void quantize_q8_0_f32_kernel(
int8_t *xq,
float *xscale,
const float *x,
uint64_t in_dim,
uint64_t blocks) {
uint64_t b = blockIdx.x;
uint64_t tok = blockIdx.y;
if (b >= blocks) return;
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const float *xr = x + tok * in_dim + i0;
float a = 0.0f;
if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]);
__shared__ float vals[32];
vals[threadIdx.x] = a;
__syncthreads();
for (uint32_t stride = 16; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) vals[threadIdx.x] = fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]);
__syncthreads();
}
const float d = vals[0] / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
if (threadIdx.x == 0) xscale[tok * blocks + b] = d;
int8_t *dst = xq + (tok * blocks + b) * 32;
if (threadIdx.x < bn) {
int v = (int)lrintf(xr[threadIdx.x] * id);
v = v > 127 ? 127 : (v < -128 ? -128 : v);
dst[threadIdx.x] = (int8_t)v;
} else {
dst[threadIdx.x] = 0;
}
}
__global__ static void quantize_q8_0_group_slice_rows_kernel(
int8_t *xq,
float *xscale,
const float *x,
uint64_t group_dim,
uint64_t blocks,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt) {
const uint64_t b = blockIdx.x;
const uint64_t packed_row = blockIdx.y;
if (b >= blocks) return;
const uint64_t token = packed_row / group_cnt;
const uint64_t group = group0 + packed_row - token * group_cnt;
const uint64_t i0 = b * 32u;
const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u;
const float *xr = x +
(token * n_groups_total + group) * group_dim + i0;
float a = 0.0f;
if (threadIdx.x < bn) a = fabsf(xr[threadIdx.x]);
__shared__ float vals[32];
vals[threadIdx.x] = a;
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
vals[threadIdx.x] =
fmaxf(vals[threadIdx.x], vals[threadIdx.x + stride]);
}
__syncthreads();
}
const float d = vals[0] / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
if (threadIdx.x == 0u) xscale[packed_row * blocks + b] = d;
int8_t *dst = xq + (packed_row * blocks + b) * 32u;
if (threadIdx.x < bn) {
int v = (int)lrintf(xr[threadIdx.x] * id);
v = v > 127 ? 127 : (v < -128 ? -128 : v);
dst[threadIdx.x] = (int8_t)v;
} else {
dst[threadIdx.x] = 0;
}
}
__global__ static void matmul_q8_0_preq_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
uint64_t row = (uint64_t)blockIdx.x;
uint64_t tok = (uint64_t)blockIdx.y;
if (row >= out_dim || tok >= n_tok) return;
const unsigned char *wr = w + row * blocks * 34;
const int8_t *xqr = xq + tok * blocks * 32;
const float *xsr = xscale + tok * blocks;
float acc = 0.0f;
for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) {
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xqr + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
__shared__ float partial[256];
partial[threadIdx.x] = acc;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
}
__global__ static void matmul_q8_0_preq_warp8_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks,
int use_dp4a) {
uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const unsigned char *wr = w + row * blocks * 34;
const int8_t *xqr = xq + tok * blocks * 32u;
const float *xsr = xscale + tok * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xqr + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) out[tok * out_dim + row] = acc;
}
__device__ __forceinline__ static uint32_t q8_top1_float_ordered_key(float v) {
const uint32_t u = __float_as_uint(v);
return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u);
}
__device__ __forceinline__ static uint64_t q8_top1_pack_key(float v, uint32_t idx) {
return ((uint64_t)q8_top1_float_ordered_key(v) << 32u) |
(uint64_t)(0xffffffffu - idx);
}
__device__ __forceinline__ static float q8_top1_unpack_value(uint32_t ordered) {
const uint32_t u = (ordered & 0x80000000u)
? (ordered ^ 0x80000000u)
: ~ordered;
return __uint_as_float(u);
}
__global__ static void matmul_q8_0_top1_preq_warp8_kernel(
unsigned long long *best_key,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks,
uint32_t index_offset,
int use_dp4a) {
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
const uint64_t row = (uint64_t)blockIdx.x * 8u + warp;
const bool valid = row < out_dim;
float acc = 0.0f;
if (valid) {
const unsigned char *wr = w + row * blocks * 34;
for (uint64_t b = lane; b < blocks; b += 32u) {
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xq + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xscale[b] * (float)dot;
}
}
acc = warp_sum_f32(acc);
__shared__ unsigned long long keys[8];
if (lane == 0u) {
keys[warp] = valid
? (unsigned long long)q8_top1_pack_key(acc, index_offset + (uint32_t)row)
: 0ull;
}
__syncthreads();
if (threadIdx.x == 0u) {
unsigned long long block_best = keys[0];
#pragma unroll
for (uint32_t i = 1u; i < 8u; i++) {
if (keys[i] > block_best) block_best = keys[i];
}
(void)atomicMax(best_key, block_best);
}
}
__global__ static void matmul_q8_0_top1_unpack_kernel(
uint32_t *selected,
float *values,
const unsigned long long *best_key) {
if (threadIdx.x != 0u || blockIdx.x != 0u) return;
const uint64_t key = (uint64_t)best_key[0];
const uint32_t ordered = (uint32_t)(key >> 32u);
const uint32_t idx = 0xffffffffu - (uint32_t)key;
selected[0] = idx;
values[0] = q8_top1_unpack_value(ordered);
}
__global__ static void matmul_q8_0_kslice_preq_warp8_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t slice_dim,
uint64_t out_dim,
uint64_t full_blocks,
uint64_t block_start,
uint64_t slice_blocks,
int use_dp4a) {
const uint64_t row =
(uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = blockIdx.y;
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
out += tok * out_dim;
xq += tok * slice_blocks * 32u;
xscale += tok * slice_blocks;
const unsigned char *wr =
w + row * full_blocks * 34u + block_start * 34u;
float acc = 0.0f;
for (uint64_t b = lane; b < slice_blocks; b += 32u) {
uint64_t i0 = b * 32u;
uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u;
const __half *scale_h = (const __half *)(wr + b * 34u);
const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
const int8_t *xqb = xq + b * 32u;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xscale[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) out[row] = acc;
}
__global__ static void matmul_q8_0_pair_preq_warp8_kernel(
float *out0,
float *out1,
const unsigned char *w0,
const unsigned char *w1,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
uint64_t blocks,
int use_dp4a) {
uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
uint32_t lane = threadIdx.x & 31u;
if (row >= out0_dim && row >= out1_dim) return;
float acc0 = 0.0f;
float acc1 = 0.0f;
const unsigned char *wr0 = row < out0_dim ? w0 + row * blocks * 34 : NULL;
const unsigned char *wr1 = row < out1_dim ? w1 + row * blocks * 34 : NULL;
const int8_t *xqr = xq + tok * blocks * 32u;
const float *xsr = xscale + tok * blocks;
for (uint64_t b = lane; b < blocks; b += 32u) {
uint64_t i0 = b * 32;
uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const int8_t *xqb = xqr + b * 32;
const float xs = xsr[b];
if (wr0) {
const __half *scale_h = (const __half *)(wr0 + b * 34);
const int8_t *qs = (const int8_t *)(wr0 + b * 34 + 2);
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc0 += __half2float(*scale_h) * xs * (float)dot;
}
if (wr1) {
const __half *scale_h = (const __half *)(wr1 + b * 34);
const int8_t *qs = (const int8_t *)(wr1 + b * 34 + 2);
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc1 += __half2float(*scale_h) * xs * (float)dot;
}
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
if (lane == 0) {
if (row < out0_dim) out0[tok * out0_dim + row] = acc0;
if (row < out1_dim) out1[tok * out1_dim + row] = acc1;
}
}
__global__ static void shared_mid_q8_0_preq_warp8_exact_kernel(
float *mid,
const unsigned char *gate_w,
const unsigned char *up_w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks,
float clamp,
const int32_t *selected,
uint32_t expert_split,
bool home_rank,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
if (selected) {
/* Complementary predicates select exactly one writer; ties stay on
* the home rank to avoid an unnecessary peer store. */
uint32_t home_count = 0u;
uint32_t peer_count = 0u;
#pragma unroll
for (uint32_t i = 0; i < 6u; i++) {
const int32_t expert = selected[i];
if (expert >= 0 && (uint32_t)expert < expert_split) {
home_count++;
} else if (expert >= 0 &&
(uint32_t)expert < 2u * expert_split) {
peer_count++;
}
}
const bool assigned = home_rank
? home_count <= peer_count : peer_count < home_count;
if (!assigned) return;
}
const unsigned char *gate_row = gate_w + row * blocks * 34u;
const unsigned char *up_row = up_w + row * blocks * 34u;
float gate = 0.0f;
float up = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32u;
const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
const int8_t *xqb = xq + b * 32u;
const float xs = xscale[b];
const unsigned char *gb = gate_row + b * 34u;
const unsigned char *ub = up_row + b * 34u;
gate += __half2float(*(const __half *)gb) * xs *
(float)dot_i8_block((const int8_t *)(gb + 2u), xqb, bn,
use_dp4a);
up += __half2float(*(const __half *)ub) * xs *
(float)dot_i8_block((const int8_t *)(ub + 2u), xqb, bn,
use_dp4a);
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
gate = fminf(gate, clamp);
up = fminf(fmaxf(up, -clamp), clamp);
}
const float silu = gate / (1.0f + expf(-gate));
mid[row] = silu * up * 1.0f;
}
}
__global__ static void matmul_q8_0_pair_preq_batch_kernel(
float *out0,
float *out1,
const unsigned char *w0,
const unsigned char *w1,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x;
const uint64_t tok = (uint64_t)blockIdx.y;
if (tok >= n_tok) return;
const int has0 = row < out0_dim;
const int has1 = row < out1_dim;
if (!has0 && !has1) return;
const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL;
const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL;
const int8_t *xqr = xq + tok * blocks * 32u;
const float *xsr = xscale + tok * blocks;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) {
const uint64_t i0 = b * 32u;
const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
const int8_t *xqb = xqr + b * 32u;
const float xs = xsr[b];
if (has0) {
const __half *scale_h = (const __half *)(wr0 + b * 34u);
const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u);
const int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc0 += __half2float(*scale_h) * xs * (float)dot;
}
if (has1) {
const __half *scale_h = (const __half *)(wr1 + b * 34u);
const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u);
const int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc1 += __half2float(*scale_h) * xs * (float)dot;
}
}
__shared__ float partial0[256];
__shared__ float partial1[256];
partial0[threadIdx.x] = acc0;
partial1[threadIdx.x] = acc1;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
partial0[threadIdx.x] += partial0[threadIdx.x + stride];
partial1[threadIdx.x] += partial1[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
if (has0) out0[tok * out0_dim + row] = partial0[0];
if (has1) out1[tok * out1_dim + row] = partial1[0];
}
}
__global__ static void matmul_q8_0_pair_preq_batch_tok2_exact_kernel(
float *out0,
float *out1,
const unsigned char *w0,
const unsigned char *w1,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x;
const uint64_t tok0 = (uint64_t)blockIdx.y * 2u;
if (tok0 >= n_tok) return;
const int has0 = row < out0_dim;
const int has1 = row < out1_dim;
if (!has0 && !has1) return;
const int valid1 = tok0 + 1u < n_tok;
const unsigned char *wr0 = has0 ? w0 + row * blocks * 34u : NULL;
const unsigned char *wr1 = has1 ? w1 + row * blocks * 34u : NULL;
const int8_t *xqr0 = xq + tok0 * blocks * 32u;
const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0;
const float *xsr0 = xscale + tok0 * blocks;
const float *xsr1 = valid1 ? xsr0 + blocks : xsr0;
float acc00 = 0.0f;
float acc01 = 0.0f;
float acc10 = 0.0f;
float acc11 = 0.0f;
for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) {
const uint64_t i0 = b * 32u;
const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
const int8_t *xqb0 = xqr0 + b * 32u;
const int8_t *xqb1 = xqr1 + b * 32u;
const float xs0 = xsr0[b];
const float xs1 = valid1 ? xsr1[b] : 0.0f;
if (has0) {
const __half *scale_h = (const __half *)(wr0 + b * 34u);
const int8_t *qs = (const int8_t *)(wr0 + b * 34u + 2u);
const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
int dot1 = 0;
if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
const float ws = __half2float(*scale_h);
acc00 += ws * xs0 * (float)dot0;
if (valid1) acc01 += ws * xs1 * (float)dot1;
}
if (has1) {
const __half *scale_h = (const __half *)(wr1 + b * 34u);
const int8_t *qs = (const int8_t *)(wr1 + b * 34u + 2u);
const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
int dot1 = 0;
if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
const float ws = __half2float(*scale_h);
acc10 += ws * xs0 * (float)dot0;
if (valid1) acc11 += ws * xs1 * (float)dot1;
}
}
__shared__ float partial00[256];
__shared__ float partial01[256];
__shared__ float partial10[256];
__shared__ float partial11[256];
partial00[threadIdx.x] = acc00;
partial01[threadIdx.x] = acc01;
partial10[threadIdx.x] = acc10;
partial11[threadIdx.x] = acc11;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
partial00[threadIdx.x] += partial00[threadIdx.x + stride];
partial01[threadIdx.x] += partial01[threadIdx.x + stride];
partial10[threadIdx.x] += partial10[threadIdx.x + stride];
partial11[threadIdx.x] += partial11[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
if (has0) {
out0[tok0 * out0_dim + row] = partial00[0];
if (valid1) out0[(tok0 + 1u) * out0_dim + row] = partial01[0];
}
if (has1) {
out1[tok0 * out1_dim + row] = partial10[0];
if (valid1) out1[(tok0 + 1u) * out1_dim + row] = partial11[0];
}
}
}
__device__ static float moe_owned_packed_combine_row(
const float *home_slots,
const float *peer_packed,
const int32_t *selected,
uint32_t row,
uint32_t out_dim,
uint32_t expert_split);
__global__ static void matmul_q8_0_hc_expand_preq_warp8_kernel(
float *out_hc,
float *block_out,
const float *block_add,
const float *block_add2,
const float *owned_home_slots,
const float *owned_peer_packed,
const int32_t *owned_selected,
const float *residual_hc,
const float *split,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint32_t n_embd,
uint32_t n_hc,
uint64_t blocks,
int has_add,
int has_add2,
int has_owned_slots,
uint32_t owned_expert_split,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const unsigned char *wr = w + row * blocks * 34;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xq + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xscale[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) {
const uint32_t d = (uint32_t)row;
block_out[d] = acc;
float block_v = acc;
if (has_owned_slots) {
const float routed = moe_owned_packed_combine_row(
owned_home_slots,
owned_peer_packed,
owned_selected,
d,
(uint32_t)out_dim,
owned_expert_split);
block_v = __fadd_rn(block_v, routed);
} else if (has_add) {
float add_v = block_add[d];
if (has_add2) add_v += block_add2[d];
block_v += add_v;
}
const float *post = split + n_hc;
const float *comb = split + 2u * n_hc;
for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
float hc_acc = block_v * post[dst_hc];
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
hc_acc += comb_v * res_v;
}
out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
}
}
}
__global__ static void matmul_q8_0_kslice_hc_expand_add_preq_warp8_kernel(
float *out_hc,
float *block_out,
const float *block_add,
const float *residual_hc,
const float *split,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t slice_dim,
uint64_t out_dim,
uint64_t full_blocks,
uint64_t block_start,
uint64_t slice_blocks,
uint32_t n_embd,
uint32_t n_hc,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const unsigned char *wr = w + row * full_blocks * 34u + block_start * 34u;
float acc = 0.0f;
for (uint64_t b = lane; b < slice_blocks; b += 32u) {
const uint64_t i0 = b * 32u;
const uint64_t bn = slice_dim - i0 < 32u ? slice_dim - i0 : 32u;
const __half *scale_h = (const __half *)(wr + b * 34u);
const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
const int8_t *xqb = xq + b * 32u;
const int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xscale[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) {
const uint32_t d = (uint32_t)row;
block_out[d] = acc;
const float block_v = acc + block_add[d];
const float *post = split + n_hc;
const float *comb = split + 2u * n_hc;
for (uint32_t dst_hc = 0; dst_hc < n_hc; dst_hc++) {
float hc_acc = block_v * post[dst_hc];
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
const float comb_v = comb[dst_hc + (uint64_t)src_hc * n_hc];
const float res_v = residual_hc[(uint64_t)src_hc * n_embd + d];
hc_acc += comb_v * res_v;
}
out_hc[(uint64_t)dst_hc * n_embd + d] = hc_acc;
}
}
}
__global__ static void matmul_q8_0_preq_batch_warp8_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim || tok >= n_tok) return;
const unsigned char *wr = w + row * blocks * 34;
const int8_t *xqr = xq + tok * blocks * 32;
const float *xsr = xscale + tok * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xqr + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) out[tok * out_dim + row] = acc;
}
__global__ static void matmul_q8_0_preq_batch_warp8_tok2_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim) return;
const unsigned char *wr = w + row * blocks * 34u;
const int8_t *xqr0 = xq;
const int8_t *xqr1 = xq + blocks * 32u;
const float *xsr0 = xscale;
const float *xsr1 = xscale + blocks;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32u;
const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
const __half *scale_h = (const __half *)(wr + b * 34u);
const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
const int8_t *xqb0 = xqr0 + b * 32u;
const int8_t *xqb1 = xqr1 + b * 32u;
int dot0 = 0;
int dot1 = 0;
if (use_dp4a && bn == 32u) {
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
const int32_t w4 = load_i8x4_i32_unaligned(qs + i);
dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0);
dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1);
}
} else {
dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
}
const float ws = __half2float(*scale_h);
acc0 += ws * xsr0[b] * (float)dot0;
acc1 += ws * xsr1[b] * (float)dot1;
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
if (lane == 0) {
out[row] = acc0;
out[out_dim + row] = acc1;
}
}
__global__ static void matmul_q8_0_preq_batch_warp8_tok4_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok0 = (uint64_t)blockIdx.y * 4u;
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim || tok0 >= n_tok) return;
const unsigned char *wr = w + row * blocks * 34;
const int8_t *xqr0 = xq + tok0 * blocks * 32;
const int8_t *xqr1 = xqr0 + blocks * 32;
const int8_t *xqr2 = xqr1 + blocks * 32;
const int8_t *xqr3 = xqr2 + blocks * 32;
const float *xsr0 = xscale + tok0 * blocks;
const float *xsr1 = xsr0 + blocks;
const float *xsr2 = xsr1 + blocks;
const float *xsr3 = xsr2 + blocks;
const int valid1 = tok0 + 1u < n_tok;
const int valid2 = tok0 + 2u < n_tok;
const int valid3 = tok0 + 3u < n_tok;
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb0 = xqr0 + b * 32;
const int8_t *xqb1 = xqr1 + b * 32;
const int8_t *xqb2 = xqr2 + b * 32;
const int8_t *xqb3 = xqr3 + b * 32;
int dot0 = 0;
int dot1 = 0;
int dot2 = 0;
int dot3 = 0;
if (use_dp4a && bn == 32u) {
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
const int32_t w4 = load_i8x4_i32_unaligned(qs + i);
dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0);
if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1);
if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2);
if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3);
}
} else {
dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a);
if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a);
}
const float ws = __half2float(*scale_h);
acc0 += ws * xsr0[b] * (float)dot0;
if (valid1) acc1 += ws * xsr1[b] * (float)dot1;
if (valid2) acc2 += ws * xsr2[b] * (float)dot2;
if (valid3) acc3 += ws * xsr3[b] * (float)dot3;
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
acc2 = warp_sum_f32(acc2);
acc3 = warp_sum_f32(acc3);
if (lane == 0) {
out[tok0 * out_dim + row] = acc0;
if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1;
if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2;
if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3;
}
}
__global__ static void matmul_q8_0_preq_batch_warp8_tok8_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok0 = (uint64_t)blockIdx.y * 8u;
const uint32_t lane = threadIdx.x & 31u;
if (row >= out_dim || tok0 >= n_tok) return;
const unsigned char *wr = w + row * blocks * 34;
const uint64_t xq_stride = blocks * 32u;
const int8_t *xqr0 = xq + tok0 * xq_stride;
const int valid1 = tok0 + 1u < n_tok;
const int valid2 = tok0 + 2u < n_tok;
const int valid3 = tok0 + 3u < n_tok;
const int valid4 = tok0 + 4u < n_tok;
const int valid5 = tok0 + 5u < n_tok;
const int valid6 = tok0 + 6u < n_tok;
const int valid7 = tok0 + 7u < n_tok;
const int8_t *xqr1 = valid1 ? xqr0 + xq_stride : xqr0;
const int8_t *xqr2 = valid2 ? xqr1 + xq_stride : xqr0;
const int8_t *xqr3 = valid3 ? xqr2 + xq_stride : xqr0;
const int8_t *xqr4 = valid4 ? xqr3 + xq_stride : xqr0;
const int8_t *xqr5 = valid5 ? xqr4 + xq_stride : xqr0;
const int8_t *xqr6 = valid6 ? xqr5 + xq_stride : xqr0;
const int8_t *xqr7 = valid7 ? xqr6 + xq_stride : xqr0;
const float *xsr0 = xscale + tok0 * blocks;
const float *xsr1 = valid1 ? xsr0 + blocks : xsr0;
const float *xsr2 = valid2 ? xsr1 + blocks : xsr0;
const float *xsr3 = valid3 ? xsr2 + blocks : xsr0;
const float *xsr4 = valid4 ? xsr3 + blocks : xsr0;
const float *xsr5 = valid5 ? xsr4 + blocks : xsr0;
const float *xsr6 = valid6 ? xsr5 + blocks : xsr0;
const float *xsr7 = valid7 ? xsr6 + blocks : xsr0;
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
float acc4 = 0.0f;
float acc5 = 0.0f;
float acc6 = 0.0f;
float acc7 = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = in_dim - i0 < 32 ? in_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb0 = xqr0 + b * 32;
const int8_t *xqb1 = xqr1 + b * 32;
const int8_t *xqb2 = xqr2 + b * 32;
const int8_t *xqb3 = xqr3 + b * 32;
const int8_t *xqb4 = xqr4 + b * 32;
const int8_t *xqb5 = xqr5 + b * 32;
const int8_t *xqb6 = xqr6 + b * 32;
const int8_t *xqb7 = xqr7 + b * 32;
int dot0 = 0;
int dot1 = 0;
int dot2 = 0;
int dot3 = 0;
int dot4 = 0;
int dot5 = 0;
int dot6 = 0;
int dot7 = 0;
if (use_dp4a && bn == 32u) {
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
const int32_t w4 = load_i8x4_i32_unaligned(qs + i);
dot0 = __dp4a(w4, load_i8x4_i32_aligned(xqb0 + i), dot0);
if (valid1) dot1 = __dp4a(w4, load_i8x4_i32_aligned(xqb1 + i), dot1);
if (valid2) dot2 = __dp4a(w4, load_i8x4_i32_aligned(xqb2 + i), dot2);
if (valid3) dot3 = __dp4a(w4, load_i8x4_i32_aligned(xqb3 + i), dot3);
if (valid4) dot4 = __dp4a(w4, load_i8x4_i32_aligned(xqb4 + i), dot4);
if (valid5) dot5 = __dp4a(w4, load_i8x4_i32_aligned(xqb5 + i), dot5);
if (valid6) dot6 = __dp4a(w4, load_i8x4_i32_aligned(xqb6 + i), dot6);
if (valid7) dot7 = __dp4a(w4, load_i8x4_i32_aligned(xqb7 + i), dot7);
}
} else {
dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
if (valid2) dot2 = dot_i8_block(qs, xqb2, bn, use_dp4a);
if (valid3) dot3 = dot_i8_block(qs, xqb3, bn, use_dp4a);
if (valid4) dot4 = dot_i8_block(qs, xqb4, bn, use_dp4a);
if (valid5) dot5 = dot_i8_block(qs, xqb5, bn, use_dp4a);
if (valid6) dot6 = dot_i8_block(qs, xqb6, bn, use_dp4a);
if (valid7) dot7 = dot_i8_block(qs, xqb7, bn, use_dp4a);
}
const float ws = __half2float(*scale_h);
acc0 += ws * xsr0[b] * (float)dot0;
if (valid1) acc1 += ws * xsr1[b] * (float)dot1;
if (valid2) acc2 += ws * xsr2[b] * (float)dot2;
if (valid3) acc3 += ws * xsr3[b] * (float)dot3;
if (valid4) acc4 += ws * xsr4[b] * (float)dot4;
if (valid5) acc5 += ws * xsr5[b] * (float)dot5;
if (valid6) acc6 += ws * xsr6[b] * (float)dot6;
if (valid7) acc7 += ws * xsr7[b] * (float)dot7;
}
acc0 = warp_sum_f32(acc0);
acc1 = warp_sum_f32(acc1);
acc2 = warp_sum_f32(acc2);
acc3 = warp_sum_f32(acc3);
acc4 = warp_sum_f32(acc4);
acc5 = warp_sum_f32(acc5);
acc6 = warp_sum_f32(acc6);
acc7 = warp_sum_f32(acc7);
if (lane == 0) {
out[tok0 * out_dim + row] = acc0;
if (valid1) out[(tok0 + 1u) * out_dim + row] = acc1;
if (valid2) out[(tok0 + 2u) * out_dim + row] = acc2;
if (valid3) out[(tok0 + 3u) * out_dim + row] = acc3;
if (valid4) out[(tok0 + 4u) * out_dim + row] = acc4;
if (valid5) out[(tok0 + 5u) * out_dim + row] = acc5;
if (valid6) out[(tok0 + 6u) * out_dim + row] = acc6;
if (valid7) out[(tok0 + 7u) * out_dim + row] = acc7;
}
}
__global__ static void matmul_q8_0_preq_batch_tok2_exact_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x;
const uint64_t tok0 = (uint64_t)blockIdx.y * 2u;
if (row >= out_dim || tok0 >= n_tok) return;
const int valid1 = tok0 + 1u < n_tok;
const unsigned char *wr = w + row * blocks * 34u;
const int8_t *xqr0 = xq + tok0 * blocks * 32u;
const int8_t *xqr1 = valid1 ? xqr0 + blocks * 32u : xqr0;
const float *xsr0 = xscale + tok0 * blocks;
const float *xsr1 = valid1 ? xsr0 + blocks : xsr0;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint64_t b = threadIdx.x; b < blocks; b += blockDim.x) {
const uint64_t i0 = b * 32u;
const uint64_t bn = in_dim - i0 < 32u ? in_dim - i0 : 32u;
const __half *scale_h = (const __half *)(wr + b * 34u);
const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
const int8_t *xqb0 = xqr0 + b * 32u;
const int8_t *xqb1 = xqr1 + b * 32u;
const int dot0 = dot_i8_block(qs, xqb0, bn, use_dp4a);
int dot1 = 0;
if (valid1) dot1 = dot_i8_block(qs, xqb1, bn, use_dp4a);
const float ws = __half2float(*scale_h);
acc0 += ws * xsr0[b] * (float)dot0;
if (valid1) acc1 += ws * xsr1[b] * (float)dot1;
}
__shared__ float partial0[256];
__shared__ float partial1[256];
partial0[threadIdx.x] = acc0;
partial1[threadIdx.x] = acc1;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
partial0[threadIdx.x] += partial0[threadIdx.x + stride];
partial1[threadIdx.x] += partial1[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
out[tok0 * out_dim + row] = partial0[0];
if (valid1) out[(tok0 + 1u) * out_dim + row] = partial1[0];
}
}
/* ---- INT8 tensor-core exact Q8_0 batch matmul --------------------------
* Bit-identical replacement for the exact tok2/warp8-family batched Q8_0
* kernels. Each output element's reduction is the reference's strided
* halving tree over T slots (T = reduction width: 32 for the warp kernels,
* cuda_q8_exact_threads(blocks) for the exact kernels; slots >= blocks hold
* +0.0f). The kernel decomposes that tree as: 32 streams at stride T/32
* whose 32 terms per outer step j combine via an adjacent-pairwise static
* register stack taken in bit-reversed stream order (== the top five strided
* tree levels), plus per-(j&3) sequential accumulators and a fixed tail for
* the remaining levels. Fuzz-verified bitwise against both reference
* kernels across shapes, including blocks < T and ragged out_dim/n_tok.
* Rollback: DS4_CUDA_NO_Q8_MMA=1. */
__device__ __forceinline__ static uint32_t ldu32_unaligned(const uint8_t *p) {
const uintptr_t addr = (uintptr_t)p;
if (!(addr & 3u)) return *(const uint32_t *)p;
/* Q8 blocks are two-byte aligned. Read exactly four bytes, including at
* the end of a tensor with no padding after its final quant block. */
const uint16_t *half = (const uint16_t *)p;
return (uint32_t)half[0] | ((uint32_t)half[1] << 16u);
}
__device__ __forceinline__ static void mma_m16n8k32_s8(
int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3,
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
uint32_t b0, uint32_t b1) {
#if __CUDA_ARCH__ >= 800
asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
: "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3)
: "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
#else
(void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3;
#endif
}
__device__ __forceinline__ static uint32_t bitrev5(uint32_t i) {
return ((i & 1u) << 4) | ((i & 2u) << 2) | (i & 4u) | ((i & 8u) >> 2) | ((i & 16u) >> 4);
}
template <uint32_t T>
__global__ static void matmul_q8_0_mma_exact_kernel(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
uint64_t a_stride_blocks, /* activation row stride in blocks (>= blocks) */
uint64_t out_stride) { /* output token stride in floats (>= out_dim) */
extern __shared__ unsigned char q8mma_sh[];
/* Different output rows/tokens read the same block simultaneously.
* Padded pitches keep their scales in distinct shared-memory banks. */
const uint32_t ws_pitch = (uint32_t)blocks + 2u;
const uint32_t xs_pitch = (uint32_t)blocks + 1u;
__half *sh_ws = (__half *)q8mma_sh;
float *sh_xs = (float *)(q8mma_sh + 64u * ws_pitch * sizeof(__half));
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const bool grouped = in_dim >= 4096u && out_dim >= 4096u && n_tok >= 128u;
const uint64_t row_base = (grouped ? (uint64_t)blockIdx.y * 8u + blockIdx.x % 8u : blockIdx.x) * 64u;
const uint64_t tok_base = (grouped ? blockIdx.x / 8u : blockIdx.y) * 16u;
if (row_base >= out_dim) return;
/* stage weight scales (64 rows) and activation scales (16 tokens) */
for (uint32_t idx = threadIdx.x; idx < 64u * (uint32_t)blocks; idx += blockDim.x) {
const uint32_t rl = idx / (uint32_t)blocks;
const uint32_t b = idx - rl * (uint32_t)blocks;
uint64_t row = row_base + rl;
if (row >= out_dim) row = out_dim - 1u;
sh_ws[rl * ws_pitch + b] = *(const __half *)(w + row * blocks * 34u + (uint64_t)b * 34u);
}
for (uint32_t idx = threadIdx.x; idx < 16u * (uint32_t)blocks; idx += blockDim.x) {
const uint32_t tl = idx / (uint32_t)blocks;
const uint32_t b = idx - tl * (uint32_t)blocks;
const uint64_t tok = tok_base + tl;
sh_xs[tl * xs_pitch + b] = tok < n_tok ? xscale[tok * a_stride_blocks + b] : 0.0f;
}
__syncthreads();
const uint64_t row0 = row_base + (uint64_t)warp * 8u;
/* thread's C elements: rows n0,n0+1; tokens mt0, mt0+8 */
const uint32_t n0 = (lane & 3u) * 2u;
const uint32_t mt0 = lane >> 2u;
const uint64_t tokA = tok_base + mt0;
const uint64_t tokB = tok_base + mt0 + 8u;
/* A source rows for loads (fragment layout): rows lane>>2 and (lane>>2)+8 */
const uint64_t a_tok_lo = tok_base + (lane >> 2u);
const uint64_t a_tok_hi = a_tok_lo + 8u;
const int8_t *aq_lo = xq + (a_tok_lo < n_tok ? a_tok_lo : 0u) * a_stride_blocks * 32u;
const int8_t *aq_hi = xq + (a_tok_hi < n_tok ? a_tok_hi : 0u) * a_stride_blocks * 32u;
const bool a_lo_ok = a_tok_lo < n_tok;
const bool a_hi_ok = a_tok_hi < n_tok;
/* B source row for loads: row lane>>2 within the warp tile */
uint64_t b_row = row0 + (lane >> 2u);
if (b_row >= out_dim) b_row = out_dim - 1u;
const unsigned char *b_wr = w + b_row * blocks * 34u;
/* per-element (4) x per-(j&3) accumulators */
float acc00 = 0.0f, acc01 = 0.0f, acc02 = 0.0f, acc03 = 0.0f;
float acc10 = 0.0f, acc11 = 0.0f, acc12 = 0.0f, acc13 = 0.0f;
float acc20 = 0.0f, acc21 = 0.0f, acc22 = 0.0f, acc23 = 0.0f;
float acc30 = 0.0f, acc31 = 0.0f, acc32 = 0.0f, acc33 = 0.0f;
const uint32_t stride = T / 32u;
const uint32_t rl_ws0 = warp * 8u + n0; /* local row for ws of element cols */
const uint32_t tl_xsA = mt0; /* local token rows for xs */
const uint32_t tl_xsB = mt0 + 8u;
for (uint32_t j = 0; j < stride; j++) {
/* adjacent-pairwise static stack over 32 terms in bitrev5 m order */
float s0e0 = 0, s1e0 = 0, s2e0 = 0, s3e0 = 0, s4e0 = 0;
float s0e1 = 0, s1e1 = 0, s2e1 = 0, s3e1 = 0, s4e1 = 0;
float s0e2 = 0, s1e2 = 0, s2e2 = 0, s3e2 = 0, s4e2 = 0;
float s0e3 = 0, s1e3 = 0, s2e3 = 0, s3e3 = 0, s4e3 = 0;
#pragma unroll
for (uint32_t i = 0; i < 32u; i++) {
const uint32_t m = bitrev5(i);
const uint32_t s = j + m * stride; /* slot index */
float t0 = 0.0f, t1 = 0.0f, t2 = 0.0f, t3 = 0.0f;
/* slot s sums blocks {s + k*T} sequentially (multi-term when
* blocks > T, exactly like the per-lane strided walk). */
for (uint32_t b = s; b < blocks; b += T) {
const uint32_t koff = (lane & 3u) * 4u;
const int8_t *ablk_lo = aq_lo + b * 32u;
const int8_t *ablk_hi = aq_hi + b * 32u;
const uint32_t a0 = a_lo_ok ? *(const uint32_t *)(ablk_lo + koff) : 0u;
const uint32_t a1 = a_hi_ok ? *(const uint32_t *)(ablk_hi + koff) : 0u;
const uint32_t a2 = a_lo_ok ? *(const uint32_t *)(ablk_lo + 16u + koff) : 0u;
const uint32_t a3 = a_hi_ok ? *(const uint32_t *)(ablk_hi + 16u + koff) : 0u;
const uint8_t *bq = (const uint8_t *)(b_wr + (uint64_t)b * 34u + 2u);
const uint32_t b0 = ldu32_unaligned(bq + koff);
const uint32_t b1 = ldu32_unaligned(bq + 16u + koff);
int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0;
mma_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1);
/* term = ws * xs * dot, same expression as reference */
const float ws0 = __half2float(sh_ws[rl_ws0 * ws_pitch + b]);
const float ws1 = __half2float(sh_ws[(rl_ws0 + 1u) * ws_pitch + b]);
const float xsA = sh_xs[tl_xsA * xs_pitch + b];
const float xsB = sh_xs[tl_xsB * xs_pitch + b];
t0 += ws0 * xsA * (float)c0;
t1 += ws1 * xsA * (float)c1;
t2 += ws0 * xsB * (float)c2;
t3 += ws1 * xsB * (float)c3;
}
/* static adjacent stack push (compile-time resolved) */
if ((i & 1u) == 0u) { s0e0 = t0; s0e1 = t1; s0e2 = t2; s0e3 = t3; }
else {
t0 = s0e0 + t0; t1 = s0e1 + t1; t2 = s0e2 + t2; t3 = s0e3 + t3;
if ((i & 2u) == 0u) { s1e0 = t0; s1e1 = t1; s1e2 = t2; s1e3 = t3; }
else {
t0 = s1e0 + t0; t1 = s1e1 + t1; t2 = s1e2 + t2; t3 = s1e3 + t3;
if ((i & 4u) == 0u) { s2e0 = t0; s2e1 = t1; s2e2 = t2; s2e3 = t3; }
else {
t0 = s2e0 + t0; t1 = s2e1 + t1; t2 = s2e2 + t2; t3 = s2e3 + t3;
if ((i & 8u) == 0u) { s3e0 = t0; s3e1 = t1; s3e2 = t2; s3e3 = t3; }
else {
t0 = s3e0 + t0; t1 = s3e1 + t1; t2 = s3e2 + t2; t3 = s3e3 + t3;
if ((i & 16u) == 0u) { s4e0 = t0; s4e1 = t1; s4e2 = t2; s4e3 = t3; }
else {
t0 = s4e0 + t0; t1 = s4e1 + t1; t2 = s4e2 + t2; t3 = s4e3 + t3;
/* i == 31: t is the finished x_j */
switch (j & 3u) {
case 0u: acc00 += t0; acc10 += t1; acc20 += t2; acc30 += t3; break;
case 1u: acc01 += t0; acc11 += t1; acc21 += t2; acc31 += t3; break;
case 2u: acc02 += t0; acc12 += t1; acc22 += t2; acc32 += t3; break;
default: acc03 += t0; acc13 += t1; acc23 += t2; acc33 += t3; break;
}
}
}
}
}
}
}
}
/* tail combine per T */
float r0, r1, r2, r3;
if (T == 32u) {
r0 = acc00; r1 = acc10; r2 = acc20; r3 = acc30;
} else if (T == 64u) {
r0 = acc00 + acc01; r1 = acc10 + acc11; r2 = acc20 + acc21; r3 = acc30 + acc31;
} else {
r0 = (acc00 + acc02) + (acc01 + acc03);
r1 = (acc10 + acc12) + (acc11 + acc13);
r2 = (acc20 + acc22) + (acc21 + acc23);
r3 = (acc30 + acc32) + (acc31 + acc33);
}
/* writes */
const uint64_t rowa = row0 + n0;
const uint64_t rowb = rowa + 1u;
if (tokA < n_tok) {
if (rowa < out_dim) out[tokA * out_stride + rowa] = r0;
if (rowb < out_dim) out[tokA * out_stride + rowb] = r1;
}
if (tokB < n_tok) {
if (rowa < out_dim) out[tokB * out_stride + rowa] = r2;
if (rowb < out_dim) out[tokB * out_stride + rowb] = r3;
}
}
static int cuda_q4_mma_ok(void) {
/* All tiers on this host use the same cubin/PTX and GPU model. */
static int cached = -1;
if (cached < 0) {
if (getenv("DS4_CUDA_MOE_NO_Q4_MMA") != NULL) {
cached = 0;
} else {
int dev = 0, major = 0, minor = 0;
cudaFuncAttributes fn_attr;
cudaGetDevice(&dev);
cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev);
cached = (major > 7 || (major == 7 && minor >= 5)) &&
cudaFuncGetAttributes(&fn_attr,
matmul_q8_0_mma_exact_kernel<32>) == cudaSuccess &&
fn_attr.ptxVersion >= 75;
}
}
return cached;
}
static int cuda_q8_mma_attr_ready[DS4_MAX_GPUS][4];
static int cuda_q8_mma_try_launch(
float *out,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t blocks,
uint64_t a_stride_blocks,
uint64_t out_stride,
uint32_t T) {
static int disabled = -1;
if (disabled < 0) disabled = getenv("DS4_CUDA_NO_Q8_MMA") != NULL ? 1 : 0;
if (disabled || !cuda_q4_mma_ok()) return 0;
if ((in_dim & 31u) != 0u || blocks > 256u || n_tok < 8u) return 0;
if (((uintptr_t)w & 1u) || ((uintptr_t)xq & 3u) || ((uintptr_t)xscale & 3u)) return 0;
const size_t shmem = (size_t)(64u * (blocks + 2u) * 2u + 16u * (blocks + 1u) * 4u);
int dev = 0;
cudaGetDevice(&dev);
if (dev < 0 || dev >= DS4_MAX_GPUS) return 0;
const int ti = T == 32u ? 0 : (T == 64u ? 1 : (T == 128u ? 2 : 3));
dim3 grid(((unsigned)out_dim + 63u) / 64u, ((unsigned)n_tok + 15u) / 16u, 1);
/* Keep eight adjacent weight tiles hot while visiting the token tiles.
* Walking the entire output matrix per token tile thrashes GB10's L2. */
if (in_dim >= 4096u && out_dim >= 4096u && n_tok >= 128u)
grid = dim3(((unsigned)n_tok + 15u) / 16u * 8u,
(((unsigned)out_dim + 63u) / 64u + 7u) / 8u, 1);
#define DS4_Q8_MMA_LAUNCH(TT) \
do { \
if (!cuda_q8_mma_attr_ready[dev][ti]) { \
cudaFuncAttributes fn_attr; \
if (cudaFuncGetAttributes(&fn_attr, matmul_q8_0_mma_exact_kernel<TT>) != cudaSuccess || \
fn_attr.binaryVersion < 80 || fn_attr.ptxVersion < 80) { \
disabled = 1; \
return 0; \
} \
if (cudaFuncSetAttribute(matmul_q8_0_mma_exact_kernel<TT>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, \
(int)(64u * 258u * 2u + 16u * 257u * 4u)) != cudaSuccess) { \
disabled = 1; \
return 0; \
} \
cuda_q8_mma_attr_ready[dev][ti] = 1; \
} \
matmul_q8_0_mma_exact_kernel<TT><<<grid, 256, shmem>>>( \
out, w, xq, xscale, in_dim, out_dim, n_tok, blocks, \
a_stride_blocks, out_stride); \
} while (0)
if (T == 32u) DS4_Q8_MMA_LAUNCH(32u);
else if (T == 64u) DS4_Q8_MMA_LAUNCH(64u);
else if (T == 128u) DS4_Q8_MMA_LAUNCH(128u);
else DS4_Q8_MMA_LAUNCH(256u);
#undef DS4_Q8_MMA_LAUNCH
return cuda_ok(cudaGetLastError(), "matmul_q8_0 mma launch") ? 1 : -1;
}
__global__ static void dequant_q8_0_to_f16_kernel(
__half *out,
const unsigned char *w,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = in_dim * out_dim;
if (gid >= n) return;
uint64_t row = gid / in_dim;
uint64_t i = gid - row * in_dim;
uint64_t b = i / 32;
uint64_t j = i - b * 32;
const unsigned char *blk = w + (row * blocks + b) * 34;
const __half scale = *(const __half *)blk;
const int8_t q = *(const int8_t *)(blk + 2 + j);
out[gid] = __hmul(scale, __float2half((float)q));
}
__global__ static void dequant_q8_0_to_f32_kernel(
float *out,
const unsigned char *w,
uint64_t in_dim,
uint64_t out_dim,
uint64_t blocks) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = in_dim * out_dim;
if (gid >= n) return;
uint64_t row = gid / in_dim;
uint64_t i = gid - row * in_dim;
uint64_t b = i / 32;
uint64_t j = i - b * 32;
const unsigned char *blk = w + (row * blocks + b) * 34;
const float scale = __half2float(*(const __half *)blk);
const int8_t q = *(const int8_t *)(blk + 2 + j);
out[gid] = scale * (float)q;
}
__global__ static void grouped_q8_0_a_preq_warp8_kernel(
float *low,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint32_t n_tokens,
uint64_t blocks,
int use_dp4a) {
const uint64_t row = (uint64_t)blockIdx.x * 8u + (threadIdx.x >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y;
const uint32_t lane = threadIdx.x & 31u;
const uint64_t low_dim = (uint64_t)n_groups * rank;
if (row >= low_dim || tok >= n_tokens) return;
const uint64_t group = row / rank;
const uint64_t row_in_group = row - group * rank;
const unsigned char *wr = w + (group * rank + row_in_group) * blocks * 34;
const uint64_t xrow = tok * (uint64_t)n_groups + group;
const int8_t *xqr = xq + xrow * blocks * 32;
const float *xsr = xscale + xrow * blocks;
float acc = 0.0f;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32;
const uint64_t bn = group_dim - i0 < 32 ? group_dim - i0 : 32;
const __half *scale_h = (const __half *)(wr + b * 34);
const int8_t *qs = (const int8_t *)(wr + b * 34 + 2);
const int8_t *xqb = xqr + b * 32;
int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
acc = warp_sum_f32(acc);
if (lane == 0) low[tok * low_dim + row] = acc;
}
__global__ static void grouped_q8_0_a_preq_warp8_tok2_kernel(
float *low,
const unsigned char *w,
const int8_t *xq,
const float *xscale,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint32_t n_tokens,
uint64_t blocks,
int use_dp4a) {
const uint32_t tid_in_tok = threadIdx.x & 255u;
const uint64_t row = (uint64_t)blockIdx.x * 8u + (tid_in_tok >> 5u);
const uint64_t tok = (uint64_t)blockIdx.y * 2u + (threadIdx.x >> 8u);
const uint32_t lane = threadIdx.x & 31u;
const uint64_t low_dim = (uint64_t)n_groups * rank;
float acc = 0.0f;
if (row < low_dim && tok < n_tokens) {
const uint64_t group = row / rank;
const uint64_t row_in_group = row - group * rank;
const unsigned char *wr = w + (group * rank + row_in_group) * blocks * 34u;
const uint64_t xrow = tok * (uint64_t)n_groups + group;
const int8_t *xqr = xq + xrow * blocks * 32u;
const float *xsr = xscale + xrow * blocks;
for (uint64_t b = lane; b < blocks; b += 32u) {
const uint64_t i0 = b * 32u;
const uint64_t bn = group_dim - i0 < 32u ? group_dim - i0 : 32u;
const __half *scale_h = (const __half *)(wr + b * 34u);
const int8_t *qs = (const int8_t *)(wr + b * 34u + 2u);
const int8_t *xqb = xqr + b * 32u;
const int dot = dot_i8_block(qs, xqb, bn, use_dp4a);
acc += __half2float(*scale_h) * xsr[b] * (float)dot;
}
}
acc = warp_sum_f32(acc);
if (lane == 0 && row < low_dim && tok < n_tokens) {
low[tok * low_dim + row] = acc;
}
}
__global__ static void rms_norm_plain_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) {
uint32_t row = blockIdx.x;
if (row >= rows) return;
const float *xr = x + (uint64_t)row * n;
float *orow = out + (uint64_t)row * n;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
float scale = rsqrtf(partial[0] / (float)n + eps);
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale;
}
}
/* Latency-optimized RMS norm for the common n==4096 decode shape: one global
* read pass with register-batched loads, same per-thread accumulation order
* and shared-memory tree as rms_norm_plain_kernel (bit-identical, fuzz
* checked). */
__global__ static void rms_norm_plain_fast4096_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) {
uint32_t row = blockIdx.x;
if (row >= rows) return;
const float *xr = x + (uint64_t)row * n;
float *orow = out + (uint64_t)row * n;
float v[16];
#pragma unroll
for (uint32_t j = 0; j < 16u; j++) v[j] = xr[threadIdx.x + j * 256u];
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 16u; j++) sum += v[j] * v[j];
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
float scale = rsqrtf(partial[0] / (float)n + eps);
#pragma unroll
for (uint32_t j = 0; j < 16u; j++) orow[threadIdx.x + j * 256u] = v[j] * scale;
}
/* Batched-load RMS norm for larger rows (n multiple of 2048, e.g. the 16384
* HC-concatenated decode rows). Two passes like the reference kernel, but
* eight independent loads are issued per accumulation group; the per-thread
* accumulation order (ascending i with stride 256) is unchanged, so results
* are bit-identical. */
__global__ static void rms_norm_plain_batch8_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) {
uint32_t row = blockIdx.x;
if (row >= rows) return;
const float *xr = x + (uint64_t)row * n;
float *orow = out + (uint64_t)row * n;
float sum = 0.0f;
#pragma unroll 1
for (uint32_t i = threadIdx.x; i < n; i += 2048u) {
const float v0 = xr[i];
const float v1 = xr[i + 256u];
const float v2 = xr[i + 512u];
const float v3 = xr[i + 768u];
const float v4 = xr[i + 1024u];
const float v5 = xr[i + 1280u];
const float v6 = xr[i + 1536u];
const float v7 = xr[i + 1792u];
sum += v0 * v0;
sum += v1 * v1;
sum += v2 * v2;
sum += v3 * v3;
sum += v4 * v4;
sum += v5 * v5;
sum += v6 * v6;
sum += v7 * v7;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
float scale = rsqrtf(partial[0] / (float)n + eps);
#pragma unroll 1
for (uint32_t i = threadIdx.x; i < n; i += 2048u) {
const float v0 = xr[i];
const float v1 = xr[i + 256u];
const float v2 = xr[i + 512u];
const float v3 = xr[i + 768u];
const float v4 = xr[i + 1024u];
const float v5 = xr[i + 1280u];
const float v6 = xr[i + 1536u];
const float v7 = xr[i + 1792u];
orow[i] = v0 * scale;
orow[i + 256u] = v1 * scale;
orow[i + 512u] = v2 * scale;
orow[i + 768u] = v3 * scale;
orow[i + 1024u] = v4 * scale;
orow[i + 1280u] = v5 * scale;
orow[i + 1536u] = v6 * scale;
orow[i + 1792u] = v7 * scale;
}
}
/* Emit the same normalized values as rms_norm_plain_batch8_kernel directly
* as FP16 activations for the following cuBLAS projection. */
__global__ static void rms_norm_plain_f16_batch8_kernel(
__half *out,
const float *x,
uint32_t n,
uint32_t rows,
float eps) {
const uint32_t row = blockIdx.x;
if (row >= rows) return;
const float *xr = x + (uint64_t)row * n;
__half *orow = out + (uint64_t)row * n;
float sum = 0.0f;
#pragma unroll 1
for (uint32_t i = threadIdx.x; i < n; i += 2048u) {
const float v0 = xr[i];
const float v1 = xr[i + 256u];
const float v2 = xr[i + 512u];
const float v3 = xr[i + 768u];
const float v4 = xr[i + 1024u];
const float v5 = xr[i + 1280u];
const float v6 = xr[i + 1536u];
const float v7 = xr[i + 1792u];
sum += v0 * v0;
sum += v1 * v1;
sum += v2 * v2;
sum += v3 * v3;
sum += v4 * v4;
sum += v5 * v5;
sum += v6 * v6;
sum += v7 * v7;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial[threadIdx.x] += partial[threadIdx.x + stride];
}
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)n + eps);
#pragma unroll 1
for (uint32_t i = threadIdx.x; i < n; i += 2048u) {
orow[i] = __float2half(xr[i] * scale);
orow[i + 256u] = __float2half(xr[i + 256u] * scale);
orow[i + 512u] = __float2half(xr[i + 512u] * scale);
orow[i + 768u] = __float2half(xr[i + 768u] * scale);
orow[i + 1024u] = __float2half(xr[i + 1024u] * scale);
orow[i + 1280u] = __float2half(xr[i + 1280u] * scale);
orow[i + 1536u] = __float2half(xr[i + 1536u] * scale);
orow[i + 1792u] = __float2half(xr[i + 1792u] * scale);
}
}
__global__ static void rms_norm_weight_kernel(float *out, const float *x, const float *w, uint32_t n, uint32_t rows, float eps) {
uint32_t row = blockIdx.x;
if (row >= rows) return;
const float *xr = x + (uint64_t)row * n;
float *orow = out + (uint64_t)row * n;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
sum = block_sum_f32_256(sum, partial);
float scale = rsqrtf(sum / (float)n + eps);
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale * w[i];
}
}
__global__ static void dsv4_qkv_rms_norm_rows_kernel(
float *q_out,
const float *q,
const float *q_w,
uint32_t q_n,
float *kv_out,
const float *kv,
const float *kv_w,
uint32_t kv_n,
uint32_t rows,
float eps) {
const uint32_t row = blockIdx.x;
const uint32_t which = blockIdx.y;
if (row >= rows || which > 1u) return;
const uint32_t n = which == 0u ? q_n : kv_n;
const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
const float *w = which == 0u ? q_w : kv_w;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
const float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)n + eps);
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale * w[i];
}
}
__global__ static void head_rms_norm_kernel(float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
uint32_t row = blockIdx.x;
if (row >= n_tok * n_head) return;
float *xr = x + (uint64_t)row * head_dim;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
float scale = rsqrtf(partial[0] / (float)head_dim + eps);
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) xr[i] *= scale;
}
__device__ static float rope_yarn_ramp_dev(float low, float high, int i0);
__global__ static void dsv4_qkv_rms_norm_rows_kv_rope_kernel(
float *q_out,
const float *q,
const float *q_w,
uint32_t q_n,
float *kv_out,
const float *kv,
const float *kv_w,
uint32_t kv_n,
uint32_t rows,
uint32_t kv_n_head,
uint32_t kv_head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
int inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
const uint32_t row = blockIdx.x;
const uint32_t which = blockIdx.y;
if (row >= rows || which > 1u) return;
const uint32_t n = which == 0u ? q_n : kv_n;
const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
const float *w = which == 0u ? q_w : kv_w;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
const float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)n + eps);
if (which == 0u) {
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
orow[i] = xr[i] * scale * w[i];
}
return;
}
const uint32_t n_nope = kv_head_dim - n_rot;
for (uint32_t h = 0; h < kv_n_head; h++) {
const uint32_t head_base = h * kv_head_dim;
for (uint32_t d = threadIdx.x; d < n_nope; d += blockDim.x) {
const uint32_t i = head_base + d;
orow[i] = xr[i] * scale * w[i];
}
}
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
const uint32_t pairs_per_head = n_rot / 2u;
const uint32_t total_pairs = kv_n_head * pairs_per_head;
for (uint32_t p = threadIdx.x; p < total_pairs; p += blockDim.x) {
const uint32_t h = p / pairs_per_head;
const uint32_t pair = p - h * pairs_per_head;
const uint32_t d = n_nope + pair * 2u;
const uint32_t i0 = h * kv_head_dim + d;
const uint32_t i = pair * 2u;
float theta_extrap = (float)(pos0 + row) * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
const float x0 = xr[i0] * scale * w[i0];
const float x1 = xr[i0 + 1u] * scale * w[i0 + 1u];
orow[i0] = x0 * c - x1 * s;
orow[i0 + 1u] = x0 * s + x1 * c;
}
}
__global__ static void head_rms_norm_rope_tail_kernel(
float *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
int inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
uint32_t row = blockIdx.x;
if (row >= n_tok * n_head) return;
uint32_t t = row / n_head;
float *xr = x + (uint64_t)row * head_dim;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
float v = xr[i];
sum += v * v;
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float scale = rsqrtf(partial[0] / (float)head_dim + eps);
const uint32_t n_nope = head_dim - n_rot;
for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) {
xr[i] *= scale;
}
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) {
uint32_t i = pair * 2u;
float theta_extrap = (float)(pos0 + t) * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float *tail = xr + n_nope;
float x0 = tail[i] * scale;
float x1 = tail[i + 1] * scale;
tail[i] = x0 * c - x1 * s;
tail[i + 1] = x0 * s + x1 * c;
}
}
__device__ static float rope_yarn_ramp_dev(float low, float high, int i0) {
float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
}
__global__ static void rope_tail_kernel(
float *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t pos_stride,
uint32_t n_ctx_orig,
int inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t pairs = n_tok * n_head * (n_rot / 2);
if (gid >= pairs) return;
uint32_t pair = gid % (n_rot / 2);
uint32_t tmp = gid / (n_rot / 2);
uint32_t h = tmp % n_head;
uint32_t t = tmp / n_head;
uint32_t n_nope = head_dim - n_rot;
uint32_t i = pair * 2;
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1), corr1);
}
float theta_extrap = (float)(pos0 + t * pos_stride) * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float *tail = x + ((uint64_t)t * n_head + h) * head_dim + n_nope;
float x0 = tail[i];
float x1 = tail[i + 1];
tail[i] = x0 * c - x1 * s;
tail[i + 1] = x0 * s + x1 * c;
}
__global__ static void rope_tail_decode_rows_kernel(
float *x,
cuda_attention_decode_row_table rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
int inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t pairs = n_rows * n_head * (n_rot / 2u);
if (gid >= pairs) return;
const uint32_t pair = gid % (n_rot / 2u);
const uint32_t tmp = gid / (n_rot / 2u);
const uint32_t h = tmp % n_head;
const uint32_t row = tmp / n_head;
const uint32_t n_nope = head_dim - n_rot;
const uint32_t i = pair * 2u;
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
const float denom = 2.0f * logf(freq_base);
corr0 = floorf((float)n_rot *
logf((float)n_ctx_orig /
(beta_fast * 2.0f * (float)M_PI)) / denom);
corr1 = ceilf((float)n_rot *
logf((float)n_ctx_orig /
(beta_slow * 2.0f * (float)M_PI)) / denom);
corr0 = fmaxf(0.0f, corr0);
corr1 = fminf((float)(n_rot - 1u), corr1);
}
const float theta_extrap = (float)rows.row[row].pos *
powf(freq_base, -((float)i) / (float)n_rot);
const float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
const float ramp_mix =
rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
const float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
if (inverse) s = -s;
float *tail = x + ((uint64_t)row * n_head + h) * head_dim + n_nope;
const float x0 = tail[i];
const float x1 = tail[i + 1u];
tail[i] = x0 * c - x1 * s;
tail[i + 1u] = x0 * s + x1 * c;
}
__device__ static float dsv4_e4m3fn_value_dev(int i) {
int exp = (i >> 3) & 15;
int mant = i & 7;
if (exp == 0) return (float)mant * 0.001953125f;
return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
}
__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
float sign = x < 0.0f ? -1.0f : 1.0f;
float ax = fminf(fabsf(x), 448.0f);
int lo = 0, hi = 126;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
else hi = mid - 1;
}
int best = lo;
if (best < 126) {
float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
}
return sign * dsv4_e4m3fn_value_dev(best);
}
__device__ static float dsv4_e2m1fn_value_dev(int i) {
switch (i & 7) {
case 0: return 0.0f;
case 1: return 0.5f;
case 2: return 1.0f;
case 3: return 1.5f;
case 4: return 2.0f;
case 5: return 3.0f;
case 6: return 4.0f;
default: return 6.0f;
}
}
__device__ static float dsv4_e2m1fn_dequant_dev(float x) {
float sign = x < 0.0f ? -1.0f : 1.0f;
float ax = fminf(fabsf(x), 6.0f);
int best = 0;
float best_diff = fabsf(ax - dsv4_e2m1fn_value_dev(0));
for (int i = 1; i < 8; i++) {
float diff = fabsf(ax - dsv4_e2m1fn_value_dev(i));
if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
best = i;
best_diff = diff;
}
}
return sign * dsv4_e2m1fn_value_dev(best);
}
__device__ static float model_scalar_dev(const void *base, uint64_t offset, uint32_t type, uint64_t idx) {
const char *p = (const char *)base + offset;
if (type == 1u) return __half2float(((const __half *)p)[idx]);
return ((const float *)p)[idx];
}
__device__ static float rope_yarn_ramp_cpu_equiv_dev(float low, float high, int i0) {
float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
}
__device__ static DS4_CUDA_UNUSED void rope_tail_one_dev(float *x, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint32_t n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) {
uint32_t n_nope = head_dim - n_rot;
float corr0 = 0.0f, corr1 = 0.0f;
if (ext_factor != 0.0f) {
float denom = 2.0f * logf(freq_base);
corr0 = fmaxf(0.0f, floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom));
corr1 = fminf((float)(n_rot - 1), ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom));
}
for (uint32_t i = 0; i < n_rot; i += 2) {
float theta_extrap = (float)pos * powf(freq_base, -((float)i) / (float)n_rot);
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
float mscale = attn_factor;
if (ext_factor != 0.0f) {
float mix = rope_yarn_ramp_cpu_equiv_dev(corr0, corr1, (int)i) * ext_factor;
theta = theta_interp * (1.0f - mix) + theta_extrap * mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
float c = cosf(theta) * mscale;
float s = sinf(theta) * mscale;
float x0 = x[n_nope + i];
float x1 = x[n_nope + i + 1];
x[n_nope + i] = x0 * c - x1 * s;
x[n_nope + i + 1] = x0 * s + x1 * c;
}
}
__device__ static void fp8_kv_quantize_row(
float *xr,
uint32_t head_dim,
uint32_t n_rot,
float *scratch) {
uint32_t tid = threadIdx.x;
uint32_t n_nope = head_dim - n_rot;
for (uint32_t off = 0; off < n_nope; off += 64) {
float v = 0.0f;
if (off + tid < n_nope) v = xr[off + tid];
scratch[tid] = off + tid < n_nope ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
__syncthreads();
}
float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
if (off + tid < n_nope) {
float q = dsv4_e4m3fn_dequant_dev(fminf(448.0f, fmaxf(-448.0f, v / scale))) * scale;
xr[off + tid] = q;
}
__syncthreads();
}
}
__global__ static void fp8_kv_quantize_kernel(
float *x,
uint32_t n_tok,
uint32_t head_dim,
uint32_t n_rot) {
uint32_t row = blockIdx.x;
if (row >= n_tok) return;
__shared__ float scratch[64];
fp8_kv_quantize_row(
x + (uint64_t)row * head_dim, head_dim, n_rot, scratch);
}
__global__ static void fp8_kv_quantize_store_rows_kernel(
float *x,
cuda_attention_decode_row_table rows,
uint32_t n_rows,
uint32_t head_dim,
uint32_t n_rot) {
const uint32_t row = blockIdx.x;
if (row >= n_rows) return;
__shared__ float scratch[64];
float *xr = x + (uint64_t)row * head_dim;
fp8_kv_quantize_row(xr, head_dim, n_rot, scratch);
const ds4_gpu_attention_decode_row dsc = rows.row[row];
float *raw = (float *)(uintptr_t)dsc.raw_kv;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
raw[(uint64_t)dsc.raw_start * head_dim + d] =
__half2float(__float2half(xr[d]));
}
}
__global__ static void indexer_hadamard_fp4_kernel(float *x, uint32_t n_rows, uint32_t head_dim) {
uint32_t row = blockIdx.x;
uint32_t tid = threadIdx.x;
if (row >= n_rows || head_dim != 128u || tid >= 128u) return;
__shared__ float vals[128];
__shared__ float absbuf[128];
float *xr = x + (uint64_t)row * head_dim;
vals[tid] = xr[tid];
__syncthreads();
for (uint32_t stride = 1u; stride < 128u; stride <<= 1u) {
if ((tid & stride) == 0u) {
uint32_t base = (tid & ~(2u * stride - 1u)) + (tid & (stride - 1u));
float a = vals[base];
float b = vals[base + stride];
vals[base] = a + b;
vals[base + stride] = a - b;
}
__syncthreads();
}
float v = vals[tid] * 0.08838834764831845f;
uint32_t fp4_block = tid >> 5u;
uint32_t lane = tid & 31u;
uint32_t block_base = fp4_block * 32u;
absbuf[tid] = fabsf(v);
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride) {
absbuf[block_base + lane] = fmaxf(absbuf[block_base + lane],
absbuf[block_base + lane + stride]);
}
__syncthreads();
}
float amax = fmaxf(absbuf[block_base], 7.052966104933725e-38f);
float scale = exp2f(ceilf(log2f(amax / 6.0f)));
xr[tid] = dsv4_e2m1fn_dequant_dev(fminf(6.0f, fmaxf(-6.0f, v / scale))) * scale;
}
__global__ static void store_raw_kv_batch_kernel(float *raw, const float *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * head_dim;
if (gid >= n) return;
uint32_t d = gid % head_dim;
uint32_t t = gid / head_dim;
uint32_t row = (pos0 + t) % raw_cap;
raw[(uint64_t)row * head_dim + d] = __half2float(__float2half(kv[(uint64_t)t * head_dim + d]));
}
__global__ static void attention_prefill_raw_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
uint32_t n_tokens,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || h >= n_head) return;
uint32_t raw_count = t + 1 < window ? t + 1 : window;
uint32_t raw_start = t + 1 - raw_count;
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
__shared__ float scores[256];
__shared__ float partial[128];
__shared__ float max_s;
__shared__ float denom;
float scale = rsqrtf((float)head_dim);
float local_max = sinks[h];
__syncthreads();
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
const float *kv = raw_kv + (uint64_t)(raw_start + r) * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d];
scores[r] = dot * scale;
local_max = fmaxf(local_max, scores[r]);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
if (threadIdx.x == 0) {
float den = expf(sinks[h] - max_s);
for (uint32_t r = 0; r < raw_count; r++) {
scores[r] = expf(scores[r] - max_s);
den += scores[r];
}
denom = den;
}
__syncthreads();
float *oh = heads + ((uint64_t)t * n_head + h) * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r];
}
oh[d] = acc / denom;
}
}
__global__ static void attention_prefill_mixed_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || h >= n_head) return;
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
uint32_t raw_start = (window != 0 && t + 1u > window) ? t + 1u - window : 0u;
uint32_t raw_count = t + 1u - raw_start;
uint32_t visible_comp = (t + 1u) / ratio;
if (visible_comp > n_comp) visible_comp = n_comp;
__shared__ float scores[512];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
float scale = rsqrtf((float)head_dim);
float local_max = sinks[h];
uint32_t n_score = raw_count + visible_comp;
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
const float *kvrow = raw_kv + (uint64_t)(raw_start + r) * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
scores[r] = dot * scale;
local_max = fmaxf(local_max, scores[r]);
}
for (uint32_t c = threadIdx.x; c < visible_comp; c += blockDim.x) {
float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
float s = -INFINITY;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)c * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale + add;
}
scores[raw_count + c] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
float *oh = heads + ((uint64_t)t * n_head + h) * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)(raw_start + r) * head_dim + d] * scores[r];
for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c];
oh[d] = acc / denom;
}
}
__global__ static void attention_prefill_raw_softmax_kernel(
float *scores,
const float *sinks,
uint32_t n_tokens,
uint32_t window,
uint32_t n_keys) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens) return;
float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys;
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
float local_max = sinks[h];
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
bool valid = k <= t && (window == 0 || t - k < window);
float s = valid ? row[k] : -INFINITY;
row[k] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f;
row[k] = p;
den_local += p;
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom;
}
__global__ static void attention_prefill_mixed_softmax_kernel(
float *scores,
const float *sinks,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_keys) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || ratio == 0) return;
float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys;
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
float local_max = sinks[h];
const uint32_t visible_comp = (t + 1u) / ratio;
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
float s = -INFINITY;
if (k < n_tokens) {
if (k <= t && (window == 0 || t - k < window)) s = row[k];
} else {
uint32_t c = k - n_tokens;
if (c < n_comp && c < visible_comp) {
float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
if (add > -1.0e20f) s = row[k] + add;
}
}
row[k] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
float p = isfinite(row[k]) ? expf(row[k] - max_s) : 0.0f;
row[k] = p;
den_local += p;
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) row[k] /= denom;
}
__global__ static void attention_prefill_pack_mixed_kv_kernel(
float *dst,
const float *raw_kv,
const float *comp_kv,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t head_dim) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)(n_tokens + n_comp) * head_dim;
if (gid >= n) return;
uint32_t d = gid % head_dim;
uint32_t r = gid / head_dim;
dst[gid] = r < n_tokens ? raw_kv[(uint64_t)r * head_dim + d]
: comp_kv[(uint64_t)(r - n_tokens) * head_dim + d];
}
__global__ static void attention_visual_pack_mixed_kv_kernel(
float *dst,
const float *raw_kv,
const float *comp_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t head_dim) {
const uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t count = (uint64_t)(n_raw + n_comp) * head_dim;
if (gid >= count) return;
const uint32_t d = gid % head_dim;
const uint32_t row = gid / head_dim;
if (row < n_raw) {
const uint32_t physical = (raw_start + row) % raw_cap;
dst[gid] = raw_kv[(uint64_t)physical * head_dim + d];
} else {
dst[gid] = comp_kv[(uint64_t)(row - n_raw) * head_dim + d];
}
}
__global__ static void attention_visual_mixed_softmax_kernel(
float *scores,
const float *sinks,
const float *comp_mask,
const uint32_t *raw_bounds,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t first_raw_pos,
uint32_t n_raw,
uint32_t n_comp,
uint32_t ratio,
uint32_t n_keys) {
const uint32_t t = blockIdx.x;
const uint32_t h = blockIdx.y;
if (t >= n_tokens) return;
float *row = scores + ((uint64_t)h * n_tokens + t) * n_keys;
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
const uint32_t raw_lo = raw_bounds[2u * t];
const uint32_t raw_hi = raw_bounds[2u * t + 1u];
const uint32_t qpos = pos0 + t;
uint32_t visible_comp = ratio ? (qpos + 1u) / ratio : 0u;
if (visible_comp > n_comp) visible_comp = n_comp;
float local_max = sinks[h];
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
float score = -INFINITY;
if (k < n_raw) {
const uint32_t kpos = first_raw_pos + k;
if (kpos >= raw_lo && kpos <= raw_hi) score = row[k];
} else {
const uint32_t c = k - n_raw;
if (c < visible_comp) {
const float add = use_comp_mask
? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
if (add > -1.0e20f) score = row[k] + add;
}
}
row[k] = score;
local_max = fmaxf(local_max, score);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
if (threadIdx.x < stride)
partial[threadIdx.x] = fmaxf(partial[threadIdx.x],
partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0u) max_s = partial[0];
__syncthreads();
float sum = 0.0f;
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x) {
const float probability = isfinite(row[k])
? expf(row[k] - max_s) : 0.0f;
row[k] = probability;
sum += probability;
}
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x / 2u; stride != 0u; stride >>= 1u) {
if (threadIdx.x < stride)
partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0u)
denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
for (uint32_t k = threadIdx.x; k < n_keys; k += blockDim.x)
row[k] /= denom;
}
__global__ static void attention_prefill_unpack_heads_kernel(
float *heads,
const float *tmp,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_head * head_dim;
if (gid >= n) return;
uint32_t d = gid % head_dim;
uint64_t q = gid / head_dim;
uint32_t h = q % n_head;
uint32_t t = q / n_head;
heads[gid] = tmp[((uint64_t)h * n_tokens + t) * head_dim + d];
}
__global__ static void attention_pack_group_heads_f16_kernel(
__half *dst,
const float *heads,
uint32_t n_tokens,
uint32_t n_groups,
uint32_t group_dim) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_groups * n_tokens * group_dim;
if (gid >= n) return;
uint32_t d = gid % group_dim;
uint64_t q = gid / group_dim;
uint32_t t = q % n_tokens;
uint32_t g = q / n_tokens;
dst[gid] = __float2half(heads[((uint64_t)t * n_groups + g) * group_dim + d]);
}
__global__ static void attention_unpack_group_low_kernel(
float *low,
const float *tmp,
uint32_t n_tokens,
uint32_t n_groups,
uint32_t rank) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_groups * n_tokens * rank;
if (gid >= n) return;
uint32_t r = gid % rank;
uint64_t q = gid / rank;
uint32_t t = q % n_tokens;
uint32_t g = q / n_tokens;
uint32_t low_dim = n_groups * rank;
low[(uint64_t)t * low_dim + (uint64_t)g * rank + r] = tmp[gid];
}
__global__ static void attention_decode_mixed_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t score_lanes_single) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || h >= n_head) return;
const bool single_all = (n_tokens == 1u && ratio == 0u);
uint32_t qpos = pos0 + t;
uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
__shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP];
__shared__ uint32_t raw_rows[256];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x;
const bool score_thread = threadIdx.x < score_threads;
float scale = rsqrtf((float)head_dim);
if (threadIdx.x == 0) {
raw_count = 0;
raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
if (score_thread) {
for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
}
__syncthreads();
uint32_t n_score = raw_count + visible_comp;
float local_max = sinks[h];
if (score_thread) {
if (visible_comp == 0 || (n_tokens == 1u && score_lanes_single == 0u)) {
for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) {
const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
scores[r] = dot * scale;
local_max = fmaxf(local_max, scores[r]);
}
for (uint32_t c = threadIdx.x; c < visible_comp; c += score_threads) {
float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
float s = -INFINITY;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)c * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale + add;
}
scores[raw_count + c] = s;
local_max = fmaxf(local_max, s);
}
} else if (n_tokens == 1u && score_lanes_single == 4u) {
uint32_t qlane = threadIdx.x & 3u;
uint32_t qgroup = threadIdx.x >> 2u;
for (uint32_t row0 = 0; row0 < n_score; row0 += 64u) {
uint32_t row = row0 + qgroup;
if (row < n_score) {
float add = 0.0f;
const float *kvrow = NULL;
if (row < raw_count) {
kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim;
} else {
uint32_t c = row - raw_count;
add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim;
}
float s = -INFINITY;
if (kvrow) {
float dot = 0.0f;
for (uint32_t d = qlane; d < head_dim; d += 4u) dot += qh[d] * kvrow[d];
const uint32_t mask = 0xfu << (threadIdx.x & 28u);
dot += __shfl_down_sync(mask, dot, 2, 4);
dot += __shfl_down_sync(mask, dot, 1, 4);
s = dot * scale + add;
}
if (qlane == 0) scores[row] = s;
}
}
__syncthreads();
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
local_max = fmaxf(local_max, scores[i]);
}
} else {
uint32_t qlane = threadIdx.x & 7u;
uint32_t qgroup = threadIdx.x >> 3u;
for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) {
uint32_t row = row0 + qgroup;
if (row < n_score) {
float add = 0.0f;
const float *kvrow = NULL;
if (row < raw_count) {
kvrow = raw_kv + (uint64_t)raw_rows[row] * head_dim;
} else {
uint32_t c = row - raw_count;
add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
if (add > -1.0e20f) kvrow = comp_kv + (uint64_t)c * head_dim;
}
float s = -INFINITY;
if (kvrow) {
float dot = 0.0f;
for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d];
const uint32_t mask = 0xffu << (threadIdx.x & 24u);
for (uint32_t off = 4u; off > 0u; off >>= 1u) {
dot += __shfl_down_sync(mask, dot, off, 8);
}
s = dot * scale + add;
}
if (qlane == 0) scores[row] = s;
}
}
__syncthreads();
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
local_max = fmaxf(local_max, scores[i]);
}
}
}
if (score_thread) partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
}
if (score_thread) partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
float *oh = heads + ((uint64_t)t * n_head + h) * head_dim;
if (head_dim == 512u && blockDim.x >= 512u) {
uint32_t d = threadIdx.x;
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc += kv[d] * s;
}
for (uint32_t c = 0; c < visible_comp; c++) {
float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc += kv[d] * s;
}
oh[d] = acc / denom;
} else if (head_dim == 512u && blockDim.x == 256u) {
uint32_t d0 = threadIdx.x;
uint32_t d1 = d0 + 256u;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
for (uint32_t c = 0; c < visible_comp; c++) {
float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
oh[d0] = acc0 / denom;
oh[d1] = acc1 / denom;
} else {
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r];
for (uint32_t c = 0; c < visible_comp; c++) acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c];
oh[d] = acc / denom;
}
}
}
__global__ static void attention_decode_score_split_scores_kernel(
float *score_out,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
const uint32_t h = blockIdx.y;
const uint32_t j = blockIdx.z;
if (h >= n_head || j >= S) return;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return;
const uint32_t qbase = n_score / S;
const uint32_t rem = n_score % S;
const uint32_t g0 = j * qbase + (j < rem ? j : rem);
const uint32_t cnt = qbase + (j < rem ? 1u : 0u);
const uint32_t g1 = g0 + cnt;
const float *qh = q + (uint64_t)h * head_dim;
float *row_scores = score_out + (uint64_t)h * n_score;
const float scale = rsqrtf((float)head_dim);
for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) {
float s = -INFINITY;
if (g < raw_count) {
const uint32_t raw_row =
(raw_start + raw_first_idx + g) % raw_cap;
const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale;
} else {
const uint32_t cidx = g - raw_count;
const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)cidx * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale + add;
}
}
row_scores[g] = s;
}
}
__device__ __forceinline__ float ds4_dot_scalar_ldg(
const float *a,
const float *b,
uint32_t n) {
float dot = 0.0f;
for (uint32_t d = 0; d < n; d++) dot += __ldg(a + d) * __ldg(b + d);
return dot;
}
__global__ static void attention_decode_score_split_scores_ldg_kernel(
float *score_out,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
const uint32_t h = blockIdx.y;
const uint32_t j = blockIdx.z;
if (h >= n_head || j >= S) return;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return;
const uint32_t qbase = n_score / S;
const uint32_t rem = n_score % S;
const uint32_t g0 = j * qbase + (j < rem ? j : rem);
const uint32_t cnt = qbase + (j < rem ? 1u : 0u);
const uint32_t g1 = g0 + cnt;
const float *qh = q + (uint64_t)h * head_dim;
float *row_scores = score_out + (uint64_t)h * n_score;
const float scale = rsqrtf((float)head_dim);
for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) {
float s = -INFINITY;
if (g < raw_count) {
const uint32_t raw_row =
(raw_start + raw_first_idx + g) % raw_cap;
const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim;
const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim);
s = dot * scale;
} else {
const uint32_t cidx = g - raw_count;
const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)cidx * head_dim;
const float dot = ds4_dot_scalar_ldg(qh, kvrow, head_dim);
s = dot * scale + add;
}
}
row_scores[g] = s;
}
}
/* Head-tiled exact score kernel for head_dim==512.
*
* The reference score kernel assigns one (head, row-chunk) per block and lets
* every thread walk one KV row with a scalar sequential dot. Because MQA
* shares the same KV rows across all 64 heads, that reference layout re-reads
* every KV row once per head, and the per-thread row walk is fully
* uncoalesced (threads stride 2KB apart), which multiplies L2 traffic again.
*
* This kernel keeps the per-score arithmetic bit-identical (same ascending-d
* scalar accumulation `dot += q[d] * kv[d]`, same `dot * scale [+ add]`
* epilogue, same masked-row/raw-window classification) but stages a 16-row KV
* tile and a 16-head Q tile in shared memory with coalesced global loads, so
* each KV row is read from L2 once per 16 heads instead of once per head.
* Scores are independent outputs, so retiling the (head, row) space cannot
* change any output bit as long as each individual dot keeps its order. */
#define DS4_SCORE_TILE_HEADS 16u
#define DS4_SCORE_TILE_ROWS 16u
#define DS4_SCORE_TILE_STRIDE 516u /* 512 + 4 floats: 16B-aligned rows, banks shifted by 4 */
__global__ static void attention_decode_score_split_scores_tile512_kernel(
float *score_out,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return;
extern __shared__ float score_tile_shared[];
float *sh_q = score_tile_shared; /* 16 x 516 */
float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE; /* 16 x 516 */
__shared__ float sh_add[DS4_SCORE_TILE_ROWS];
const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS;
const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS;
if (g_base >= n_score || h_base >= n_head) return;
/* Cooperative Q tile load: 16 heads x 512 floats, float4 coalesced. */
{
const float4 *q4 = (const float4 *)(q + (uint64_t)h_base * 512u);
const uint32_t tile_heads =
n_head - h_base < DS4_SCORE_TILE_HEADS ? n_head - h_base : DS4_SCORE_TILE_HEADS;
for (uint32_t idx = threadIdx.x; idx < tile_heads * 128u; idx += blockDim.x) {
const uint32_t hh = idx >> 7u; /* head within tile */
const uint32_t dd = idx & 127u; /* float4 within row */
const float4 v = q4[hh * 128u + dd];
float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u;
dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w;
}
}
/* Row classification + mask staging (thread per row). */
if (threadIdx.x < DS4_SCORE_TILE_ROWS) {
const uint32_t g = g_base + threadIdx.x;
float add = -INFINITY;
if (g < n_score) {
if (g < raw_count) {
add = 0.0f; /* raw rows are always visible */
} else {
const uint32_t cidx = g - raw_count;
add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f;
}
}
sh_add[threadIdx.x] = add;
}
__syncthreads();
/* Cooperative KV tile load: two rows at a time, float4 coalesced.
* Masked rows (add <= -1e20) are skipped; their scores never read KV. */
{
const uint32_t rows_per_pass = blockDim.x >> 7u; /* 128 threads per row */
const uint32_t rr0 = threadIdx.x >> 7u;
const uint32_t dd = threadIdx.x & 127u;
for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) {
const uint32_t g = g_base + r;
if (g >= n_score) continue;
const bool visible = g < raw_count || sh_add[r] > -1.0e20f;
if (!visible) continue;
const float4 *src;
if (g < raw_count) {
const uint32_t raw_row = (raw_start + raw_first_idx + g) % raw_cap;
src = (const float4 *)(raw_kv + (uint64_t)raw_row * 512u);
} else {
const uint32_t cidx = g - raw_count;
src = (const float4 *)(comp_kv + (uint64_t)cidx * 512u);
}
const float4 v = src[dd];
float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u;
dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w;
}
}
__syncthreads();
/* One score per thread: r = tid&15 (consecutive threads, coalesced score
* writes), h = tid>>4. The dot keeps the reference kernel's exact scalar
* ascending-d accumulation. */
const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u);
const uint32_t h = h_base + (threadIdx.x >> 4u);
const uint32_t g = g_base + r;
if (h >= n_head || g >= n_score) return;
const float scale = rsqrtf((float)head_dim);
float *row_scores = score_out + (uint64_t)h * n_score;
const float *qh = sh_q + (uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE;
const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE;
float s = -INFINITY;
const bool need_dot = g < raw_count || sh_add[r] > -1.0e20f;
if (need_dot) {
/* The reference kernel's runtime-trip loop compiles to one sequential
* FFMA chain. Keep exactly that accumulation order here: batched loads
* for latency hiding, but a single explicit ascending fma chain. */
float dot = 0.0f;
#pragma unroll 1
for (uint32_t d = 0; d < 512u; d += 8u) {
const float a0 = qh[d + 0u], a1 = qh[d + 1u];
const float a2 = qh[d + 2u], a3 = qh[d + 3u];
const float a4 = qh[d + 4u], a5 = qh[d + 5u];
const float a6 = qh[d + 6u], a7 = qh[d + 7u];
const float b0 = kvrow[d + 0u], b1 = kvrow[d + 1u];
const float b2 = kvrow[d + 2u], b3 = kvrow[d + 3u];
const float b4 = kvrow[d + 4u], b5 = kvrow[d + 5u];
const float b6 = kvrow[d + 6u], b7 = kvrow[d + 7u];
dot = __fmaf_rn(a0, b0, dot);
dot = __fmaf_rn(a1, b1, dot);
dot = __fmaf_rn(a2, b2, dot);
dot = __fmaf_rn(a3, b3, dot);
dot = __fmaf_rn(a4, b4, dot);
dot = __fmaf_rn(a5, b5, dot);
dot = __fmaf_rn(a6, b6, dot);
dot = __fmaf_rn(a7, b7, dot);
}
if (g < raw_count) {
s = dot * scale;
} else {
/* The reference expression `dot * scale + add` contracts to one
* FFMA; keep that exact contraction explicit. */
s = __fmaf_rn(dot, scale, sh_add[r]);
}
}
row_scores[g] = s;
}
/* Multi-session form of the exact tiled score kernel. Each z-slice selects a
* private KV table entry, while every individual score keeps the same scalar
* ascending-d FMA chain as the one-session kernel. */
__global__ static void attention_decode_score_split_scores_tile512_rows_kernel(
float *score_out,
const float *q,
cuda_attention_decode_row_table rows,
uint32_t n_rows,
uint32_t score_stride,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t row = blockIdx.z;
if (row >= n_rows) return;
const ds4_gpu_attention_decode_row dsc = rows.row[row];
if (dsc.indexed) return;
const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv;
const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv;
const bool single_all = dsc.ratio == 0u;
const uint32_t qpos = dsc.pos;
const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw;
uint32_t visible_comp = single_all
? dsc.n_comp
: (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u);
if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp;
uint32_t raw_count = 0u;
uint32_t raw_first_idx = 0u;
if (dsc.n_raw != 0u) {
const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u;
if (single_all) {
raw_count = dsc.n_raw > 256u ? 256u : dsc.n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (dsc.window != 0u && qpos + 1u > dsc.window) {
const uint32_t wlo = qpos + 1u - dsc.window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
extern __shared__ float score_tile_shared[];
float *sh_q = score_tile_shared;
float *sh_kv = sh_q + DS4_SCORE_TILE_HEADS * DS4_SCORE_TILE_STRIDE;
const uint32_t g_base = blockIdx.x * DS4_SCORE_TILE_ROWS;
const uint32_t h_base = blockIdx.y * DS4_SCORE_TILE_HEADS;
if (g_base >= n_score || h_base >= n_head) return;
{
const float4 *q4 = (const float4 *)(
q + ((uint64_t)row * n_head + h_base) * head_dim);
const uint32_t tile_heads =
n_head - h_base < DS4_SCORE_TILE_HEADS
? n_head - h_base : DS4_SCORE_TILE_HEADS;
for (uint32_t idx = threadIdx.x;
idx < tile_heads * 128u;
idx += blockDim.x) {
const uint32_t hh = idx >> 7u;
const uint32_t dd = idx & 127u;
const float4 v = q4[hh * 128u + dd];
float *dst = sh_q + hh * DS4_SCORE_TILE_STRIDE + dd * 4u;
dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w;
}
}
__syncthreads();
{
const uint32_t rows_per_pass = blockDim.x >> 7u;
const uint32_t rr0 = threadIdx.x >> 7u;
const uint32_t dd = threadIdx.x & 127u;
for (uint32_t r = rr0; r < DS4_SCORE_TILE_ROWS; r += rows_per_pass) {
const uint32_t g = g_base + r;
if (g >= n_score) continue;
const float4 *src;
if (g < raw_count) {
const uint32_t raw_row =
(dsc.raw_start + raw_first_idx + g) % dsc.raw_cap;
src = (const float4 *)(raw_kv + (uint64_t)raw_row * head_dim);
} else {
src = (const float4 *)(comp_kv +
(uint64_t)(g - raw_count) * head_dim);
}
const float4 v = src[dd];
float *dst = sh_kv + r * DS4_SCORE_TILE_STRIDE + dd * 4u;
dst[0] = v.x; dst[1] = v.y; dst[2] = v.z; dst[3] = v.w;
}
}
__syncthreads();
const uint32_t r = threadIdx.x & (DS4_SCORE_TILE_ROWS - 1u);
const uint32_t h = h_base + (threadIdx.x >> 4u);
const uint32_t g = g_base + r;
if (h >= n_head || g >= n_score) return;
const float scale = rsqrtf((float)head_dim);
float *row_scores = score_out +
((uint64_t)row * n_head + h) * score_stride;
const float *qh = sh_q +
(uint64_t)(threadIdx.x >> 4u) * DS4_SCORE_TILE_STRIDE;
const float *kvrow = sh_kv + (uint64_t)r * DS4_SCORE_TILE_STRIDE;
float dot = 0.0f;
#pragma unroll 1
for (uint32_t dd = 0; dd < 512u; dd += 8u) {
const float a0 = qh[dd + 0u], a1 = qh[dd + 1u];
const float a2 = qh[dd + 2u], a3 = qh[dd + 3u];
const float a4 = qh[dd + 4u], a5 = qh[dd + 5u];
const float a6 = qh[dd + 6u], a7 = qh[dd + 7u];
const float b0 = kvrow[dd + 0u], b1 = kvrow[dd + 1u];
const float b2 = kvrow[dd + 2u], b3 = kvrow[dd + 3u];
const float b4 = kvrow[dd + 4u], b5 = kvrow[dd + 5u];
const float b6 = kvrow[dd + 6u], b7 = kvrow[dd + 7u];
dot = __fmaf_rn(a0, b0, dot);
dot = __fmaf_rn(a1, b1, dot);
dot = __fmaf_rn(a2, b2, dot);
dot = __fmaf_rn(a3, b3, dot);
dot = __fmaf_rn(a4, b4, dot);
dot = __fmaf_rn(a5, b5, dot);
dot = __fmaf_rn(a6, b6, dot);
dot = __fmaf_rn(a7, b7, dot);
}
row_scores[g] = g < raw_count
? dot * scale
: __fmaf_rn(dot, scale, 0.0f);
}
__device__ __forceinline__ float ds4_dot512_float4_ordered(
const float *a,
const float *b) {
const float4 *a4 = (const float4 *)a;
const float4 *b4 = (const float4 *)b;
float dot = 0.0f;
#pragma unroll 1
for (uint32_t i = 0; i < 128u; i++) {
const float4 av = a4[i];
const float4 bv = b4[i];
dot = __fadd_rn(dot, __fmul_rn(av.x, bv.x));
dot = __fadd_rn(dot, __fmul_rn(av.y, bv.y));
dot = __fadd_rn(dot, __fmul_rn(av.z, bv.z));
dot = __fadd_rn(dot, __fmul_rn(av.w, bv.w));
}
return dot;
}
__device__ __forceinline__ float ds4_dot512_float4_plain(
const float *a,
const float *b) {
const float4 *a4 = (const float4 *)a;
const float4 *b4 = (const float4 *)b;
float dot = 0.0f;
#pragma unroll 1
for (uint32_t i = 0; i < 128u; i++) {
const float4 av = a4[i];
const float4 bv = b4[i];
dot += av.x * bv.x;
dot += av.y * bv.y;
dot += av.z * bv.z;
dot += av.w * bv.w;
}
return dot;
}
__global__ static void attention_decode_score_split_scores_vec4_kernel(
float *score_out,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t S) {
const uint32_t h = blockIdx.y;
const uint32_t j = blockIdx.z;
if (h >= n_head || j >= S) return;
const uint32_t head_dim = 512u;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return;
const uint32_t qbase = n_score / S;
const uint32_t rem = n_score % S;
const uint32_t g0 = j * qbase + (j < rem ? j : rem);
const uint32_t cnt = qbase + (j < rem ? 1u : 0u);
const uint32_t g1 = g0 + cnt;
const float *qh = q + (uint64_t)h * head_dim;
float *row_scores = score_out + (uint64_t)h * n_score;
const float scale = rsqrtf((float)head_dim);
for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) {
float s = -INFINITY;
if (g < raw_count) {
const uint32_t raw_row =
(raw_start + raw_first_idx + g) % raw_cap;
const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim;
const float dot = ds4_dot512_float4_ordered(qh, kvrow);
s = dot * scale;
} else {
const uint32_t cidx = g - raw_count;
const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)cidx * head_dim;
const float dot = ds4_dot512_float4_ordered(qh, kvrow);
s = dot * scale + add;
}
}
row_scores[g] = s;
}
}
__global__ static void attention_decode_score_split_scores_vec4_plain_kernel(
float *score_out,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t S) {
const uint32_t h = blockIdx.y;
const uint32_t j = blockIdx.z;
if (h >= n_head || j >= S) return;
const uint32_t head_dim = 512u;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return;
const uint32_t qbase = n_score / S;
const uint32_t rem = n_score % S;
const uint32_t g0 = j * qbase + (j < rem ? j : rem);
const uint32_t cnt = qbase + (j < rem ? 1u : 0u);
const uint32_t g1 = g0 + cnt;
const float *qh = q + (uint64_t)h * head_dim;
float *row_scores = score_out + (uint64_t)h * n_score;
const float scale = rsqrtf((float)head_dim);
for (uint32_t g = g0 + threadIdx.x; g < g1; g += blockDim.x) {
float s = -INFINITY;
if (g < raw_count) {
const uint32_t raw_row =
(raw_start + raw_first_idx + g) % raw_cap;
const float *kvrow = raw_kv + (uint64_t)raw_row * head_dim;
const float dot = ds4_dot512_float4_plain(qh, kvrow);
s = dot * scale;
} else {
const uint32_t cidx = g - raw_count;
const float add = use_comp_mask ? comp_mask[(uint64_t)cidx] : 0.0f;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)cidx * head_dim;
const float dot = ds4_dot512_float4_plain(qh, kvrow);
s = dot * scale + add;
}
}
row_scores[g] = s;
}
}
__global__ static void attention_decode_score_split_finalize_kernel(
float *heads,
const float *sinks,
const float *score_in,
const float *raw_kv,
const float *comp_kv,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t h = blockIdx.y;
if (h >= n_head) return;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
__shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP];
__shared__ uint32_t raw_rows[256];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count_s;
__shared__ uint32_t raw_first_idx_s;
const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x;
const bool score_thread = threadIdx.x < score_threads;
if (threadIdx.x == 0) {
raw_count_s = 0;
raw_first_idx_s = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count_s = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx_s = lo - first_raw_pos;
raw_count_s = hi - lo + 1u;
if (raw_count_s > 256u) raw_count_s = 256u;
}
}
}
}
__syncthreads();
const uint32_t raw_count = raw_count_s;
const uint32_t raw_first_idx = raw_first_idx_s;
if (score_thread) {
for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
}
__syncthreads();
const uint32_t n_score = raw_count + visible_comp;
const float *row_scores = score_in + (uint64_t)h * n_score;
float local_max = sinks[h];
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
const float s = row_scores[i];
scores[i] = s;
local_max = fmaxf(local_max, s);
}
}
if (score_thread) partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial[threadIdx.x] =
fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
}
if (score_thread) partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
float *oh = heads + (uint64_t)h * head_dim;
if (head_dim == 512u && blockDim.x >= 512u) {
const uint32_t d = threadIdx.x;
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
const float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc += kv[d] * s;
}
for (uint32_t c = 0; c < visible_comp; c++) {
const float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc += kv[d] * s;
}
oh[d] = acc / denom;
} else if (head_dim == 512u && blockDim.x == 256u) {
const uint32_t d0 = threadIdx.x;
const uint32_t d1 = d0 + 256u;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
const float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
for (uint32_t c = 0; c < visible_comp; c++) {
const float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
oh[d0] = acc0 / denom;
oh[d1] = acc1 / denom;
} else {
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r];
}
for (uint32_t c = 0; c < visible_comp; c++) {
acc += comp_kv[(uint64_t)c * head_dim + d] * scores[raw_count + c];
}
oh[d] = acc / denom;
}
}
}
__global__ static void attention_decode_score_split_finalize_rows_kernel(
float *heads,
const float *sinks,
const float *score_in,
cuda_attention_decode_row_table rows,
uint32_t n_rows,
uint32_t score_stride,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t row = blockIdx.x;
const uint32_t h = blockIdx.y;
if (row >= n_rows || h >= n_head) return;
const ds4_gpu_attention_decode_row dsc = rows.row[row];
if (dsc.indexed) return;
const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv;
const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv;
const bool single_all = dsc.ratio == 0u;
const uint32_t qpos = dsc.pos;
const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw;
uint32_t visible_comp = single_all
? dsc.n_comp
: (dsc.n_comp ? (qpos + 1u) / dsc.ratio : 0u);
if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp;
__shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP];
__shared__ uint32_t raw_rows[256];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count_s;
__shared__ uint32_t raw_first_idx_s;
const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x;
const bool score_thread = threadIdx.x < score_threads;
if (threadIdx.x == 0u) {
raw_count_s = 0u;
raw_first_idx_s = 0u;
if (dsc.n_raw != 0u) {
const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u;
if (single_all) {
raw_count_s = dsc.n_raw > 256u ? 256u : dsc.n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (dsc.window != 0u && qpos + 1u > dsc.window) {
const uint32_t wlo = qpos + 1u - dsc.window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx_s = lo - first_raw_pos;
raw_count_s = hi - lo + 1u;
if (raw_count_s > 256u) raw_count_s = 256u;
}
}
}
}
__syncthreads();
const uint32_t raw_count = raw_count_s;
const uint32_t raw_first_idx = raw_first_idx_s;
if (score_thread) {
for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) {
raw_rows[r] =
(dsc.raw_start + raw_first_idx + r) % dsc.raw_cap;
}
}
__syncthreads();
const uint32_t n_score = raw_count + visible_comp;
const float *row_scores = score_in +
((uint64_t)row * n_head + h) * score_stride;
float local_max = sinks[h];
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
const float s = row_scores[i];
scores[i] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
}
__syncthreads();
for (uint32_t stride = score_threads >> 1u;
stride > 0u;
stride >>= 1u) {
if (threadIdx.x < stride) {
partial[threadIdx.x] =
fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0u) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
partial[threadIdx.x] = den_local;
}
__syncthreads();
for (uint32_t stride = score_threads >> 1u;
stride > 0u;
stride >>= 1u) {
if (threadIdx.x < stride) {
partial[threadIdx.x] += partial[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0u) {
denom = partial[0] + expf(sinks[h] - max_s);
}
__syncthreads();
float *oh = heads + ((uint64_t)row * n_head + h) * head_dim;
if (head_dim == 512u && blockDim.x >= 512u) {
const uint32_t dim = threadIdx.x;
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc += kv[dim] * scores[r];
}
for (uint32_t c = 0; c < visible_comp; c++) {
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc += kv[dim] * scores[raw_count + c];
}
oh[dim] = acc / denom;
} else {
for (uint32_t dim = threadIdx.x;
dim < head_dim;
dim += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] *
scores[r];
}
for (uint32_t c = 0; c < visible_comp; c++) {
acc += comp_kv[(uint64_t)c * head_dim + dim] *
scores[raw_count + c];
}
oh[dim] = acc / denom;
}
}
}
__global__ static void attention_decode_score_split_finalize_dim2_kernel(
float *heads,
const float *sinks,
const float *score_in,
const float *raw_kv,
const float *comp_kv,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t dim_half = blockIdx.x;
const uint32_t h = blockIdx.y;
if (h >= n_head || head_dim != 512u || dim_half >= 2u) return;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
__shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP];
__shared__ uint32_t raw_rows[256];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count_s;
__shared__ uint32_t raw_first_idx_s;
const uint32_t score_threads = 256u;
if (threadIdx.x == 0) {
raw_count_s = 0;
raw_first_idx_s = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count_s = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx_s = lo - first_raw_pos;
raw_count_s = hi - lo + 1u;
if (raw_count_s > 256u) raw_count_s = 256u;
}
}
}
}
__syncthreads();
const uint32_t raw_count = raw_count_s;
const uint32_t raw_first_idx = raw_first_idx_s;
for (uint32_t r = threadIdx.x; r < raw_count; r += score_threads) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
__syncthreads();
const uint32_t n_score = raw_count + visible_comp;
const float *row_scores = score_in + (uint64_t)h * n_score;
float local_max = sinks[h];
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
const float s = row_scores[i];
scores[i] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial[threadIdx.x] =
fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
const uint32_t d = dim_half * 256u + threadIdx.x;
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
const float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc += kv[d] * s;
}
for (uint32_t c = 0; c < visible_comp; c++) {
const float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)c * head_dim;
acc += kv[d] * s;
}
heads[(uint64_t)h * head_dim + d] = acc / denom;
}
__global__ static void attention_decode_global_softmax_kernel(
float *score_inout,
float *denom_out,
const float *sinks,
uint32_t n_score,
uint32_t n_head) {
const uint32_t h = blockIdx.x;
if (h >= n_head || n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return;
__shared__ float scores[DS4_CUDA_ATTENTION_SCORE_CAP];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom_s;
const uint32_t score_threads = blockDim.x > 256u ? 256u : blockDim.x;
const bool score_thread = threadIdx.x < score_threads;
float *row_scores = score_inout + (uint64_t)h * n_score;
float local_max = sinks[h];
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
const float s = row_scores[i];
scores[i] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
}
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial[threadIdx.x] =
fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
const float e = expf(scores[i] - max_s);
scores[i] = e;
den_local += e;
}
partial[threadIdx.x] = den_local;
}
__syncthreads();
for (uint32_t stride = score_threads >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom_s = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
if (score_thread) {
for (uint32_t i = threadIdx.x; i < n_score; i += score_threads) {
row_scores[i] = scores[i];
}
}
if (threadIdx.x == 0) denom_out[h] = denom_s;
}
__global__ static void attention_decode_split_value_kernel(
float *partials,
const float *score_exp,
const float *raw_kv,
const float *comp_kv,
uint32_t raw_count,
uint32_t raw_first_idx,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_score,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
const uint32_t h = blockIdx.y;
const uint32_t j = blockIdx.z;
if (h >= n_head || j >= S || n_score == 0u) return;
const uint32_t qbase = n_score / S;
const uint32_t rem = n_score % S;
const uint32_t g0 = j * qbase + (j < rem ? j : rem);
const uint32_t cnt = qbase + (j < rem ? 1u : 0u);
const uint32_t g1 = g0 + cnt;
const float *row_scores = score_exp + (uint64_t)h * n_score;
float *pout = partials + ((uint64_t)h * S + j) * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t g = g0; g < g1; g++) {
const float s = row_scores[g];
if (g < raw_count) {
const uint32_t raw_row =
(raw_start + raw_first_idx + g) % raw_cap;
acc += raw_kv[(uint64_t)raw_row * head_dim + d] * s;
} else {
const uint32_t c = g - raw_count;
acc += comp_kv[(uint64_t)c * head_dim + d] * s;
}
}
pout[d] = acc;
}
}
__global__ static void attention_decode_split_value_combine_kernel(
float *heads,
const float *partials,
const float *denom,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
const uint32_t h = blockIdx.y;
if (h >= n_head) return;
const float *base = partials + (uint64_t)h * S * head_dim;
const float den = denom[h];
float *oh = heads + (uint64_t)h * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t j = 0; j < S; j++) {
acc += base[(uint64_t)j * head_dim + d];
}
oh[d] = acc / den;
}
}
typedef struct {
uint32_t n_rot;
uint32_t pos0;
uint32_t n_ctx_orig;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
} cuda_attention_inv_rope_params;
static void attention_decode_score_split_graph_destroy_one(int logical_tier) {
if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return;
cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier];
if (c->exec) (void)cudaGraphExecDestroy(c->exec);
if (c->graph) (void)cudaGraphDestroy(c->graph);
memset(c, 0, sizeof(*c));
}
static int attention_decode_score_split_graph_launch(
int logical_tier,
float *heads,
const float *sinks,
float *scores,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t final_threads,
uint32_t S,
const cuda_attention_inv_rope_params *inv_rope) {
if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0;
cuda_score_split_graph_cache *c = &g_score_split_graph[logical_tier];
const bool graph_inv_rope =
inv_rope &&
head_dim == 512u &&
inv_rope->n_rot != 0u &&
inv_rope->n_rot <= head_dim &&
(inv_rope->n_rot & 1u) == 0u;
const bool shape_match =
c->valid &&
c->n_head == n_head &&
c->head_dim == head_dim &&
c->S == S &&
c->final_threads == final_threads &&
c->fuses_inv_rope == (graph_inv_rope ? 1 : 0) &&
(!graph_inv_rope || c->n_rot == inv_rope->n_rot);
if (c->valid && !shape_match) {
attention_decode_score_split_graph_destroy_one(logical_tier);
c = &g_score_split_graph[logical_tier];
}
dim3 score_grid(1, n_head, S);
dim3 final_grid(1, n_head, 1);
dim3 score_block(256, 1, 1);
dim3 final_block(final_threads, 1, 1);
void *score_args[] = {
&scores, &q, &raw_kv, &comp_kv, &comp_mask, &use_comp_mask,
&pos0, &n_raw, &raw_cap, &raw_start, &n_comp, &window, &ratio,
&n_head, &head_dim, &S
};
cudaKernelNodeParams score_params;
memset(&score_params, 0, sizeof(score_params));
score_params.func = (void *)attention_decode_score_split_scores_kernel;
score_params.gridDim = score_grid;
score_params.blockDim = score_block;
score_params.sharedMemBytes = 0;
score_params.kernelParams = score_args;
score_params.extra = NULL;
void *final_args[] = {
&heads, &sinks, &scores, &raw_kv, &comp_kv, &pos0, &n_raw,
&raw_cap, &raw_start, &n_comp, &window, &ratio, &n_head, &head_dim
};
cudaKernelNodeParams final_params;
memset(&final_params, 0, sizeof(final_params));
final_params.func = (void *)attention_decode_score_split_finalize_kernel;
final_params.gridDim = final_grid;
final_params.blockDim = final_block;
final_params.sharedMemBytes = 0;
final_params.kernelParams = final_args;
final_params.extra = NULL;
uint32_t rope_n_tok = 1u;
uint32_t rope_pos_stride = 1u;
int rope_inverse = 1;
uint32_t rope_n_rot = graph_inv_rope ? inv_rope->n_rot : 0u;
uint32_t rope_pos0 = graph_inv_rope ? inv_rope->pos0 : 0u;
uint32_t rope_n_ctx_orig = graph_inv_rope ? inv_rope->n_ctx_orig : 0u;
float rope_freq_base = graph_inv_rope ? inv_rope->freq_base : 0.0f;
float rope_freq_scale = graph_inv_rope ? inv_rope->freq_scale : 0.0f;
float rope_ext_factor = graph_inv_rope ? inv_rope->ext_factor : 0.0f;
float rope_attn_factor = graph_inv_rope ? inv_rope->attn_factor : 0.0f;
float rope_beta_fast = graph_inv_rope ? inv_rope->beta_fast : 0.0f;
float rope_beta_slow = graph_inv_rope ? inv_rope->beta_slow : 0.0f;
void *rope_args[] = {
&heads, &rope_n_tok, &n_head, &head_dim, &rope_n_rot,
&rope_pos0, &rope_pos_stride, &rope_n_ctx_orig, &rope_inverse,
&rope_freq_base, &rope_freq_scale, &rope_ext_factor,
&rope_attn_factor, &rope_beta_fast, &rope_beta_slow
};
cudaKernelNodeParams rope_params;
memset(&rope_params, 0, sizeof(rope_params));
if (graph_inv_rope) {
const uint32_t pairs = n_head * (rope_n_rot / 2u);
rope_params.func = (void *)rope_tail_kernel;
rope_params.gridDim = dim3((pairs + 255u) / 256u, 1, 1);
rope_params.blockDim = dim3(256, 1, 1);
rope_params.sharedMemBytes = 0;
rope_params.kernelParams = rope_args;
rope_params.extra = NULL;
}
if (!c->valid) {
cudaError_t err = cudaGraphCreate(&c->graph, 0);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: attention score-split graph create failed: %s\n",
cudaGetErrorString(err));
attention_decode_score_split_graph_destroy_one(logical_tier);
return -1;
}
err = cudaGraphAddKernelNode(&c->score_node, c->graph, NULL, 0,
&score_params);
if (err == cudaSuccess) {
err = cudaGraphAddKernelNode(&c->final_node, c->graph,
&c->score_node, 1, &final_params);
}
if (err == cudaSuccess && graph_inv_rope) {
err = cudaGraphAddKernelNode(&c->rope_node, c->graph,
&c->final_node, 1, &rope_params);
}
if (err == cudaSuccess) {
err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0);
}
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: attention score-split graph instantiate failed: %s\n",
cudaGetErrorString(err));
attention_decode_score_split_graph_destroy_one(logical_tier);
return -1;
}
c->n_head = n_head;
c->head_dim = head_dim;
c->S = S;
c->final_threads = final_threads;
c->n_rot = graph_inv_rope ? inv_rope->n_rot : 0u;
c->fuses_inv_rope = graph_inv_rope ? 1 : 0;
c->valid = 1;
} else {
cudaError_t err =
cudaGraphExecKernelNodeSetParams(c->exec, c->score_node,
&score_params);
if (err == cudaSuccess) {
err = cudaGraphExecKernelNodeSetParams(c->exec, c->final_node,
&final_params);
}
if (err == cudaSuccess && graph_inv_rope) {
err = cudaGraphExecKernelNodeSetParams(c->exec, c->rope_node,
&rope_params);
}
if (err != cudaSuccess) {
fprintf(stderr,
"ds4: attention score-split graph update failed: %s\n",
cudaGetErrorString(err));
attention_decode_score_split_graph_destroy_one(logical_tier);
return -1;
}
}
cudaError_t err = cudaGraphLaunch(c->exec, 0);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: attention score-split graph launch failed: %s\n",
cudaGetErrorString(err));
attention_decode_score_split_graph_destroy_one(logical_tier);
return -1;
}
return 1;
}
static int attention_decode_score_split_launch(
int logical_tier,
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t final_threads,
const cuda_attention_inv_rope_params *inv_rope) {
if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0)) return 0;
const int explicit_exact =
cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 0);
if (!cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1)) return 0;
if (!explicit_exact && cuda_splitkv_decode_requested()) return 0;
if (g_cuda_decode_score4 || g_cuda_decode_score8) return 0;
if (head_dim == 0u || n_head == 0u) return 0;
const bool single_all = (ratio == 0u);
const uint32_t qpos = pos0;
const uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
const uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) return 0;
/* With the head-tiled score kernel the exact score-split path beats the
* one-block mixed kernel even for short score counts, so the gate that
* used to protect short contexts (512) now defaults to 1. */
const uint32_t min_score = cuda_parse_u32_env_clamped(
"DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE", 1u, 0u,
DS4_CUDA_ATTENTION_SCORE_CAP, NULL);
if (n_score < min_score) return 0;
uint32_t chunk = cuda_parse_u32_env_clamped(
"DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK", DS4_CUDA_SPLITKV_CHUNK,
1u, DS4_CUDA_ATTENTION_SCORE_CAP, NULL);
uint32_t s_floor = cuda_parse_u32_env_clamped(
"DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR", 6u,
1u, DS4_CUDA_SPLITKV_S_MAX, NULL);
uint32_t s_max = cuda_parse_u32_env_clamped(
"DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX", DS4_CUDA_SPLITKV_S_MAX,
1u, DS4_CUDA_SPLITKV_S_MAX, NULL);
int exact_present = 0;
uint32_t S = cuda_parse_u32_env_clamped(
"DS4_CUDA_EXACT_SCORE_SPLIT_S", 0u, 1u,
DS4_CUDA_SPLITKV_S_MAX, &exact_present);
if (!exact_present) {
S = (n_score + chunk - 1u) / chunk;
if (S < s_floor) S = s_floor < n_score ? s_floor : n_score;
if (S > s_max) S = s_max;
}
if (S > n_score) S = n_score;
if (S <= 1u) return 0;
const bool graph_inv_rope =
g_cuda_exact_score_split_fuse_inv_rope &&
inv_rope &&
head_dim == 512u &&
final_threads >= 512u &&
inv_rope->n_rot != 0u &&
inv_rope->n_rot <= 512u &&
(inv_rope->n_rot & 1u) == 0u;
const uint64_t score_count = (uint64_t)n_head * n_score;
float *scores = (float *)cuda_tmp_alloc_on(logical_tier,
score_count * sizeof(float),
"attention exact score split");
if (!scores) return 0;
const bool use_ldg_scores = g_cuda_exact_score_split_ldg;
const bool use_vec4_plain_scores =
!use_ldg_scores &&
g_cuda_exact_score_split_vec4_plain &&
head_dim == 512u;
const bool use_vec4_scores =
!use_ldg_scores &&
!use_vec4_plain_scores &&
(g_cuda_exact_score_split_vec4 || g_decode_score_vec4) &&
head_dim == 512u;
const bool use_dim2_finalize =
g_cuda_exact_score_split_dim2 &&
head_dim == 512u &&
final_threads >= 512u &&
!graph_inv_rope;
if ((g_cuda_exact_score_split_graph || graph_inv_rope) &&
!use_dim2_finalize &&
!use_ldg_scores &&
!use_vec4_plain_scores &&
!use_vec4_scores)
{
int rc = attention_decode_score_split_graph_launch(
logical_tier, heads, sinks, scores, q, raw_kv, comp_kv, comp_mask,
use_comp_mask, pos0, n_raw, raw_cap, raw_start, n_comp, window,
ratio, n_head, head_dim, final_threads, S,
graph_inv_rope ? inv_rope : NULL);
if (rc == 1) return 1;
if (rc < 0) return -1;
}
if (graph_inv_rope) return 0;
static int score_tile_disabled = -1;
if (score_tile_disabled < 0) {
score_tile_disabled = getenv("DS4_CUDA_NO_SCORE_TILE") != NULL ? 1 : 0;
}
if (!score_tile_disabled &&
head_dim == 512u &&
!use_ldg_scores &&
!use_vec4_plain_scores &&
!use_vec4_scores) {
/* cudaFuncSetAttribute() applies to the current device only, so opt in
* to >48KB dynamic shared memory once per device. */
static int tile_shmem_ready[DS4_MAX_GPUS] = {0};
const size_t tile_shmem =
(size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) *
DS4_SCORE_TILE_STRIDE * sizeof(float);
int tile_dev = 0;
cudaGetDevice(&tile_dev);
if (tile_dev >= 0 && tile_dev < DS4_MAX_GPUS &&
!tile_shmem_ready[tile_dev]) {
if (!cuda_ok(cudaFuncSetAttribute(
attention_decode_score_split_scores_tile512_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)tile_shmem),
"attention score tile shared-memory opt-in")) {
score_tile_disabled = 1;
}
tile_shmem_ready[tile_dev] = 1;
}
if (score_tile_disabled) {
return 0; /* retry via the generic path on the next call */
}
dim3 tile_grid((n_score + DS4_SCORE_TILE_ROWS - 1u) / DS4_SCORE_TILE_ROWS,
(n_head + DS4_SCORE_TILE_HEADS - 1u) / DS4_SCORE_TILE_HEADS,
1);
attention_decode_score_split_scores_tile512_kernel<<<tile_grid, 256, tile_shmem>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim);
if (!cuda_ok(cudaGetLastError(), "attention exact score split tile launch")) return -1;
} else {
dim3 score_grid(1, n_head, S);
if (use_ldg_scores) {
attention_decode_score_split_scores_ldg_kernel<<<score_grid, 256>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim, S);
} else if (use_vec4_plain_scores) {
attention_decode_score_split_scores_vec4_plain_kernel<<<score_grid, 256>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, S);
} else if (use_vec4_scores) {
attention_decode_score_split_scores_vec4_kernel<<<score_grid, 256>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, S);
} else {
attention_decode_score_split_scores_kernel<<<score_grid, 256>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim, S);
}
if (!cuda_ok(cudaGetLastError(), "attention exact score split scores launch")) return -1;
}
if (use_dim2_finalize) {
dim3 final_grid(2, n_head, 1);
attention_decode_score_split_finalize_dim2_kernel<<<final_grid, 256>>>(
heads, sinks, scores, raw_kv, comp_kv,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim);
if (!cuda_ok(cudaGetLastError(), "attention exact score split dim2 finalize launch")) return -1;
} else {
dim3 final_grid(1, n_head, 1);
attention_decode_score_split_finalize_kernel<<<final_grid, final_threads>>>(
heads, sinks, scores, raw_kv, comp_kv,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim);
if (!cuda_ok(cudaGetLastError(), "attention exact score split finalize launch")) return -1;
}
return 1;
}
/* ---- perf-02 split-KV / flash-decode (opt-in, default OFF) ----------------
*
* attention_decode_splitkv_kernel computes a partial online-softmax over a
* contiguous chunk of the flattened logical row set [0, n_score) used by
* attention_decode_mixed_kernel (raw rows first, then compressed rows, same
* ascending ordering). Each block handles (t = blockIdx.x, h = blockIdx.y,
* chunk = blockIdx.z) and writes a partial (m_j, l_j, acc_j[head_dim]) WITHOUT
* the sink term. attention_decode_splitkv_combine_kernel merges the S partials
* per (t,h), folds the sink once, and writes the final normalized head output.
*
* The math is the standard flash-attention online-softmax rescale and is
* algebraically identical to attention_decode_mixed_kernel; it is NOT
* guaranteed bit-identical in FP32 (different expf inputs + add/mul grouping),
* hence default-OFF behind DS4_CUDA_SPLITKV_DECODE and the S==1 dispatch to the
* old kernel as the bit-exact anchor (handled in the launch helper).
*
* Partials scratch layout (per logical tier), contiguous floats:
* stride = head_dim + 2
* base(t,h,j) = ((t*n_head + h)*S + j) * stride
* [0] = m_j (chunk running max; -INF if empty/all-masked)
* [1] = l_j (chunk denominator sum exp(s - m_j))
* [2 .. 2+head_dim) = acc_j[head_dim] (chunk weighted value sum)
*/
__global__ static void attention_decode_splitkv_kernel(
float *partials,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
uint32_t j = blockIdx.z;
if (t >= n_tokens || h >= n_head || j >= S) return;
const bool single_all = (n_tokens == 1u && ratio == 0u);
uint32_t qpos = pos0 + t;
uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
/* scores buffer holds only this chunk's rows. The launch helper guarantees
* cnt <= DS4_CUDA_SPLITKV_SCORE_CAP, including env-tuned split counts. */
__shared__ float scores[DS4_CUDA_SPLITKV_SCORE_CAP];
__shared__ uint32_t raw_rows[256];
__shared__ float partial[256];
__shared__ float m_s;
__shared__ float l_s;
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
float scale = rsqrtf((float)head_dim);
if (threadIdx.x == 0) {
raw_count = 0;
raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
uint32_t n_score = raw_count + visible_comp;
/* even split of [0, n_score) across S chunks: first (n_score % S) chunks
* get base+1, identical deterministic partition for every block. */
uint32_t qbase = n_score / S;
uint32_t rem = n_score % S;
uint32_t g0 = j * qbase + (j < rem ? j : rem);
uint32_t cnt = qbase + (j < rem ? 1u : 0u);
uint32_t g1 = g0 + cnt; /* exclusive end of this chunk */
/* Map raw rows that fall in this chunk into shared raw_rows[]. The chunk's
* raw portion is [raw_lo, raw_hi). cnt <= CHUNK and raw rows <= 256, so
* the slice fits raw_rows[256]. */
uint32_t raw_lo = g0 < raw_count ? g0 : raw_count;
uint32_t raw_hi = g1 < raw_count ? g1 : raw_count;
for (uint32_t r = raw_lo + threadIdx.x; r < raw_hi; r += blockDim.x) {
raw_rows[r - raw_lo] = (raw_start + raw_first_idx + r) % raw_cap;
}
__syncthreads();
float *pout = partials + (((uint64_t)t * n_head + h) * S + j) * (head_dim + 2u);
/* Pass 1: scores for this chunk's rows into shared scores[0..cnt). */
float local_max = -INFINITY;
for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) {
uint32_t g = g0 + i;
float s;
if (g < raw_count) {
const float *kvrow = raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale;
} else {
uint32_t c = g - raw_count;
float add = use_comp_mask ? comp_mask[(uint64_t)t * n_comp + c] : 0.0f;
s = -INFINITY;
if (add > -1.0e20f) {
const float *kvrow = comp_kv + (uint64_t)c * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
s = dot * scale + add;
}
}
scores[i] = s;
local_max = fmaxf(local_max, s);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) m_s = partial[0];
__syncthreads();
float chunk_max = m_s;
/* All-masked / empty-chunk guard: never evaluate exp(-INF - -INF) -> NaN.
* Write zero partial (m=-INF, l=0, acc=0) and return. */
if (!isfinite(chunk_max)) {
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) pout[2u + d] = 0.0f;
if (threadIdx.x == 0) { pout[0] = -INFINITY; pout[1] = 0.0f; }
return;
}
/* Pass 2: exponentiate in place and reduce denominator. */
float den_local = 0.0f;
for (uint32_t i = threadIdx.x; i < cnt; i += blockDim.x) {
float e = expf(scores[i] - chunk_max);
scores[i] = e;
den_local += e;
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) l_s = partial[0];
__syncthreads();
/* Pass 3: weighted value accumulation over this chunk's rows (ascending g),
* preserving raw-then-comp ordering to match the reference accumulation. */
if (head_dim == 512u && blockDim.x == 256u) {
uint32_t d0 = threadIdx.x;
uint32_t d1 = d0 + 256u;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t i = 0; i < cnt; i++) {
uint32_t g = g0 + i;
float s = scores[i];
const float *kv = (g < raw_count)
? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim
: comp_kv + (uint64_t)(g - raw_count) * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
pout[2u + d0] = acc0;
pout[2u + d1] = acc1;
} else {
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t i = 0; i < cnt; i++) {
uint32_t g = g0 + i;
float s = scores[i];
const float *kv = (g < raw_count)
? raw_kv + (uint64_t)raw_rows[g - raw_lo] * head_dim
: comp_kv + (uint64_t)(g - raw_count) * head_dim;
acc += kv[d] * s;
}
pout[2u + d] = acc;
}
}
if (threadIdx.x == 0) {
pout[0] = chunk_max;
pout[1] = l_s;
}
}
__global__ static void attention_decode_splitkv_combine_kernel(
float *heads,
const float *sinks,
const float *partials,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t S) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || h >= n_head) return;
const float *base = partials + (((uint64_t)t * n_head + h) * S) * (head_dim + 2u);
uint32_t stride = head_dim + 2u;
__shared__ float M_s;
__shared__ float L_s;
if (threadIdx.x == 0) {
/* Global max M = max(sink, max_j m_j); sink placed first to match the
* reference (sink seeds local_max). */
float M = sinks[h];
for (uint32_t jj = 0; jj < S; jj++) {
float m_j = base[(uint64_t)jj * stride];
M = fmaxf(M, m_j); /* -INF partials never raise M */
}
M_s = M;
/* L = Σ_j exp(m_j - M) * l_j + exp(sink - M); sink term added last to
* mirror the reference's denom = Σ scores + expf(sink - max). Chunks
* with l_j == 0 / m_j == -INF contribute exactly 0 (guarded to avoid
* exp(-INF - finite) * 0 edge cases). */
float L = 0.0f;
for (uint32_t jj = 0; jj < S; jj++) {
float m_j = base[(uint64_t)jj * stride];
float l_j = base[(uint64_t)jj * stride + 1u];
if (l_j != 0.0f && isfinite(m_j)) L += expf(m_j - M) * l_j;
}
L += expf(sinks[h] - M);
L_s = L;
}
__syncthreads();
float M = M_s;
float L = L_s;
float *oh = heads + ((uint64_t)t * n_head + h) * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float A = 0.0f;
for (uint32_t jj = 0; jj < S; jj++) {
float m_j = base[(uint64_t)jj * stride];
float l_j = base[(uint64_t)jj * stride + 1u];
if (l_j != 0.0f && isfinite(m_j)) {
A += expf(m_j - M) * base[(uint64_t)jj * stride + 2u + d];
}
}
oh[d] = A / L;
}
}
__device__ __forceinline__ void attention_compact_topk_stable(
uint32_t *comp_rows,
uint32_t *comp_count,
uint32_t *warp_offsets,
const int32_t *topk,
uint32_t top_k,
uint32_t visible_comp) {
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t n_warp = blockDim.x >> 5u;
if (threadIdx.x == 0u) *comp_count = 0u;
__syncthreads();
for (uint32_t base = 0u; base < 512u; base += blockDim.x) {
const uint32_t i = base + threadIdx.x;
const int32_t c = i < top_k ? topk[i] : -1;
const bool valid = c >= 0 && (uint32_t)c < visible_comp;
const uint32_t mask = __ballot_sync(0xffffffffu, valid);
if (lane == 0u) warp_offsets[warp] = __popc(mask);
__syncthreads();
if (threadIdx.x == 0u) {
uint32_t out = *comp_count;
for (uint32_t w = 0u; w < n_warp; w++) {
const uint32_t count = warp_offsets[w];
warp_offsets[w] = out;
out += count;
}
*comp_count = out;
}
__syncthreads();
if (valid) {
const uint32_t lanes_before = lane == 0u
? 0u : ((1u << lane) - 1u);
const uint32_t slot = warp_offsets[warp] +
__popc(mask & lanes_before);
if (slot < 512u) comp_rows[slot] = (uint32_t)c;
}
__syncthreads();
}
}
__global__ static void attention_indexed_mixed_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const int32_t *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t h = blockIdx.y;
if (t >= n_tokens || h >= n_head) return;
uint32_t qpos = pos0 + t;
uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t visible_comp = n_comp;
if (ratio != 0) {
visible_comp = (qpos + 1u) / ratio;
if (visible_comp > n_comp) visible_comp = n_comp;
}
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
__shared__ float scores[768];
__shared__ uint32_t raw_rows[256];
__shared__ uint32_t comp_rows[512];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
__shared__ uint32_t comp_count;
__shared__ uint32_t comp_warp_offsets[8];
float scale = rsqrtf((float)head_dim);
if (threadIdx.x == 0) {
raw_count = 0;
raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
attention_compact_topk_stable(
comp_rows, &comp_count, comp_warp_offsets,
topk + (uint64_t)t * top_k, top_k, visible_comp);
uint32_t n_score = raw_count + comp_count;
float local_max = sinks[h];
if (comp_count == 0) {
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kvrow[d];
scores[r] = dot * scale;
local_max = fmaxf(local_max, scores[r]);
}
} else {
uint32_t qlane = threadIdx.x & 7u;
uint32_t qgroup = threadIdx.x >> 3u;
for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) {
uint32_t row = row0 + qgroup;
if (row < n_score) {
const float *kvrow = row < raw_count
? raw_kv + (uint64_t)raw_rows[row] * head_dim
: comp_kv + (uint64_t)comp_rows[row - raw_count] * head_dim;
float dot = 0.0f;
for (uint32_t d = qlane; d < head_dim; d += 8u) dot += qh[d] * kvrow[d];
const uint32_t mask = 0xffu << (threadIdx.x & 24u);
for (uint32_t off = 4u; off > 0u; off >>= 1u) {
dot += __shfl_down_sync(mask, dot, off, 8);
}
if (qlane == 0) scores[row] = dot * scale;
}
}
__syncthreads();
for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) {
local_max = fmaxf(local_max, scores[i]);
}
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
float *oh = heads + ((uint64_t)t * n_head + h) * head_dim;
if (head_dim == 512u && blockDim.x == 256u) {
uint32_t d0 = threadIdx.x;
uint32_t d1 = d0 + 256u;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
for (uint32_t c = 0; c < comp_count; c++) {
float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
oh[d0] = acc0 / denom;
oh[d1] = acc1 / denom;
} else {
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + d] * scores[r];
for (uint32_t s = 0; s < comp_count; s++) acc += comp_kv[(uint64_t)comp_rows[s] * head_dim + d] * scores[raw_count + s];
oh[d] = acc / denom;
}
}
}
__global__ static void attention_indexed_mixed_decode_rows_kernel(
float *heads,
const float *sinks,
const float *q,
cuda_attention_decode_row_table rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t row = blockIdx.x;
const uint32_t h = blockIdx.y;
if (row >= n_rows || h >= n_head) return;
const ds4_gpu_attention_decode_row dsc = rows.row[row];
if (!dsc.indexed) return;
const float *raw_kv = (const float *)(uintptr_t)dsc.raw_kv;
const float *comp_kv = (const float *)(uintptr_t)dsc.comp_kv;
const int32_t *topk = (const int32_t *)(uintptr_t)dsc.topk;
const uint32_t qpos = dsc.pos;
const uint32_t first_raw_pos = dsc.pos + 1u - dsc.n_raw;
uint32_t visible_comp = dsc.n_comp;
if (dsc.ratio != 0u) {
visible_comp = (qpos + 1u) / dsc.ratio;
if (visible_comp > dsc.n_comp) visible_comp = dsc.n_comp;
}
const float *qh = q + ((uint64_t)row * n_head + h) * head_dim;
__shared__ float scores[768];
__shared__ uint32_t raw_rows[256];
__shared__ uint32_t comp_rows[512];
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
__shared__ uint32_t comp_count;
__shared__ uint32_t comp_warp_offsets[8];
const float scale = rsqrtf((float)head_dim);
if (threadIdx.x == 0u) {
raw_count = 0u;
raw_first_idx = 0u;
if (dsc.n_raw != 0u) {
const uint32_t raw_last_pos = first_raw_pos + dsc.n_raw - 1u;
if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (dsc.window != 0u && qpos + 1u > dsc.window) {
const uint32_t wlo = qpos + 1u - dsc.window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
raw_rows[r] =
(dsc.raw_start + raw_first_idx + r) % dsc.raw_cap;
}
attention_compact_topk_stable(
comp_rows, &comp_count, comp_warp_offsets,
topk, dsc.top_k, visible_comp);
const uint32_t n_score = raw_count + comp_count;
float local_max = sinks[h];
if (comp_count == 0u) {
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
const float *kvrow = raw_kv + (uint64_t)raw_rows[r] * head_dim;
float dot = 0.0f;
for (uint32_t dim = 0; dim < head_dim; dim++) {
dot += qh[dim] * kvrow[dim];
}
scores[r] = dot * scale;
local_max = fmaxf(local_max, scores[r]);
}
} else {
const uint32_t qlane = threadIdx.x & 7u;
const uint32_t qgroup = threadIdx.x >> 3u;
for (uint32_t row0 = 0; row0 < n_score; row0 += 32u) {
const uint32_t score_row = row0 + qgroup;
if (score_row < n_score) {
const float *kvrow = score_row < raw_count
? raw_kv + (uint64_t)raw_rows[score_row] * head_dim
: comp_kv +
(uint64_t)comp_rows[score_row - raw_count] * head_dim;
float dot = 0.0f;
for (uint32_t dim = qlane; dim < head_dim; dim += 8u) {
dot += qh[dim] * kvrow[dim];
}
const uint32_t mask = 0xffu << (threadIdx.x & 24u);
for (uint32_t off = 4u; off > 0u; off >>= 1u) {
dot += __shfl_down_sync(mask, dot, off, 8);
}
if (qlane == 0u) scores[score_row] = dot * scale;
}
}
__syncthreads();
for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) {
local_max = fmaxf(local_max, scores[i]);
}
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u;
stride > 0u;
stride >>= 1u) {
if (threadIdx.x < stride) {
partial[threadIdx.x] =
fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0u) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t i = threadIdx.x; i < n_score; i += blockDim.x) {
scores[i] = expf(scores[i] - max_s);
den_local += scores[i];
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u;
stride > 0u;
stride >>= 1u) {
if (threadIdx.x < stride) {
partial[threadIdx.x] += partial[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0u) {
denom = partial[0] + expf(sinks[h] - max_s);
}
__syncthreads();
float *oh = heads + ((uint64_t)row * n_head + h) * head_dim;
if (head_dim == 512u && blockDim.x == 256u) {
const uint32_t d0 = threadIdx.x;
const uint32_t d1 = d0 + 256u;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
const float s = scores[r];
const float *kv = raw_kv + (uint64_t)raw_rows[r] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
for (uint32_t c = 0; c < comp_count; c++) {
const float s = scores[raw_count + c];
const float *kv = comp_kv + (uint64_t)comp_rows[c] * head_dim;
acc0 += kv[d0] * s;
acc1 += kv[d1] * s;
}
oh[d0] = acc0 / denom;
oh[d1] = acc1 / denom;
} else {
for (uint32_t dim = threadIdx.x;
dim < head_dim;
dim += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < raw_count; r++) {
acc += raw_kv[(uint64_t)raw_rows[r] * head_dim + dim] *
scores[r];
}
for (uint32_t c = 0; c < comp_count; c++) {
acc += comp_kv[(uint64_t)comp_rows[c] * head_dim + dim] *
scores[raw_count + c];
}
oh[dim] = acc / denom;
}
}
}
__global__ static void attention_indexed_mixed_heads8_rb4_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const int32_t *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t head_group = blockIdx.y;
if (t >= n_tokens || head_dim != 512u) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t head = head_group * 8u + warp;
const bool valid_head = head < n_head;
__shared__ uint32_t raw_rows[256];
__shared__ uint32_t comp_rows[512];
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
__shared__ uint32_t comp_count;
__shared__ float4 kv_shared[4 * 128];
__shared__ float scores[8 * 768];
uint32_t qpos = pos0 + t;
uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t visible_comp = n_comp;
if (ratio != 0) {
visible_comp = (qpos + 1u) / ratio;
if (visible_comp > n_comp) visible_comp = n_comp;
}
if (threadIdx.x == 0) {
raw_count = 0;
raw_first_idx = 0;
comp_count = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
if (threadIdx.x == 0) {
for (uint32_t i = 0; i < top_k && comp_count < 512u; i++) {
int32_t c = topk[(uint64_t)t * top_k + i];
if (c >= 0 && (uint32_t)c < visible_comp) comp_rows[comp_count++] = (uint32_t)c;
}
}
__syncthreads();
const uint32_t n_score = raw_count + comp_count;
const float scale = rsqrtf((float)head_dim);
const float4 *q4 = valid_head
? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim)
: NULL;
float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 q1 = q0, q2 = q0, q3 = q0;
if (valid_head) {
q0 = q4[lane + 0u];
q1 = q4[lane + 32u];
q2 = q4[lane + 64u];
q3 = q4[lane + 96u];
}
for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) {
const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u;
for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) {
const uint32_t rr = off >> 7u;
const uint32_t c4 = off & 127u;
const uint32_t sr = row0 + rr;
const float4 *src = sr < raw_count
? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim)
: (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim);
kv_shared[off] = src[c4];
}
__syncthreads();
if (valid_head) {
for (uint32_t rr = 0; rr < nr; rr++) {
const float4 *kv4 = kv_shared + rr * 128u;
float dot = dot4_f32(q0, kv4[lane + 0u]) +
dot4_f32(q1, kv4[lane + 32u]) +
dot4_f32(q2, kv4[lane + 64u]) +
dot4_f32(q3, kv4[lane + 96u]);
dot = warp_sum_f32(dot);
if (lane == 0) scores[warp * 768u + row0 + rr] = dot * scale;
}
}
__syncthreads();
}
float max_s = valid_head ? sinks[head] : -INFINITY;
if (valid_head) {
const float *score_row = scores + warp * 768u;
for (uint32_t i = lane; i < n_score; i += 32u) max_s = fmaxf(max_s, score_row[i]);
max_s = warp_max_f32(max_s);
max_s = __shfl_sync(0xffffffffu, max_s, 0);
}
float den = 0.0f;
if (valid_head) {
float *score_row = scores + warp * 768u;
for (uint32_t i = lane; i < n_score; i += 32u) {
float p = expf(score_row[i] - max_s);
score_row[i] = p;
den += p;
}
den = warp_sum_f32(den);
den += expf(sinks[head] - max_s);
den = __shfl_sync(0xffffffffu, den, 0);
}
float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 o1 = o0, o2 = o0, o3 = o0;
for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) {
const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u;
for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) {
const uint32_t rr = off >> 7u;
const uint32_t c4 = off & 127u;
const uint32_t sr = row0 + rr;
const float4 *src = sr < raw_count
? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim)
: (const float4 *)(comp_kv + (uint64_t)comp_rows[sr - raw_count] * head_dim);
kv_shared[off] = src[c4];
}
__syncthreads();
if (valid_head) {
const float *score_row = scores + warp * 768u;
for (uint32_t rr = 0; rr < nr; rr++) {
const float p = den == 0.0f ? 0.0f : score_row[row0 + rr] / den;
const float4 *kv4 = kv_shared + rr * 128u;
float4 k0 = kv4[lane + 0u];
float4 k1 = kv4[lane + 32u];
float4 k2 = kv4[lane + 64u];
float4 k3 = kv4[lane + 96u];
o0.x += k0.x * p; o0.y += k0.y * p; o0.z += k0.z * p; o0.w += k0.w * p;
o1.x += k1.x * p; o1.y += k1.y * p; o1.z += k1.z * p; o1.w += k1.w * p;
o2.x += k2.x * p; o2.y += k2.y * p; o2.z += k2.z * p; o2.w += k2.w * p;
o3.x += k3.x * p; o3.y += k3.y * p; o3.z += k3.z * p; o3.w += k3.w * p;
}
}
__syncthreads();
}
if (valid_head) {
float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim);
out4[lane + 0u] = o0;
out4[lane + 32u] = o1;
out4[lane + 64u] = o2;
out4[lane + 96u] = o3;
}
}
template <uint32_t ROWS_PER_STAGE, uint32_t HEADS_PER_GROUP>
__global__ static void __launch_bounds__(512, 2)
attention_indexed_mixed_heads8_online_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const int32_t *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t head_group = blockIdx.y;
if (t >= n_tokens || head_dim != 512u) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t head = head_group * HEADS_PER_GROUP + warp;
const bool valid_head = head < n_head;
__shared__ uint32_t raw_rows[256];
__shared__ uint32_t raw_count;
__shared__ uint32_t raw_first_idx;
__shared__ float4 kv_shared[ROWS_PER_STAGE * 128];
uint32_t qpos = pos0 + t;
uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t visible_comp = n_comp;
if (ratio != 0) {
visible_comp = (qpos + 1u) / ratio;
if (visible_comp > n_comp) visible_comp = n_comp;
}
if (threadIdx.x == 0) {
raw_count = 0;
raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
}
__syncthreads();
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
__syncthreads();
const uint32_t comp_count = top_k < 512u ? top_k : 512u;
const uint32_t n_score = raw_count + comp_count;
const float scale = rsqrtf((float)head_dim);
const float4 *q4 = valid_head
? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim)
: NULL;
float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 q1 = q0, q2 = q0, q3 = q0;
if (valid_head) {
q0 = q4[lane + 0u];
q1 = q4[lane + 32u];
q2 = q4[lane + 64u];
q3 = q4[lane + 96u];
}
float max_s = -INFINITY;
float sum_s = 0.0f;
float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 o1 = o0, o2 = o0, o3 = o0;
for (uint32_t row0 = 0; row0 < n_score; row0 += ROWS_PER_STAGE) {
const uint32_t nr = n_score - row0 < ROWS_PER_STAGE ? n_score - row0 : ROWS_PER_STAGE;
for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) {
const uint32_t rr = off >> 7u;
const uint32_t c4 = off & 127u;
const uint32_t sr = row0 + rr;
const uint32_t comp_idx = sr < raw_count
? 0u
: (uint32_t)topk[(uint64_t)t * top_k + (sr - raw_count)];
const bool visible = sr < raw_count || comp_idx < visible_comp;
const float4 *src = sr < raw_count
? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim)
: (const float4 *)(comp_kv + (uint64_t)(visible ? comp_idx : 0u) * head_dim);
kv_shared[off] = visible ? src[c4] : make_float4(0, 0, 0, 0);
}
__syncthreads();
if (valid_head) {
for (uint32_t rr = 0; rr < nr; rr++) {
const uint32_t sr = row0 + rr;
if (sr >= raw_count &&
(uint32_t)topk[(uint64_t)t * top_k + sr - raw_count] >= visible_comp)
continue;
const float4 *kv4 = kv_shared + rr * 128u;
float4 k0 = kv4[lane + 0u];
float4 k1 = kv4[lane + 32u];
float4 k2 = kv4[lane + 64u];
float4 k3 = kv4[lane + 96u];
float score = dot4_f32(q0, k0) +
dot4_f32(q1, k1) +
dot4_f32(q2, k2) +
dot4_f32(q3, k3);
score = warp_sum_f32(score) * scale;
score = __shfl_sync(0xffffffffu, score, 0);
const float new_m = fmaxf(max_s, score);
const float old_scale = expf(max_s - new_m);
const float row_scale = expf(score - new_m);
sum_s = sum_s * old_scale + row_scale;
o0.x = o0.x * old_scale + k0.x * row_scale;
o0.y = o0.y * old_scale + k0.y * row_scale;
o0.z = o0.z * old_scale + k0.z * row_scale;
o0.w = o0.w * old_scale + k0.w * row_scale;
o1.x = o1.x * old_scale + k1.x * row_scale;
o1.y = o1.y * old_scale + k1.y * row_scale;
o1.z = o1.z * old_scale + k1.z * row_scale;
o1.w = o1.w * old_scale + k1.w * row_scale;
o2.x = o2.x * old_scale + k2.x * row_scale;
o2.y = o2.y * old_scale + k2.y * row_scale;
o2.z = o2.z * old_scale + k2.z * row_scale;
o2.w = o2.w * old_scale + k2.w * row_scale;
o3.x = o3.x * old_scale + k3.x * row_scale;
o3.y = o3.y * old_scale + k3.y * row_scale;
o3.z = o3.z * old_scale + k3.z * row_scale;
o3.w = o3.w * old_scale + k3.w * row_scale;
max_s = new_m;
}
}
__syncthreads();
}
if (valid_head) {
const float sink = sinks[head];
const float new_m = fmaxf(max_s, sink);
const float old_scale = expf(max_s - new_m);
const float sink_scale = expf(sink - new_m);
sum_s = sum_s * old_scale + sink_scale;
o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale;
o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale;
o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale;
o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale;
const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s;
o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s;
o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s;
o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s;
o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s;
float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim);
out4[lane + 0u] = o0;
out4[lane + 32u] = o1;
out4[lane + 64u] = o2;
out4[lane + 96u] = o3;
}
}
__global__ static void attention_static_mixed_heads8_online_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t head_group = blockIdx.y;
if (t >= n_tokens || head_dim != 512u) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t head = head_group * 8u + warp;
const bool valid_head = head < n_head;
__shared__ float4 kv_shared[4 * 128];
const uint32_t raw_count = window != 0u && t + 1u > window ? window : t + 1u;
const uint32_t raw_start = t + 1u - raw_count;
uint32_t comp_count = 0;
if (n_comp != 0u && ratio != 0u) {
comp_count = (t + 1u) / ratio;
if (comp_count > n_comp) comp_count = n_comp;
}
const uint32_t n_score = raw_count + comp_count;
const float scale = rsqrtf((float)head_dim);
const float4 *q4 = valid_head
? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim)
: NULL;
float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 q1 = q0, q2 = q0, q3 = q0;
if (valid_head) {
q0 = q4[lane + 0u];
q1 = q4[lane + 32u];
q2 = q4[lane + 64u];
q3 = q4[lane + 96u];
}
float max_s = -INFINITY;
float sum_s = 0.0f;
float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 o1 = o0, o2 = o0, o3 = o0;
for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) {
const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u;
for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) {
const uint32_t rr = off >> 7u;
const uint32_t c4 = off & 127u;
const uint32_t sr = row0 + rr;
const float4 *src = sr < raw_count
? (const float4 *)(raw_kv + (uint64_t)(raw_start + sr) * head_dim)
: (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim);
kv_shared[off] = src[c4];
}
__syncthreads();
if (valid_head) {
for (uint32_t rr = 0; rr < nr; rr++) {
const float4 *kv4 = kv_shared + rr * 128u;
float4 k0 = kv4[lane + 0u];
float4 k1 = kv4[lane + 32u];
float4 k2 = kv4[lane + 64u];
float4 k3 = kv4[lane + 96u];
float score = dot4_f32(q0, k0) +
dot4_f32(q1, k1) +
dot4_f32(q2, k2) +
dot4_f32(q3, k3);
score = warp_sum_f32(score) * scale;
score = __shfl_sync(0xffffffffu, score, 0);
const float new_m = fmaxf(max_s, score);
const float old_scale = expf(max_s - new_m);
const float row_scale = expf(score - new_m);
sum_s = sum_s * old_scale + row_scale;
o0.x = o0.x * old_scale + k0.x * row_scale;
o0.y = o0.y * old_scale + k0.y * row_scale;
o0.z = o0.z * old_scale + k0.z * row_scale;
o0.w = o0.w * old_scale + k0.w * row_scale;
o1.x = o1.x * old_scale + k1.x * row_scale;
o1.y = o1.y * old_scale + k1.y * row_scale;
o1.z = o1.z * old_scale + k1.z * row_scale;
o1.w = o1.w * old_scale + k1.w * row_scale;
o2.x = o2.x * old_scale + k2.x * row_scale;
o2.y = o2.y * old_scale + k2.y * row_scale;
o2.z = o2.z * old_scale + k2.z * row_scale;
o2.w = o2.w * old_scale + k2.w * row_scale;
o3.x = o3.x * old_scale + k3.x * row_scale;
o3.y = o3.y * old_scale + k3.y * row_scale;
o3.z = o3.z * old_scale + k3.z * row_scale;
o3.w = o3.w * old_scale + k3.w * row_scale;
max_s = new_m;
}
}
__syncthreads();
}
if (valid_head) {
const float sink = sinks[head];
const float new_m = fmaxf(max_s, sink);
const float old_scale = expf(max_s - new_m);
const float sink_scale = expf(sink - new_m);
sum_s = sum_s * old_scale + sink_scale;
o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale;
o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale;
o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale;
o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale;
const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s;
o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s;
o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s;
o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s;
o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s;
float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim);
out4[lane + 0u] = o0;
out4[lane + 32u] = o1;
out4[lane + 64u] = o2;
out4[lane + 96u] = o3;
}
}
/* -------------------------------------------------------------------------
* DS4_ATTN_TOKENTILE: token-tile HMMA indexed attention (opt-in).
* Ported from the pass-12 standalone prototype with fixed STAGE_ROWS=32.
* Four-token tiles reduce the selected-row union while eight heads keep the
* 32-row MMA shape full.
*/
static constexpr uint32_t kTTTileTokens = 4u;
static constexpr uint32_t kTTG = 8u;
static constexpr uint32_t kTTM = 32u;
static constexpr uint32_t kTTStageRows = 32u;
static constexpr uint32_t kTTRawWindow = 128u;
static constexpr uint32_t kTTHeadDim = 512u;
static constexpr uint32_t kTTWarps = 16u;
static constexpr uint32_t kTTThreads = 512u;
static constexpr uint32_t kTTScoreKQuarters = 4u;
static constexpr uint32_t kTTScoreKSliceDim = kTTHeadDim / kTTScoreKQuarters;
static constexpr uint32_t kTTScoreKStepsPerQuarter = kTTScoreKSliceDim / 16u;
static constexpr uint32_t kTTRecordRingPlanes = 4u;
static constexpr uint32_t kTTProbStride = 40u;
static constexpr uint32_t kTTRingChunkBytes = 16u;
static constexpr uint32_t kTTRingChunksPerRow =
(kTTHeadDim * sizeof(half)) / kTTRingChunkBytes;
static constexpr size_t kTTSmemHardCap = 90ull * 1024ull;
static_assert(kTTM == kTTTileTokens * kTTG, "token-tile M must be 4 tokens x G8");
static_assert(kTTProbStride == kTTStageRows + 8u, "token-tile prob stride changed");
static_assert(kTTScoreKSliceDim % 16u == 0, "token-tile score K split changed");
static_assert(kTTRingChunksPerRow == 64u, "token-tile KV ring expects 64 chunks");
static_assert(sizeof(int2) == 8u, "token-tile union record must remain 8 bytes");
template <uint32_t TT_STAGE_ROWS>
struct tt_TokentileLayout {
static constexpr uint32_t prob_stride = TT_STAGE_ROWS + 8u;
static constexpr uint32_t ring_plane_bytes =
TT_STAGE_ROWS * kTTRingChunksPerRow * kTTRingChunkBytes;
static constexpr uint32_t ring_plane_elems = ring_plane_bytes / sizeof(half);
};
__device__ static float tt_warp_sum_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v += __shfl_xor_sync(0xffffffffu, v, offset);
}
return v;
}
__device__ static float tt_warp_max_f32(float v) {
for (int offset = 16; offset > 0; offset >>= 1) {
v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, offset));
}
return v;
}
__device__ __forceinline__ uint32_t tt_lane_id(void) {
return threadIdx.x & 31u;
}
__device__ __forceinline__ uint32_t tt_warp_id(void) {
return threadIdx.x >> 5u;
}
__device__ __forceinline__ int tt_mma_c_i(uint32_t lane, int l) {
return ((l >> 1) << 3) + (int)(lane >> 2);
}
__device__ __forceinline__ int tt_mma_c_j(uint32_t lane, int l) {
return (int)((lane & 3u) << 1) + (l & 1);
}
__device__ __forceinline__ unsigned tt_smem_addr(const void *p) {
return static_cast<unsigned>(__cvta_generic_to_shared(p));
}
__device__ __forceinline__ uint32_t tt_ring_off_bytes(uint32_t row, uint32_t c) {
return (row * kTTRingChunksPerRow + (c ^ (row & 7u))) * kTTRingChunkBytes;
}
__device__ __forceinline__ void tt_ldmatrix_x4_addr(uint32_t (&r)[4], unsigned a) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
: "r"(a));
#else
(void)a;
r[0] = r[1] = r[2] = r[3] = 0;
#endif
}
__device__ __forceinline__ void tt_ldmatrix_x2_addr(uint32_t (&r)[2], unsigned a) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];"
: "=r"(r[0]), "=r"(r[1])
: "r"(a));
#else
(void)a;
r[0] = r[1] = 0;
#endif
}
__device__ __forceinline__ void tt_ldmatrix_x2_trans_addr(uint32_t (&r)[2], unsigned a) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.b16 {%0, %1}, [%2];"
: "=r"(r[0]), "=r"(r[1])
: "r"(a));
#else
(void)a;
r[0] = r[1] = 0;
#endif
}
__device__ __forceinline__ void tt_ldmatrix_x4(uint32_t (&r)[4], const void *p) {
tt_ldmatrix_x4_addr(r, tt_smem_addr(p));
}
__device__ __forceinline__ void tt_ldmatrix_x2(uint32_t (&r)[2], const void *p) {
tt_ldmatrix_x2_addr(r, tt_smem_addr(p));
}
__device__ __forceinline__ void tt_ldmatrix_x2_trans(uint32_t (&r)[2], const void *p) {
tt_ldmatrix_x2_trans_addr(r, tt_smem_addr(p));
}
__device__ __forceinline__ void tt_mma_m16n8k16_f16_f32(
float *d,
const uint32_t (&a)[4],
const uint32_t (&b)[2]) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
uint32_t d0 = __float_as_uint(d[0]);
uint32_t d1 = __float_as_uint(d[1]);
uint32_t d2 = __float_as_uint(d[2]);
uint32_t d3 = __float_as_uint(d[3]);
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
: "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3)
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
d[0] = __uint_as_float(d0);
d[1] = __uint_as_float(d1);
d[2] = __uint_as_float(d2);
d[3] = __uint_as_float(d3);
#else
(void)a;
(void)b;
#endif
}
__device__ __forceinline__ unsigned char *tt_align16(unsigned char *p) {
uintptr_t x = reinterpret_cast<uintptr_t>(p);
x = (x + 15u) & ~uintptr_t(15u);
return reinterpret_cast<unsigned char *>(x);
}
__device__ __forceinline__ void tt_zero_16B(void *dst) {
*reinterpret_cast<int4 *>(dst) = make_int4(0, 0, 0, 0);
}
__device__ __forceinline__ void tt_zero_8B(void *dst) {
*reinterpret_cast<int2 *>(dst) = make_int2(0, 0);
}
__device__ __forceinline__ void tt_cp_async_16B(void *dst, const void *src, bool pred) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
if (pred) {
const unsigned smem = tt_smem_addr(dst);
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
:: "r"(smem), "l"(src));
} else {
tt_zero_16B(dst);
}
#else
(void)src;
(void)pred;
tt_zero_16B(dst);
#endif
}
__device__ __forceinline__ void tt_cp_async_8B(void *dst, const void *src, bool pred) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
if (pred) {
const unsigned smem = tt_smem_addr(dst);
asm volatile("cp.async.ca.shared.global [%0], [%1], 8;"
:: "r"(smem), "l"(src));
} else {
tt_zero_8B(dst);
}
#else
(void)src;
(void)pred;
tt_zero_8B(dst);
#endif
}
__device__ __forceinline__ void tt_cp_async_commit(void) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("cp.async.commit_group;");
#endif
}
template <int KeepGroups>
__device__ __forceinline__ void tt_cp_async_wait_group(void) {
static_assert(KeepGroups >= 0 && KeepGroups <= 7, "bad cp.async wait_group depth");
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("cp.async.wait_group %0;" :: "n"(KeepGroups));
#endif
}
__device__ __forceinline__ uint32_t tt_score_partial_slot(
uint32_t kq,
uint32_t m,
uint32_t r) {
return (kq + r + m) & 3u;
}
template <uint32_t TT_STAGE_ROWS>
__device__ __forceinline__ void tt_store_score_partial(
float4 * __restrict__ partials,
uint32_t kq,
uint32_t m,
uint32_t r,
float v) {
float *dst = &partials[m * TT_STAGE_ROWS + r].x;
dst[tt_score_partial_slot(kq, m, r)] = v;
}
__device__ __forceinline__ float4 tt_load_score_partial_record(
const float4 * __restrict__ p) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
float x, y, z, w;
asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];"
: "=f"(x), "=f"(y), "=f"(z), "=f"(w)
: "r"(tt_smem_addr(p)));
return make_float4(x, y, z, w);
#else
return *p;
#endif
}
__device__ __forceinline__ void tt_issue_cp_async_row(
half * __restrict__ dst,
uint32_t rr,
uint32_t lane16,
const half * __restrict__ src,
bool live) {
const char *src_b = reinterpret_cast<const char *>(src);
char *dst_b = reinterpret_cast<char *>(dst);
#pragma unroll
for (uint32_t i = 0; i < 4u; ++i) {
const uint32_t chunk = lane16 + i * 16u;
char *db = dst_b + tt_ring_off_bytes(rr, chunk);
tt_cp_async_16B(db, live ? static_cast<const void *>(src_b + chunk * 16u)
: static_cast<const void *>(db),
live);
}
}
__device__ __forceinline__ uint32_t tt_stage_raw_rows(
uint32_t row0,
uint32_t nr,
uint32_t raw_union_count) {
uint32_t raw_rows = 0;
if (row0 < raw_union_count) {
const uint32_t raw_left = raw_union_count - row0;
raw_rows = raw_left < nr ? raw_left : nr;
}
return raw_rows;
}
__global__ static void __launch_bounds__(512, 1) attention_tokentile_union_build_kernel(
int2 *records,
uint32_t *counts,
const int32_t *topk,
const int32_t *positions,
uint32_t pos0,
uint32_t n_tokens,
uint32_t top_k,
uint32_t ratio,
uint32_t n_comp,
uint32_t rec_stride) {
extern __shared__ uint32_t bitmap[];
__shared__ uint32_t scan[513];
__shared__ uint32_t running_s;
const uint32_t tid = threadIdx.x;
const uint32_t tile_base = blockIdx.x * kTTTileTokens;
if (tile_base >= n_tokens) {
if (tid == 0u) counts[blockIdx.x] = 0u;
return;
}
const uint32_t tile_count =
n_tokens - tile_base < kTTTileTokens ? n_tokens - tile_base : kTTTileTokens;
const uint32_t bitmap_words = (n_comp + 1u) >> 1u;
for (uint32_t w = tid; w < bitmap_words; w += blockDim.x) {
bitmap[w] = 0u;
}
if (tid == 0u) running_s = 0u;
__syncthreads();
const uint32_t total_slots = tile_count * top_k;
for (uint32_t idx = tid; idx < total_slots; idx += blockDim.x) {
const uint32_t tok = idx / top_k;
const uint32_t i = idx - tok * top_k;
const uint32_t t = tile_base + tok;
const uint32_t qpos = positions ? (uint32_t)positions[t] : pos0 + t;
uint32_t visible = ratio ? (positions ? qpos / ratio : (qpos + 1u) / ratio) : n_comp;
if (visible > n_comp) visible = n_comp;
const int32_t c = topk[(uint64_t)t * top_k + i];
if (c >= 0 && (uint32_t)c < visible) {
const uint32_t cu = (uint32_t)c;
const uint32_t bits = ((uint32_t)(1u << tok)) << ((cu & 1u) * 16u);
atomicOr(&bitmap[cu >> 1u], bits);
}
}
__syncthreads();
for (uint32_t base = 0u; base < n_comp; base += blockDim.x) {
const uint32_t id = base + tid;
uint32_t mask = 0u;
if (id < n_comp) {
const uint32_t word = bitmap[id >> 1u];
mask = (id & 1u) ? (word >> 16u) : (word & 0xffffu);
}
const uint32_t pred = mask != 0u ? 1u : 0u;
if (tid == 0u) scan[0] = 0u;
scan[tid + 1u] = pred;
__syncthreads();
for (uint32_t off = 1u; off < blockDim.x; off <<= 1u) {
uint32_t v = 0u;
if (tid >= off) v = scan[tid + 1u - off];
__syncthreads();
scan[tid + 1u] += v;
__syncthreads();
}
const uint32_t rank = scan[tid];
const uint32_t total = scan[blockDim.x];
const uint32_t running = running_s;
if (pred) {
records[(uint64_t)blockIdx.x * rec_stride + running + rank] =
make_int2((int)id, (int)mask);
}
__syncthreads();
if (tid == 0u) running_s = running + total;
__syncthreads();
}
if (tid == 0u) counts[blockIdx.x] = running_s;
}
__global__ static void __launch_bounds__(256, 1) attention_tokentile_raw_mirror_kernel(
half *dst,
const float *raw_kv,
const int32_t *seq_id,
uint32_t tt_run_pos0,
uint32_t n_tokens,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t first_raw_pos,
uint32_t raw_row_min,
uint32_t head_dim) {
const uint32_t row = blockIdx.x;
const uint32_t d0 = threadIdx.x << 1u;
if (d0 >= head_dim) return;
half *dst_row = dst + (uint64_t)row * head_dim;
if (row < raw_row_min) {
dst_row[d0] = __float2half(0.0f);
if (d0 + 1u < head_dim) dst_row[d0 + 1u] = __float2half(0.0f);
return;
}
const int64_t p = (int64_t)tt_run_pos0 - (int64_t)(kTTRawWindow - 1u) + (int64_t)row;
uint32_t slot = 0u;
if (seq_id) {
slot = (uint32_t)seq_id[0] * raw_cap + (uint32_t)((uint64_t)p % raw_cap);
} else {
const uint32_t rel = (uint32_t)(p - (int64_t)first_raw_pos);
slot = (raw_start + rel) % raw_cap;
}
const float *src = raw_kv + (uint64_t)slot * head_dim;
dst_row[d0] = __float2half(src[d0]);
if (d0 + 1u < head_dim) dst_row[d0 + 1u] = __float2half(src[d0 + 1u]);
}
__global__ static void __launch_bounds__(256, 1) attention_tokentile_comp_mirror_kernel(
half *dst,
const float *comp_kv,
uint32_t n_comp,
uint32_t head_dim) {
const uint32_t row = blockIdx.x;
const uint32_t c4 = threadIdx.x;
if (row >= n_comp || c4 >= (head_dim >> 2u)) return;
const float4 v = ((const float4 *)(comp_kv + (uint64_t)row * head_dim))[c4];
half *out = dst + (uint64_t)row * head_dim + (c4 << 2u);
out[0] = __float2half(v.x);
out[1] = __float2half(v.y);
out[2] = __float2half(v.z);
out[3] = __float2half(v.w);
}
/* Non-indexed layers select the causal compressed range [0, visible), so
* their per-tile records can be built directly without a bitmap or sort. */
__global__ static void __launch_bounds__(512, 1)
attention_tokentile_dense_build_kernel(
int2 *records,
uint32_t *counts,
uint32_t pos0,
uint32_t n_tokens,
uint32_t ratio,
uint32_t n_comp,
uint32_t rec_stride) {
__shared__ uint32_t visible_s[kTTTileTokens];
const uint32_t tid = threadIdx.x;
const uint32_t tile_base = blockIdx.x * kTTTileTokens;
if (tile_base >= n_tokens) {
if (tid == 0u) counts[blockIdx.x] = 0u;
return;
}
const uint32_t tile_count =
n_tokens - tile_base < kTTTileTokens
? n_tokens - tile_base : kTTTileTokens;
if (tid < kTTTileTokens) {
uint32_t visible = 0u;
if (tid < tile_count && n_comp != 0u && ratio != 0u) {
visible = (pos0 + tile_base + tid + 1u) / ratio;
if (visible > n_comp) visible = n_comp;
}
visible_s[tid] = visible;
}
__syncthreads();
uint32_t vmax = 0u;
for (uint32_t i = 0; i < tile_count; i++) {
if (visible_s[i] > vmax) vmax = visible_s[i];
}
for (uint32_t c = tid; c < vmax; c += blockDim.x) {
uint32_t mask = 0u;
for (uint32_t i = 0; i < tile_count; i++) {
if (c < visible_s[i]) mask |= 1u << i;
}
records[(uint64_t)blockIdx.x * rec_stride + c] =
make_int2((int)c, (int)mask);
}
if (tid == 0u) counts[blockIdx.x] = vmax;
}
template <uint32_t TT_STAGE_ROWS>
__device__ __forceinline__ void tt_issue_record_stage_cp_async(
int2 * __restrict__ rec_plane,
uint32_t row0,
uint32_t nr,
uint32_t raw_union_count,
const int2 * __restrict__ union_records_tile) {
static_assert(TT_STAGE_ROWS == 32u, "token-tile record issue is fixed at R32");
const uint32_t raw_rows = tt_stage_raw_rows(row0, nr, raw_union_count);
const uint32_t lane = tt_lane_id();
const uint32_t warp = tt_warp_id();
const uint32_t lane16 = lane & 15u;
const uint32_t rr = warp * 2u + (lane >> 4u);
const bool active = rr < TT_STAGE_ROWS;
const bool comp_live = active && rr >= raw_rows && rr < nr;
if (comp_live && lane16 == 0u) {
int2 *dst = rec_plane + rr;
const uint32_t ci = row0 + rr - raw_union_count;
tt_cp_async_8B(dst, union_records_tile + ci, true);
}
}
template <uint32_t TT_STAGE_ROWS, bool USE_SMEM_RECORDS>
__device__ __forceinline__ void tt_issue_kv_stage_cp_async(
half * __restrict__ dst,
const int2 * __restrict__ rec_plane,
uint32_t row0,
uint32_t nr,
uint32_t raw_union_count,
const int2 * __restrict__ union_records_tile,
uint32_t tile_base,
const half * __restrict__ raw_kv,
const half * __restrict__ comp_kv,
uint32_t tid) {
constexpr uint32_t kCp16PerRow = (kTTHeadDim * sizeof(half)) / 16u;
static_assert(kCp16PerRow == 64u, "expected 64 cp.async chunks per f16 KV row");
static_assert(TT_STAGE_ROWS == 32u, "token-tile KV issue is fixed at R32");
(void)tid;
const uint32_t raw_rows = tt_stage_raw_rows(row0, nr, raw_union_count);
const uint32_t lane = tt_lane_id();
const uint32_t warp = tt_warp_id();
const uint32_t lane16 = lane & 15u;
const uint32_t rr = warp * 2u + (lane >> 4u);
const bool active = rr < TT_STAGE_ROWS;
if (active && rr < raw_rows) {
const uint32_t sr = row0 + rr;
const half *src = raw_kv + (uint64_t)(tile_base + sr) * kTTHeadDim;
tt_issue_cp_async_row(dst, rr, lane16, src, true);
}
if (active && rr >= raw_rows && rr < nr) {
uint32_t comp_id = 0u;
if (USE_SMEM_RECORDS) {
comp_id = (uint32_t)rec_plane[rr].x;
} else {
const uint32_t ci = row0 + rr - raw_union_count;
comp_id = (uint32_t)union_records_tile[ci].x;
}
const half *src = comp_kv + (uint64_t)comp_id * kTTHeadDim;
tt_issue_cp_async_row(dst, rr, lane16, src, true);
}
if (active && rr >= nr) {
tt_issue_cp_async_row(dst, rr, lane16, NULL, false);
}
}
__device__ __forceinline__ void tt_load_score_q_frag(
uint32_t (&q_frag)[kTTScoreKStepsPerQuarter][4],
const float * __restrict__ q,
uint32_t n_tokens,
uint32_t n_head,
uint32_t tile_base,
uint32_t head_base) {
constexpr uint32_t kMtiles = kTTM / 16u;
constexpr uint32_t kScoreWarps = kMtiles * kTTScoreKQuarters;
const uint32_t warp = tt_warp_id();
if (warp >= kScoreWarps) {
return;
}
const uint32_t mtile = warp >> 2u;
const uint32_t kq = warp & 3u;
const uint32_t k_base = kq * kTTScoreKSliceDim;
const uint32_t lane = tt_lane_id();
const uint32_t a_group = lane >> 2u;
const uint32_t a_col_pair = (lane & 3u) << 1u;
#pragma unroll
for (uint32_t kt = 0; kt < kTTScoreKStepsPerQuarter; ++kt) {
const uint32_t k0 = k_base + kt * 16u;
#pragma unroll
for (uint32_t r = 0; r < 4u; ++r) {
const uint32_t m = mtile * 16u + a_group + ((r & 1u) ? 8u : 0u);
const uint32_t tok = m / kTTG;
const uint32_t h = m - tok * kTTG;
const uint32_t gt = tile_base + tok;
const uint32_t gh = head_base + h;
const uint32_t d = k0 + a_col_pair + ((r & 2u) ? 8u : 0u);
float x0 = 0.0f;
float x1 = 0.0f;
if (gt < n_tokens && gh < n_head) {
const float *q_row = q + ((uint64_t)gt * n_head + gh) * kTTHeadDim;
x0 = q_row[d];
x1 = q_row[d + 1u];
}
const half2 packed = __floats2half2_rn(x0, x1);
q_frag[kt][r] =
(uint32_t)__half_as_ushort(__low2half(packed)) |
((uint32_t)__half_as_ushort(__high2half(packed)) << 16);
}
}
}
template <uint32_t TT_STAGE_ROWS>
__device__ __forceinline__ void tt_hmma_score_stage(
float4 * __restrict__ partial_scores,
const uint32_t (&q_frag)[kTTScoreKStepsPerQuarter][4],
const half * __restrict__ kv_cur,
uint32_t nr,
float score_scale) {
constexpr uint32_t kMtiles = kTTM / 16u;
constexpr uint32_t kScoreWarps = kMtiles * kTTScoreKQuarters;
constexpr uint32_t kNtiles = TT_STAGE_ROWS / 8u;
const uint32_t warp = tt_warp_id();
if (warp >= kScoreWarps) {
return;
}
const uint32_t mtile = warp >> 2u;
const uint32_t kq = warp & 3u;
const uint32_t lane = tt_lane_id();
const unsigned kv_smem = tt_smem_addr(kv_cur);
const uint32_t score_row_lane = lane & 7u;
const uint32_t score_chunk_lane = (lane >> 3u) & 1u;
const uint32_t score_chunk_base = kq * (kTTScoreKSliceDim / 8u) + score_chunk_lane;
#pragma unroll
for (uint32_t ntile = 0; ntile < kNtiles; ++ntile) {
const uint32_t row_base = ntile * 8u;
if (row_base < nr) {
float s_frag[4] = {0.0f, 0.0f, 0.0f, 0.0f};
const uint32_t score_row = row_base + score_row_lane;
#pragma unroll
for (uint32_t kt = 0; kt < kTTScoreKStepsPerQuarter; ++kt) {
uint32_t b[2];
tt_ldmatrix_x2_addr(
b,
kv_smem + tt_ring_off_bytes(score_row, score_chunk_base + kt * 2u));
tt_mma_m16n8k16_f16_f32(s_frag, q_frag[kt], b);
}
const uint32_t m0 = mtile * 16u + (lane >> 2u);
const uint32_t r0 = row_base + ((lane & 3u) << 1u);
const uint32_t r1 = r0 + 1u;
const uint32_t m1 = m0 + 8u;
if (r0 < nr) {
tt_store_score_partial<TT_STAGE_ROWS>(
partial_scores, kq, m0, r0, s_frag[0] * score_scale);
tt_store_score_partial<TT_STAGE_ROWS>(
partial_scores, kq, m1, r0, s_frag[2] * score_scale);
}
if (r1 < nr) {
tt_store_score_partial<TT_STAGE_ROWS>(
partial_scores, kq, m0, r1, s_frag[1] * score_scale);
tt_store_score_partial<TT_STAGE_ROWS>(
partial_scores, kq, m1, r1, s_frag[3] * score_scale);
}
}
}
}
template <uint32_t TT_STAGE_ROWS>
__device__ __forceinline__ void tt_softmax_stage(
half * __restrict__ probs,
float * __restrict__ stage_rescale,
float * __restrict__ max_s,
float * __restrict__ sum_s,
const float4 * __restrict__ scores,
const int2 * __restrict__ records,
uint32_t row0,
uint32_t nr,
uint32_t raw_union_count,
uint32_t tile_count,
uint32_t tile_base,
uint32_t raw_row_min) {
constexpr uint32_t kMPerWarp = kTTM / kTTWarps;
constexpr uint32_t kProbStride = tt_TokentileLayout<TT_STAGE_ROWS>::prob_stride;
static_assert(kTTM % kTTWarps == 0u, "softmax maps integral m rows per warp");
const uint32_t lane = tt_lane_id();
const uint32_t warp = tt_warp_id();
const uint32_t sr = row0 + lane;
const bool lane_live = lane < TT_STAGE_ROWS && lane < nr;
const bool raw_slot = lane_live && sr < raw_union_count;
uint16_t comp_mask = 0u;
if (lane_live && !raw_slot) {
comp_mask = (uint16_t)records[lane].y;
}
const uint32_t prob_lane_base = warp * kProbStride + lane;
#pragma unroll
for (uint32_t mi = 0; mi < kMPerWarp; ++mi) {
const uint32_t m = warp + mi * kTTWarps;
const uint32_t tok = m / kTTG;
const bool valid_token = tok < tile_count;
const uint32_t score_idx = m * TT_STAGE_ROWS + lane;
const uint32_t prob_idx = prob_lane_base + mi * kTTWarps * kProbStride;
float score = -INFINITY;
if (lane_live && valid_token) {
const bool selected = raw_slot
? (((uint32_t)(sr - tok) < kTTRawWindow) && (tile_base + sr >= raw_row_min))
: ((comp_mask & (uint16_t)(1u << tok)) != 0u);
if (selected) {
const float4 parts = tt_load_score_partial_record(scores + score_idx);
const float s01 = parts.x + parts.y;
const float s23 = parts.z + parts.w;
score = s01 + s23;
}
}
const float stage_m = tt_warp_max_f32(score);
const float old_m = max_s[m];
const float new_m = fmaxf(old_m, stage_m);
float old_scale = 1.0f;
if (new_m != -INFINITY) {
old_scale = (old_m == -INFINITY) ? 0.0f : expf(old_m - new_m);
}
const float row_scale = (score == -INFINITY || new_m == -INFINITY)
? 0.0f
: expf(score - new_m);
const float stage_sum = tt_warp_sum_f32(row_scale);
if (lane < TT_STAGE_ROWS) {
probs[prob_idx] = __float2half(lane < nr ? row_scale : 0.0f);
}
if (lane == 0u) {
max_s[m] = new_m;
sum_s[m] = sum_s[m] * old_scale + stage_sum;
stage_rescale[m] = old_scale;
}
}
}
template <uint32_t TT_STAGE_ROWS>
__device__ __forceinline__ void tt_pv_mma_stage(
float (&o_acc)[2u * kTTTileTokens * kTTG],
const half * __restrict__ probs,
const float * __restrict__ stage_rescale,
const half * __restrict__ kv_cur) {
constexpr uint32_t kMtiles = kTTM / 16u;
constexpr uint32_t kPvWarpBase = 8u;
constexpr uint32_t kPvNTiles = 8u;
const uint32_t lane = tt_lane_id();
const uint32_t warp = tt_warp_id();
if (warp < kPvWarpBase) {
return;
}
const uint32_t pv_warp = warp - kPvWarpBase;
#pragma unroll
for (uint32_t mtile = 0; mtile < kMtiles; ++mtile) {
const float *scale0 = stage_rescale + mtile * 16u + (lane >> 2u);
const float *scale1 = scale0 + 8u;
const float rs0 = *scale0;
const float rs1 = *scale1;
#pragma unroll
for (uint32_t ntile = 0; ntile < kPvNTiles; ++ntile) {
const uint32_t idx = ((mtile * kPvNTiles + ntile) << 2);
o_acc[idx + 0u] *= rs0;
o_acc[idx + 1u] *= rs0;
o_acc[idx + 2u] *= rs1;
o_acc[idx + 3u] *= rs1;
}
}
constexpr uint32_t kProbStride = tt_TokentileLayout<TT_STAGE_ROWS>::prob_stride;
constexpr unsigned kPvAStepBytes = 16u * sizeof(half);
constexpr unsigned kPvMtileBytes = 16u * kProbStride * sizeof(half);
const unsigned probs_lane_base =
tt_smem_addr(probs) +
(unsigned)(((lane & 15u) * (kProbStride / 2u) +
(lane >> 4u) * 4u) * sizeof(uint32_t));
const unsigned kv_smem = tt_smem_addr(kv_cur);
const uint32_t pv_row_lane = lane & 15u;
const uint32_t pv_chunk_base = pv_warp * kPvNTiles;
#pragma unroll
for (uint32_t kt = 0; kt < TT_STAGE_ROWS / 16u; ++kt) {
const unsigned probs_kt_base = probs_lane_base + (unsigned)(kt * kPvAStepBytes);
const uint32_t pv_row = kt * 16u + pv_row_lane;
#pragma unroll
for (uint32_t mtile = 0; mtile < kMtiles; ++mtile) {
uint32_t a[4];
tt_ldmatrix_x4_addr(a, probs_kt_base + (unsigned)(mtile * kPvMtileBytes));
#pragma unroll
for (uint32_t ntile = 0; ntile < kPvNTiles; ++ntile) {
uint32_t b[2];
const uint32_t idx = ((mtile * kPvNTiles + ntile) << 2);
tt_ldmatrix_x2_trans_addr(
b,
kv_smem + tt_ring_off_bytes(pv_row, pv_chunk_base + ntile));
tt_mma_m16n8k16_f16_f32(o_acc + idx, a, b);
}
}
}
}
__device__ __forceinline__ void tt_pv_mma_epilogue(
const float (&o_acc)[2u * kTTTileTokens * kTTG],
const float * __restrict__ final_scale,
float * __restrict__ heads,
uint32_t n_tokens,
uint32_t n_head,
uint32_t tile_base,
uint32_t head_base) {
constexpr uint32_t kMtiles = kTTM / 16u;
constexpr uint32_t kPvWarpBase = 8u;
constexpr uint32_t kPvNTiles = 8u;
const uint32_t lane = tt_lane_id();
const uint32_t warp = tt_warp_id();
if (warp < kPvWarpBase) {
return;
}
const uint32_t pv_warp = warp - kPvWarpBase;
#pragma unroll
for (uint32_t mtile = 0; mtile < kMtiles; ++mtile) {
#pragma unroll
for (uint32_t ntile = 0; ntile < kPvNTiles; ++ntile) {
#pragma unroll
for (int l = 0; l < 4; ++l) {
const uint32_t idx = ((mtile * kPvNTiles + ntile) << 2) + (uint32_t)l;
const uint32_t m = mtile * 16u + (uint32_t)tt_mma_c_i(lane, l);
const uint32_t tok = m / kTTG;
const uint32_t h = m - tok * kTTG;
const uint32_t gt = tile_base + tok;
const uint32_t gh = head_base + h;
const uint32_t d =
pv_warp * 64u + ntile * 8u + (uint32_t)tt_mma_c_j(lane, l);
if (gt < n_tokens && gh < n_head) {
heads[((uint64_t)gt * n_head + gh) * kTTHeadDim + d] =
o_acc[idx] * final_scale[m];
}
}
}
}
}
__global__ static void __launch_bounds__(512, 1) attention_tokentile_hmma_kernel(
float *heads,
const float *sinks,
const float *q,
const half *raw_kv,
const half *comp_kv,
const int2 *union_records,
const uint32_t *union_counts,
uint32_t rec_stride,
uint32_t n_tokens,
uint32_t n_head,
uint32_t raw_row_min) {
constexpr uint32_t kKvElems = tt_TokentileLayout<kTTStageRows>::ring_plane_elems;
constexpr uint32_t kProbStride = tt_TokentileLayout<kTTStageRows>::prob_stride;
const uint32_t tid = threadIdx.x;
const uint32_t tile_idx = blockIdx.x;
const uint32_t head_group = blockIdx.y;
const uint32_t tile_base = tile_idx * kTTTileTokens;
const uint32_t head_base = head_group * kTTG;
if (tile_base >= n_tokens) {
return;
}
const uint32_t tile_count =
n_tokens - tile_base < kTTTileTokens ? n_tokens - tile_base : kTTTileTokens;
const uint32_t raw_union_count = tile_count + kTTRawWindow - 1u;
const uint32_t comp_union_count = union_counts[tile_idx];
const uint32_t n_score = raw_union_count + comp_union_count;
const uint64_t union_tile_off = (uint64_t)tile_idx * rec_stride;
const int2 * __restrict__ union_records_tile = union_records + union_tile_off;
const float score_scale = rsqrtf((float)kTTHeadDim);
extern __shared__ unsigned char smem[];
unsigned char *p = tt_align16(smem);
half *kv_h = reinterpret_cast<half *>(p);
p = tt_align16(p + 2u * tt_TokentileLayout<kTTStageRows>::ring_plane_bytes);
half *probs = reinterpret_cast<half *>(p);
p = tt_align16(p + 2u * kTTM * kProbStride * sizeof(half));
float4 *score_scratch = reinterpret_cast<float4 *>(p);
p = tt_align16(p + kTTM * kTTStageRows * sizeof(float4));
float *max_s = reinterpret_cast<float *>(p);
p = tt_align16(p + kTTM * sizeof(float));
float *sum_s = reinterpret_cast<float *>(p);
p = tt_align16(p + kTTM * sizeof(float));
float *stage_rescale = reinterpret_cast<float *>(p);
p = tt_align16(p + kTTM * sizeof(float));
float *final_scale = reinterpret_cast<float *>(p);
p = tt_align16(p + kTTM * sizeof(float));
int2 *rec_ring = reinterpret_cast<int2 *>(p);
p = tt_align16(p + kTTRecordRingPlanes * kTTStageRows * sizeof(int2));
for (uint32_t m = tid; m < kTTM; m += blockDim.x) {
max_s[m] = -INFINITY;
sum_s[m] = 0.0f;
stage_rescale[m] = 1.0f;
final_scale[m] = 0.0f;
}
__syncthreads();
union tt_TokentileRoleRegs {
uint32_t score_q_frag[kTTScoreKStepsPerQuarter][4];
float o_acc[2u * kTTTileTokens * kTTG];
};
tt_TokentileRoleRegs role_regs;
tt_load_score_q_frag(
role_regs.score_q_frag, q, n_tokens, n_head, tile_base, head_base);
if (tt_warp_id() >= 8u) {
#pragma unroll
for (uint32_t i = 0; i < 2u * kTTM; ++i) {
role_regs.o_acc[i] = 0.0f;
}
}
uint32_t cur = 0u;
uint32_t free = 1u;
if (n_score != 0u) {
const uint32_t nr0 = n_score < kTTStageRows ? n_score : kTTStageRows;
tt_issue_kv_stage_cp_async<kTTStageRows, false>(
kv_h + cur * kKvElems,
rec_ring,
0u,
nr0,
raw_union_count,
union_records_tile,
tile_base,
raw_kv,
comp_kv,
tid);
tt_issue_record_stage_cp_async<kTTStageRows>(
rec_ring,
0u,
nr0,
raw_union_count,
union_records_tile);
if (kTTStageRows < n_score) {
const uint32_t nr1 =
n_score - kTTStageRows < kTTStageRows ? n_score - kTTStageRows : kTTStageRows;
tt_issue_record_stage_cp_async<kTTStageRows>(
rec_ring + kTTStageRows,
kTTStageRows,
nr1,
raw_union_count,
union_records_tile);
}
tt_cp_async_commit();
tt_cp_async_wait_group<0>();
}
__syncthreads();
uint32_t prob_cur = 0u;
for (uint32_t row0 = 0u; row0 < n_score; row0 += kTTStageRows) {
const uint32_t nr = n_score - row0 < kTTStageRows ? n_score - row0 : kTTStageRows;
half *kv_cur = kv_h + cur * kKvElems;
half *kv_free = kv_h + free * kKvElems;
tt_hmma_score_stage<kTTStageRows>(
score_scratch, role_regs.score_q_frag, kv_cur, nr, score_scale);
if (row0 != 0u) {
const half *kv_prev = kv_h + (cur ^ 1u) * kKvElems;
const half *probs_prev = probs + (prob_cur ^ 1u) * kTTM * kProbStride;
tt_pv_mma_stage<kTTStageRows>(
role_regs.o_acc, probs_prev, stage_rescale, kv_prev);
}
__syncthreads();
const uint32_t next_row0 = row0 + kTTStageRows;
const bool has_next = next_row0 < n_score;
if (has_next) {
const uint32_t next_nr =
n_score - next_row0 < kTTStageRows ? n_score - next_row0 : kTTStageRows;
tt_issue_kv_stage_cp_async<kTTStageRows, true>(
kv_free,
rec_ring + (((row0 / kTTStageRows) + 1u) & 3u) * kTTStageRows,
next_row0,
next_nr,
raw_union_count,
union_records_tile,
tile_base,
raw_kv,
comp_kv,
tid);
const uint32_t prefetch_row0 = next_row0 + kTTStageRows;
if (prefetch_row0 < n_score) {
const uint32_t prefetch_nr =
n_score - prefetch_row0 < kTTStageRows
? n_score - prefetch_row0
: kTTStageRows;
tt_issue_record_stage_cp_async<kTTStageRows>(
rec_ring + (((row0 / kTTStageRows) + 2u) & 3u) * kTTStageRows,
prefetch_row0,
prefetch_nr,
raw_union_count,
union_records_tile);
}
tt_cp_async_commit();
}
half *probs_cur = probs + prob_cur * kTTM * kProbStride;
tt_softmax_stage<kTTStageRows>(
probs_cur,
stage_rescale,
max_s,
sum_s,
score_scratch,
rec_ring + ((row0 / kTTStageRows) & 3u) * kTTStageRows,
row0,
nr,
raw_union_count,
tile_count,
tile_base,
raw_row_min);
tt_cp_async_wait_group<0>();
__syncthreads();
cur ^= 1u;
free ^= 1u;
prob_cur ^= 1u;
}
if (n_score != 0u) {
const half *kv_prev = kv_h + (cur ^ 1u) * kKvElems;
const half *probs_prev = probs + (prob_cur ^ 1u) * kTTM * kProbStride;
tt_pv_mma_stage<kTTStageRows>(
role_regs.o_acc, probs_prev, stage_rescale, kv_prev);
}
__syncthreads();
for (uint32_t m = tid; m < kTTM; m += blockDim.x) {
const uint32_t tok = m / kTTG;
const uint32_t h = m - tok * kTTG;
const uint32_t gt = tile_base + tok;
const uint32_t gh = head_base + h;
if (gt < n_tokens && gh < n_head) {
const float sink = sinks[gh];
const float old_m = max_s[m];
const float new_m = fmaxf(old_m, sink);
const float old_scale = old_m == -INFINITY ? 0.0f : expf(old_m - new_m);
const float sink_scale = expf(sink - new_m);
const float den = sum_s[m] * old_scale + sink_scale;
final_scale[m] = den == 0.0f ? 0.0f : old_scale / den;
} else {
final_scale[m] = 0.0f;
}
}
__syncthreads();
tt_pv_mma_epilogue(
role_regs.o_acc, final_scale, heads, n_tokens, n_head, tile_base, head_base);
}
constexpr size_t tt_align16_const(size_t x) {
return (x + 15u) & ~size_t(15u);
}
template <uint32_t TT_STAGE_ROWS, uint32_t TT_G>
struct tt_TokentileSmemBudget {
static constexpr uint32_t M = kTTTileTokens * TT_G;
static constexpr uint32_t prob_stride = tt_TokentileLayout<TT_STAGE_ROWS>::prob_stride;
static constexpr size_t q_bytes = 0;
static constexpr size_t ring_bytes = 2ull * tt_TokentileLayout<TT_STAGE_ROWS>::ring_plane_bytes;
static constexpr size_t p_bytes = 2ull * M * prob_stride * sizeof(half);
static constexpr size_t partial_records = (size_t)M * TT_STAGE_ROWS;
static constexpr size_t partial_bytes = partial_records * sizeof(float4);
static constexpr size_t stats_bytes = 4ull * M * sizeof(float);
static constexpr size_t record_bytes =
(size_t)kTTRecordRingPlanes * TT_STAGE_ROWS * sizeof(int2);
static constexpr size_t total =
tt_align16_const(
tt_align16_const(
tt_align16_const(
tt_align16_const(
tt_align16_const(
tt_align16_const(q_bytes) + ring_bytes) + p_bytes) +
partial_bytes) + stats_bytes) + record_bytes);
};
static_assert(tt_TokentileSmemBudget<kTTStageRows, kTTG>::p_bytes == 5120ull,
"M32/R32 P double-buffer must use R+8 stride");
static_assert(sizeof(float4) == 16u, "score partial records must stay 16 bytes");
static_assert(tt_TokentileSmemBudget<kTTStageRows, kTTG>::partial_bytes == 16ull * 1024ull,
"M32/R32 score partial records must be M*R float4");
static_assert(tt_TokentileSmemBudget<kTTStageRows, kTTG>::record_bytes == 1024ull,
"M32/R32 record ring must be four R-row int2 planes");
static_assert(tt_TokentileSmemBudget<kTTStageRows, kTTG>::total == 88576ull,
"M32/R32 total dynamic shared memory changed unexpectedly");
static_assert(tt_TokentileSmemBudget<kTTStageRows, kTTG>::total <= kTTSmemHardCap,
"token-tile dynamic shared memory must stay under the 90 KiB pass gate");
static int ds4_cuda_attn_tokentile_arch_ok(void) {
int device = 0;
cudaDeviceProp prop;
if (cudaGetDevice(&device) != cudaSuccess ||
cudaGetDeviceProperties(&prop, device) != cudaSuccess) {
(void)cudaGetLastError();
return 0;
}
return prop.major >= 8;
}
__global__ static void __launch_bounds__(256, 4)
attention_decode_mixed_heads8_online_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t t = blockIdx.x;
uint32_t head_group = blockIdx.y;
if (t >= n_tokens || head_dim != 512u) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t head = head_group * 8u + warp;
const bool valid_head = head < n_head;
__shared__ uint32_t raw_rows[256];
__shared__ uint32_t raw_count_s;
__shared__ uint32_t raw_first_idx_s;
__shared__ float4 kv_shared[4 * 128];
const uint32_t qpos = pos0 + t;
const uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t comp_count = 0;
if (n_comp != 0u) {
if (n_tokens == 1u && ratio == 0u) {
comp_count = n_comp;
} else if (ratio != 0u) {
comp_count = (qpos + 1u) / ratio;
if (comp_count > n_comp) comp_count = n_comp;
}
}
if (threadIdx.x == 0) {
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0u) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0u && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
raw_count_s = raw_count;
raw_first_idx_s = raw_first_idx;
}
__syncthreads();
const uint32_t raw_count = raw_count_s;
const uint32_t raw_first_idx = raw_first_idx_s;
for (uint32_t r = threadIdx.x; r < raw_count; r += blockDim.x) {
raw_rows[r] = (raw_start + raw_first_idx + r) % raw_cap;
}
__syncthreads();
const uint32_t n_score = raw_count + comp_count;
const float scale = rsqrtf((float)head_dim);
const float4 *q4 = valid_head
? (const float4 *)(q + ((uint64_t)t * n_head + head) * head_dim)
: NULL;
float4 q0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 q1 = q0, q2 = q0, q3 = q0;
if (valid_head) {
q0 = q4[lane + 0u];
q1 = q4[lane + 32u];
q2 = q4[lane + 64u];
q3 = q4[lane + 96u];
}
float max_s = -INFINITY;
float sum_s = 0.0f;
float4 o0 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
float4 o1 = o0, o2 = o0, o3 = o0;
for (uint32_t row0 = 0; row0 < n_score; row0 += 4u) {
const uint32_t nr = n_score - row0 < 4u ? n_score - row0 : 4u;
for (uint32_t off = threadIdx.x; off < nr * 128u; off += blockDim.x) {
const uint32_t rr = off >> 7u;
const uint32_t c4 = off & 127u;
const uint32_t sr = row0 + rr;
const float4 *src = sr < raw_count
? (const float4 *)(raw_kv + (uint64_t)raw_rows[sr] * head_dim)
: (const float4 *)(comp_kv + (uint64_t)(sr - raw_count) * head_dim);
kv_shared[off] = src[c4];
}
__syncthreads();
if (valid_head) {
for (uint32_t rr = 0; rr < nr; rr++) {
const float4 *kv4 = kv_shared + rr * 128u;
float4 k0 = kv4[lane + 0u];
float4 k1 = kv4[lane + 32u];
float4 k2 = kv4[lane + 64u];
float4 k3 = kv4[lane + 96u];
float score = dot4_f32(q0, k0) +
dot4_f32(q1, k1) +
dot4_f32(q2, k2) +
dot4_f32(q3, k3);
score = warp_sum_f32(score) * scale;
score = __shfl_sync(0xffffffffu, score, 0);
const float new_m = fmaxf(max_s, score);
const float old_scale = expf(max_s - new_m);
const float row_scale = expf(score - new_m);
sum_s = sum_s * old_scale + row_scale;
o0.x = o0.x * old_scale + k0.x * row_scale;
o0.y = o0.y * old_scale + k0.y * row_scale;
o0.z = o0.z * old_scale + k0.z * row_scale;
o0.w = o0.w * old_scale + k0.w * row_scale;
o1.x = o1.x * old_scale + k1.x * row_scale;
o1.y = o1.y * old_scale + k1.y * row_scale;
o1.z = o1.z * old_scale + k1.z * row_scale;
o1.w = o1.w * old_scale + k1.w * row_scale;
o2.x = o2.x * old_scale + k2.x * row_scale;
o2.y = o2.y * old_scale + k2.y * row_scale;
o2.z = o2.z * old_scale + k2.z * row_scale;
o2.w = o2.w * old_scale + k2.w * row_scale;
o3.x = o3.x * old_scale + k3.x * row_scale;
o3.y = o3.y * old_scale + k3.y * row_scale;
o3.z = o3.z * old_scale + k3.z * row_scale;
o3.w = o3.w * old_scale + k3.w * row_scale;
max_s = new_m;
}
}
__syncthreads();
}
if (valid_head) {
const float sink = sinks[head];
const float new_m = fmaxf(max_s, sink);
const float old_scale = expf(max_s - new_m);
const float sink_scale = expf(sink - new_m);
sum_s = sum_s * old_scale + sink_scale;
o0.x *= old_scale; o0.y *= old_scale; o0.z *= old_scale; o0.w *= old_scale;
o1.x *= old_scale; o1.y *= old_scale; o1.z *= old_scale; o1.w *= old_scale;
o2.x *= old_scale; o2.y *= old_scale; o2.z *= old_scale; o2.w *= old_scale;
o3.x *= old_scale; o3.y *= old_scale; o3.z *= old_scale; o3.w *= old_scale;
const float inv_s = sum_s == 0.0f ? 0.0f : 1.0f / sum_s;
o0.x *= inv_s; o0.y *= inv_s; o0.z *= inv_s; o0.w *= inv_s;
o1.x *= inv_s; o1.y *= inv_s; o1.z *= inv_s; o1.w *= inv_s;
o2.x *= inv_s; o2.y *= inv_s; o2.z *= inv_s; o2.w *= inv_s;
o3.x *= inv_s; o3.y *= inv_s; o3.z *= inv_s; o3.w *= inv_s;
float4 *out4 = (float4 *)(heads + ((uint64_t)t * n_head + head) * head_dim);
out4[lane + 0u] = o0;
out4[lane + 32u] = o1;
out4[lane + 64u] = o2;
out4[lane + 96u] = o3;
}
}
__device__ static void hc4_split_one(float *out, const float *mix, const float *scale, const float *base, uint32_t sinkhorn_iters, float epsv) {
const float pre_scale = scale[0];
const float post_scale = scale[1];
const float comb_scale = scale[2];
for (int i = 0; i < 4; i++) {
float z = mix[i] * pre_scale + base[i];
out[i] = 1.0f / (1.0f + expf(-z)) + epsv;
}
for (int i = 0; i < 4; i++) {
float z = mix[4 + i] * post_scale + base[4 + i];
out[4 + i] = 2.0f / (1.0f + expf(-z));
}
float c[16];
for (int r = 0; r < 4; r++) {
float m = -INFINITY;
for (int col = 0; col < 4; col++) {
float v = mix[8 + r * 4 + col] * comb_scale + base[8 + r * 4 + col];
c[r * 4 + col] = v;
m = fmaxf(m, v);
}
float s = 0.0f;
for (int col = 0; col < 4; col++) {
float v = expf(c[r * 4 + col] - m);
c[r * 4 + col] = v;
s += v;
}
for (int col = 0; col < 4; col++) c[r * 4 + col] = c[r * 4 + col] / s + epsv;
}
for (int col = 0; col < 4; col++) {
float s = epsv;
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
}
for (uint32_t iter = 1; iter < sinkhorn_iters; iter++) {
for (int r = 0; r < 4; r++) {
float s = epsv;
for (int col = 0; col < 4; col++) s += c[r * 4 + col];
for (int col = 0; col < 4; col++) c[r * 4 + col] /= s;
}
for (int col = 0; col < 4; col++) {
float s = epsv;
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
}
}
for (int i = 0; i < 16; i++) out[8 + i] = c[i];
}
__global__ static void hc_split_sinkhorn_kernel(float *out, const float *mix, const float *scale, const float *base, uint32_t n_rows, uint32_t sinkhorn_iters, float epsv) {
uint32_t row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= n_rows) return;
hc4_split_one(out + (uint64_t)row * 24, mix + (uint64_t)row * 24, scale, base, sinkhorn_iters, epsv);
}
__global__ static void hc_weighted_sum_kernel(float *out, const float *x, const float *w, uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens, uint32_t weight_stride_f32) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_embd * n_tokens;
if (gid >= n) return;
uint32_t d = gid % n_embd;
uint32_t t = gid / n_embd;
float acc = 0.0f;
for (uint32_t h = 0; h < n_hc; h++) {
acc += x[(uint64_t)t * n_hc * n_embd + (uint64_t)h * n_embd + d] *
w[(uint64_t)t * weight_stride_f32 + h];
}
out[(uint64_t)t * n_embd + d] = acc;
}
__global__ static void hc_expand_kernel(
float *out_hc,
const float *block_out,
const float *block_add,
const float *block_add2,
const float *residual_hc,
const float *post,
const float *comb,
uint32_t n_embd,
uint32_t n_hc,
uint32_t n_tokens,
uint32_t post_stride,
uint32_t comb_stride,
int has_add,
int has_add2) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
if (gid >= n_elem) return;
uint32_t d = gid % n_embd;
uint64_t tmp = gid / n_embd;
uint32_t dst_hc = tmp % n_hc;
uint32_t t = tmp / n_hc;
float block_v = block_out[(uint64_t)t * n_embd + d];
if (has_add) {
float add_v = block_add[(uint64_t)t * n_embd + d];
if (has_add2) add_v += block_add2[(uint64_t)t * n_embd + d];
block_v += add_v;
}
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
acc += comb_v * res_v;
}
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
}
__global__ static void hc_split_weighted_sum_fused_kernel(
float *out,
float *split,
const float *mix,
const float *residual_hc,
const float *scale,
const float *base,
uint32_t n_embd,
uint32_t n_hc,
uint32_t n_rows,
uint32_t sinkhorn_iters,
float epsv) {
uint32_t t = blockIdx.x;
uint32_t d = threadIdx.x;
if (t >= n_rows || n_hc != 4) return;
const uint32_t mix_hc = 24;
float *sp = split + (uint64_t)t * mix_hc;
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
__syncthreads();
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
float acc = 0.0f;
for (uint32_t h = 0; h < 4; h++) {
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
}
out[(uint64_t)t * n_embd + col] = acc;
}
}
__global__ static void hc_split_weighted_sum_norm_fused_kernel(
float *out,
float *norm_out,
float *split,
const float *mix,
const float *residual_hc,
const float *scale,
const float *base,
const float *norm_w,
uint32_t n_embd,
uint32_t n_hc,
uint32_t n_rows,
uint32_t sinkhorn_iters,
float epsv,
float norm_eps) {
const uint32_t t = blockIdx.x;
const uint32_t d = threadIdx.x;
if (t >= n_rows || n_hc != 4) return;
const uint32_t mix_hc = 24;
float *sp = split + (uint64_t)t * mix_hc;
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
__syncthreads();
float sum = 0.0f;
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
float acc = 0.0f;
for (uint32_t h = 0; h < 4; h++) {
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
}
out[(uint64_t)t * n_embd + col] = acc;
sum += acc * acc;
}
__shared__ float partial[256];
partial[d] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (d < stride) partial[d] += partial[d + stride];
__syncthreads();
}
const float norm_scale = rsqrtf(partial[0] / (float)n_embd + norm_eps);
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
const float v = out[(uint64_t)t * n_embd + col];
norm_out[(uint64_t)t * n_embd + col] = v * norm_scale * norm_w[col];
}
}
__global__ static void output_hc_weights_kernel(
float *out,
const float *pre,
const float *scale,
const float *base,
uint32_t n_hc,
uint32_t n_tokens,
float epsv) {
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t n = n_tokens * n_hc;
if (gid >= n) return;
uint32_t h = gid % n_hc;
float z = pre[gid] * scale[0] + base[h];
out[gid] = 1.0f / (1.0f + expf(-z)) + epsv;
}
__global__ static void fill_f32_kernel(float *x, uint64_t n, float v) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) x[i] = v;
}
__global__ static void compressor_store_kernel(
const float *kv,
const float *sc,
float *state_kv,
float *state_score,
const void *model_map,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens) {
uint32_t coff = ratio == 4u ? 2u : 1u;
uint32_t width = coff * head_dim;
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * width;
if (gid >= n) return;
uint32_t t = gid / width;
uint32_t j = gid - (uint64_t)t * width;
uint32_t pos_mod = (pos0 + t) % ratio;
uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
state_kv[(uint64_t)dst_row * width + j] = kv[(uint64_t)t * width + j];
state_score[(uint64_t)dst_row * width + j] =
sc[(uint64_t)t * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)pos_mod * width + j);
}
__global__ static void compressor_set_rows_kernel(
float *state_kv,
float *state_score,
const float *kv,
const float *sc,
const void *model_map,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t width,
uint32_t ratio,
uint32_t pos0,
uint32_t src0,
uint32_t dst0,
uint32_t rows) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)rows * width;
if (gid >= n) return;
uint32_t r = gid / width;
uint32_t j = gid - (uint64_t)r * width;
uint32_t src = src0 + r;
uint32_t dst = dst0 + r;
uint32_t phase = (pos0 + src) % ratio;
state_kv[(uint64_t)dst * width + j] = kv[(uint64_t)src * width + j];
state_score[(uint64_t)dst * width + j] =
sc[(uint64_t)src * width + j] + model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)phase * width + j);
}
__global__ static void compressor_prefill_pool_kernel(
float *comp,
const float *kv,
const float *sc,
const float *state_kv,
const float *state_score,
const void *model_map,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_comp,
uint32_t replay) {
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
uint32_t c = blockIdx.y;
if (d >= head_dim || c >= n_comp) return;
uint32_t coff = ratio == 4u ? 2u : 1u;
uint32_t width = coff * head_dim;
float vals[128];
float scores[128];
float max_s = -INFINITY;
uint32_t n_cand = 0;
if (ratio == 4u) {
if (replay && c == 0) {
for (uint32_t r = 0; r < 4; r++) {
vals[n_cand] = state_kv[(uint64_t)r * width + d];
scores[n_cand] = state_score[(uint64_t)r * width + d];
max_s = fmaxf(max_s, scores[n_cand++]);
}
} else if (c > 0) {
uint32_t base = (c - 1u) * ratio;
for (uint32_t r = 0; r < 4; r++) {
uint32_t t = base + r;
float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d);
vals[n_cand] = kv[(uint64_t)t * width + d];
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
max_s = fmaxf(max_s, scores[n_cand++]);
}
}
uint32_t base = c * ratio;
for (uint32_t r = 0; r < 4; r++) {
uint32_t t = base + r;
float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + head_dim + d);
vals[n_cand] = kv[(uint64_t)t * width + head_dim + d];
scores[n_cand] = sc[(uint64_t)t * width + head_dim + d] + ape;
max_s = fmaxf(max_s, scores[n_cand++]);
}
} else {
uint32_t base = c * ratio;
for (uint32_t r = 0; r < ratio; r++) {
uint32_t t = base + r;
float ape = model_scalar_dev(model_map, ape_offset, ape_type, (uint64_t)((pos0 + t) % ratio) * width + d);
vals[n_cand] = kv[(uint64_t)t * width + d];
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
max_s = fmaxf(max_s, scores[n_cand++]);
}
}
float den = 0.0f, acc = 0.0f;
for (uint32_t i = 0; i < n_cand; i++) {
float w = expf(scores[i] - max_s);
den += w;
acc += vals[i] * w;
}
comp[(uint64_t)c * head_dim + d] = den != 0.0f ? acc / den : 0.0f;
}
__global__ static void compressor_update_pool_kernel(
float *row,
const float *state_kv,
const float *state_score,
uint32_t head_dim,
uint32_t ratio) {
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
if (d >= head_dim) return;
uint32_t coff = ratio == 4u ? 2u : 1u;
uint32_t width = coff * head_dim;
float vals[128];
float scores[128];
float max_s = -INFINITY;
uint32_t n_cand = 0;
if (ratio == 4u) {
for (uint32_t r = 0; r < 4; r++) {
vals[n_cand] = state_kv[(uint64_t)r * width + d];
scores[n_cand] = state_score[(uint64_t)r * width + d];
max_s = fmaxf(max_s, scores[n_cand++]);
}
for (uint32_t r = 0; r < 4; r++) {
vals[n_cand] = state_kv[(uint64_t)(ratio + r) * width + head_dim + d];
scores[n_cand] = state_score[(uint64_t)(ratio + r) * width + head_dim + d];
max_s = fmaxf(max_s, scores[n_cand++]);
}
} else {
for (uint32_t r = 0; r < ratio; r++) {
vals[n_cand] = state_kv[(uint64_t)r * width + d];
scores[n_cand] = state_score[(uint64_t)r * width + d];
max_s = fmaxf(max_s, scores[n_cand++]);
}
}
float den = 0.0f, acc = 0.0f;
for (uint32_t i = 0; i < n_cand; i++) {
float w = expf(scores[i] - max_s);
den += w;
acc += vals[i] * w;
}
row[d] = den != 0.0f ? acc / den : 0.0f;
}
__global__ static void compressor_shift_ratio4_kernel(float *state_kv, float *state_score, uint32_t width) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t half = 4ull * width;
if (i >= half) return;
float v = state_kv[half + i];
float s = state_score[half + i];
state_kv[i] = v;
state_score[i] = s;
state_kv[half + i] = v;
state_score[half + i] = s;
}
__device__ static float softplus_dev(float x) {
if (x > 20.0f) return x;
if (x < -20.0f) return expf(x);
return log1pf(expf(x));
}
__global__ static void router_select_kernel(
int32_t *selected,
float *weights,
float *probs,
const float *bias,
const int32_t *hash,
const float *logits,
const int32_t *tokens,
int32_t token_scalar,
uint32_t hash_rows,
uint32_t n_tokens,
int has_bias,
int hash_mode) {
uint32_t t = blockIdx.x;
if (t >= n_tokens || threadIdx.x != 0) return;
const float *log = logits + (uint64_t)t * 256;
float *prob = probs + (uint64_t)t * 256;
int32_t *sel = selected + (uint64_t)t * 6;
float *w = weights + (uint64_t)t * 6;
for (int i = 0; i < 256; i++) prob[i] = sqrtf(softplus_dev(log[i]));
if (hash_mode) {
int32_t tok = tokens ? tokens[t] : token_scalar;
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6;
for (int i = 0; i < 6; i++) sel[i] = row[i];
} else {
for (int i = 0; i < 6; i++) sel[i] = -1;
for (int i = 0; i < 256; i++) {
float score = prob[i] + (has_bias ? bias[i] : 0.0f);
for (int j = 0; j < 6; j++) {
if (sel[j] < 0 || score > prob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) {
for (int k = 5; k > j; k--) sel[k] = sel[k - 1];
sel[j] = i;
break;
}
}
}
}
float sum = 0.0f;
for (int i = 0; i < 6; i++) {
int e = sel[i];
float v = (e >= 0 && e < 256) ? prob[e] : 0.0f;
w[i] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
for (int i = 0; i < 6; i++) w[i] = w[i] / sum * 1.5f;
}
__global__ static void router_select_parallel_kernel(
int32_t *selected,
float *weights,
float *probs,
const float *bias,
const int32_t *hash,
const float *logits,
const int32_t *tokens,
int32_t token_scalar,
uint32_t hash_rows,
uint32_t n_tokens,
int has_bias,
int hash_mode) {
uint32_t t = blockIdx.x;
uint32_t i = threadIdx.x;
if (t >= n_tokens || i >= 256u) return;
const float *log = logits + (uint64_t)t * 256;
float *prob = probs + (uint64_t)t * 256;
int32_t *sel = selected + (uint64_t)t * 6;
float *w = weights + (uint64_t)t * 6;
__shared__ float sprob[256];
const float p = sqrtf(softplus_dev(log[i]));
sprob[i] = p;
prob[i] = p;
__syncthreads();
if (i != 0) return;
if (hash_mode) {
int32_t tok = tokens ? tokens[t] : token_scalar;
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6;
for (int j = 0; j < 6; j++) sel[j] = row[j];
} else {
for (int j = 0; j < 6; j++) sel[j] = -1;
for (int e = 0; e < 256; e++) {
float score = sprob[e] + (has_bias ? bias[e] : 0.0f);
for (int j = 0; j < 6; j++) {
if (sel[j] < 0 || score > sprob[sel[j]] + (has_bias ? bias[sel[j]] : 0.0f)) {
for (int k = 5; k > j; k--) sel[k] = sel[k - 1];
sel[j] = e;
break;
}
}
}
}
float sum = 0.0f;
for (int j = 0; j < 6; j++) {
int e = sel[j];
float v = (e >= 0 && e < 256) ? sprob[e] : 0.0f;
w[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
for (int j = 0; j < 6; j++) w[j] = w[j] / sum * 1.5f;
}
__global__ static void router_select_visual_parallel_kernel(
int32_t *selected,
float *weights,
float *probs,
const float *bias,
const float *visual_bias,
const int32_t *hash,
const float *logits,
const int32_t *tokens,
uint32_t hash_rows,
uint32_t vocab_size,
uint32_t n_tokens,
int has_bias,
int hash_mode) {
const uint32_t t = blockIdx.x;
const uint32_t i = threadIdx.x;
if (t >= n_tokens || i >= 256u) return;
const float *log = logits + (uint64_t)t * 256u;
float *prob = probs + (uint64_t)t * 256u;
int32_t *sel = selected + (uint64_t)t * 6u;
float *w = weights + (uint64_t)t * 6u;
__shared__ float sprob[256];
const float p = sqrtf(softplus_dev(log[i]));
sprob[i] = p;
prob[i] = p;
__syncthreads();
if (i != 0u) return;
const int32_t token = tokens[t];
const bool image = token >= 0 && (uint32_t)token >= vocab_size;
if (hash_mode && !image) {
const uint32_t row = token >= 0 && (uint32_t)token < hash_rows
? (uint32_t)token : 0u;
const int32_t *src = hash + (uint64_t)row * 6u;
for (uint32_t j = 0; j < 6u; j++) sel[j] = src[j];
} else {
for (uint32_t j = 0; j < 6u; j++) sel[j] = -1;
for (uint32_t e = 0; e < 256u; e++) {
const float score = sprob[e] +
(image ? visual_bias[e] : (has_bias ? bias[e] : 0.0f));
for (uint32_t j = 0; j < 6u; j++) {
const int32_t old = sel[j];
const float old_score = old < 0 ? -INFINITY :
sprob[old] + (image ? visual_bias[old]
: (has_bias ? bias[old] : 0.0f));
if (old < 0 || score > old_score) {
for (uint32_t k = 5u; k > j; k--) sel[k] = sel[k - 1u];
sel[j] = (int32_t)e;
break;
}
}
}
}
float sum = 0.0f;
for (uint32_t j = 0; j < 6u; j++) {
const int32_t e = sel[j];
const float v = e >= 0 && e < 256 ? sprob[e] : 0.0f;
w[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
}
__device__ __forceinline__ static bool router_score_better(float av, uint32_t ai, float bv, uint32_t bi) {
return av > bv || (av == bv && ai < bi);
}
template<uint32_t EXPERTS = 256>
__global__ static void router_select_warp_topk_kernel(
int32_t *selected,
float *weights,
float *probs,
const float *bias,
const int32_t *hash,
const float *logits,
const int32_t *tokens,
int32_t token_scalar,
uint32_t hash_rows,
uint32_t n_tokens,
int has_bias,
int hash_mode) {
const uint32_t lane = threadIdx.x;
const uint32_t row_in_block = threadIdx.y;
const uint32_t t = blockIdx.x * blockDim.y + row_in_block;
if (t >= n_tokens || lane >= 32u) return;
const float *log = logits + (uint64_t)t * EXPERTS;
float *prob = probs + (uint64_t)t * EXPERTS;
int32_t *sel = selected + (uint64_t)t * 6u;
float *w = weights + (uint64_t)t * 6u;
__shared__ float sprob[4][EXPERTS];
float local_prob[EXPERTS / 32];
float local_score[EXPERTS / 32];
#pragma unroll
for (uint32_t j = 0; j < EXPERTS / 32; j++) {
const uint32_t e = lane + j * 32u;
const float p = sqrtf(softplus_dev(log[e]));
local_prob[j] = p;
local_score[j] = p + (has_bias ? bias[e] : 0.0f);
sprob[row_in_block][e] = p;
prob[e] = p;
}
__syncwarp();
if (hash_mode) {
if (lane == 0) {
int32_t tok = tokens ? tokens[t] : token_scalar;
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
const int32_t *row = hash + (uint64_t)tok * 6u;
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
const int32_t e = row[j];
sel[j] = e;
const float v = (e >= 0 && (uint32_t)e < EXPERTS) ? sprob[row_in_block][(uint32_t)e] : 0.0f;
w[j] = v;
sum += v;
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
}
return;
}
float out_prob[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
uint32_t out_idx[6] = {0, 0, 0, 0, 0, 0};
#pragma unroll
for (uint32_t k = 0; k < 6u; k++) {
float best_score = -INFINITY;
float best_prob = 0.0f;
uint32_t best_idx = UINT32_MAX;
#pragma unroll
for (uint32_t j = 0; j < EXPERTS / 32; j++) {
const uint32_t e = lane + j * 32u;
const float s = local_score[j];
if (router_score_better(s, e, best_score, best_idx)) {
best_score = s;
best_prob = local_prob[j];
best_idx = e;
}
}
#pragma unroll
for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
const float other_score = __shfl_xor_sync(0xffffffffu, best_score, mask);
const float other_prob = __shfl_xor_sync(0xffffffffu, best_prob, mask);
const uint32_t other_idx = __shfl_xor_sync(0xffffffffu, best_idx, mask);
if (router_score_better(other_score, other_idx, best_score, best_idx)) {
best_score = other_score;
best_prob = other_prob;
best_idx = other_idx;
}
}
#pragma unroll
for (uint32_t j = 0; j < EXPERTS / 32; j++) {
const uint32_t e = lane + j * 32u;
if (e == best_idx) local_score[j] = -INFINITY;
}
if (lane == 0) {
out_idx[k] = best_idx;
out_prob[k] = best_prob;
}
}
if (lane == 0) {
float sum = 0.0f;
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) {
sel[j] = (int32_t)out_idx[j];
w[j] = out_prob[j];
sum += out_prob[j];
}
sum = fmaxf(sum, 6.103515625e-5f);
#pragma unroll
for (uint32_t j = 0; j < 6u; j++) w[j] = w[j] / sum * 1.5f;
}
}
__global__ static void swiglu_kernel(float *out, const float *gate, const float *up, uint32_t n, float clamp, float weight) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
float g = gate[i];
float u = up[i];
if (clamp > 1.0e-6f) {
g = fminf(g, clamp);
u = fminf(fmaxf(u, -clamp), clamp);
}
float s = g / (1.0f + expf(-g));
out[i] = s * u * weight;
}
__global__ static void add_kernel(float *out, const float *a, const float *b, uint32_t n) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
out[i] = a[i] + b[i];
}
__global__ static void directional_steering_project_kernel(
float *x,
const float *directions,
uint32_t layer,
uint32_t width,
uint32_t rows,
float scale) {
const uint32_t row = blockIdx.x;
if (row >= rows || width == 0) return;
float *xr = x + (uint64_t)row * width;
const float *dir = directions + (uint64_t)layer * width;
float sum = 0.0f;
for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) {
sum += xr[i] * dir[i];
}
__shared__ float partial[256];
partial[threadIdx.x] = sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
const float coeff = scale * partial[0];
for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) {
xr[i] -= coeff * dir[i];
}
}
__global__ static void zero_kernel(float *out, uint64_t n) {
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = 0.0f;
}
__global__ static void indexer_scores_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
int causal) {
uint32_t c = blockIdx.x;
uint32_t t = blockIdx.y;
if (c >= n_comp || t >= n_tokens) return;
if (causal) {
uint32_t n_visible = (pos0 + t + 1u) / ratio;
if (c >= n_visible) {
if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = -INFINITY;
return;
}
}
float total = 0.0f;
for (uint32_t h = 0; h < n_head; h++) {
const float *qh = q + ((uint64_t)t * n_head + h) * head_dim;
const float *kh = index_comp + (uint64_t)c * head_dim;
float dot = 0.0f;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) dot += qh[d] * kh[d];
__shared__ float partial[256];
partial[threadIdx.x] = dot;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
total += fmaxf(partial[0], 0.0f) * weights[(uint64_t)t * n_head + h];
__syncthreads();
}
if (threadIdx.x == 0) scores[(uint64_t)t * n_comp + c] = total * scale;
}
__global__ static void indexer_score_one_direct_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t pos0,
uint32_t ratio,
float scale,
int causal) {
const uint32_t c = blockIdx.x;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
if (c >= n_comp || tid >= 128u) return;
if (causal) {
const uint32_t visible = ratio ? (pos0 + 1u) / ratio : n_comp;
if (c >= visible) {
if (tid == 0) scores[c] = -INFINITY;
return;
}
}
__shared__ float krow[128];
__shared__ float partial[4];
if (tid < 128u) krow[tid] = index_comp[(uint64_t)c * 128u + tid];
__syncthreads();
float total = 0.0f;
for (uint32_t h0 = 0; h0 < 64u; h0 += 4u) {
const uint32_t h = h0 + warp;
const float4 qv = ((const float4 *)(q + (uint64_t)h * 128u))[lane];
const float4 kv = ((const float4 *)krow)[lane];
float dot = qv.x * kv.x + qv.y * kv.y + qv.z * kv.z + qv.w * kv.w;
dot = warp_sum_f32(dot);
if (lane == 0) partial[warp] = fmaxf(dot, 0.0f) * weights[h] * scale;
__syncthreads();
if (tid == 0) total += partial[0] + partial[1] + partial[2] + partial[3];
__syncthreads();
}
if (tid == 0) scores[c] = total;
}
__global__ static void indexer_scores_wmma_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
int causal) {
#if __CUDA_ARCH__ >= 700
namespace wmma = nvcuda::wmma;
const uint32_t tile_c = blockIdx.x * 16u;
const uint32_t tile_t = blockIdx.y * 16u;
const uint32_t tid = threadIdx.x;
if (tid >= 32u || head_dim != 128u) return;
if (causal) {
const uint32_t last_token = min(tile_t + 16u, n_tokens);
const uint32_t max_visible = last_token > tile_t
? min((pos0 + last_token) / ratio, n_comp)
: 0u;
if (tile_c >= max_visible) {
for (uint32_t i = tid; i < 16u * 16u; i += 32u) {
const uint32_t r = i >> 4u;
const uint32_t c = i & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + c;
if (token < n_tokens && comp < n_comp) {
scores[(uint64_t)token * n_comp + comp] = -INFINITY;
}
}
return;
}
}
__shared__ __half a_sh[16 * 128];
__shared__ __half b_sh[16 * 128];
__shared__ float c_sh[16 * 16];
__shared__ float acc_sh[16 * 16];
for (uint32_t i = tid; i < 16u * 16u; i += 32u) acc_sh[i] = 0.0f;
for (uint32_t i = tid; i < 16u * 128u; i += 32u) {
const uint32_t c = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t comp = tile_c + c;
float v = 0.0f;
if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d];
b_sh[d + c * 128u] = __float2half(v);
}
__syncthreads();
for (uint32_t h = 0; h < n_head; h++) {
for (uint32_t i = tid; i < 16u * 128u; i += 32u) {
const uint32_t r = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t token = tile_t + r;
float v = 0.0f;
if (token < n_tokens) {
v = q[((uint64_t)token * n_head + h) * head_dim + d];
}
a_sh[i] = __float2half(v);
}
__syncthreads();
wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
wmma::fill_fragment(c_frag, 0.0f);
for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) {
wmma::load_matrix_sync(a_frag, a_sh + k0, 128);
wmma::load_matrix_sync(b_frag, b_sh + k0, 128);
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
}
wmma::store_matrix_sync(c_sh, c_frag, 16, wmma::mem_row_major);
__syncthreads();
for (uint32_t i = tid; i < 16u * 16u; i += 32u) {
const uint32_t r = i >> 4u;
const uint32_t token = tile_t + r;
if (token < n_tokens) {
const float w = weights[(uint64_t)token * n_head + h];
acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w;
}
}
__syncthreads();
}
for (uint32_t i = tid; i < 16u * 16u; i += 32u) {
const uint32_t r = i >> 4u;
const uint32_t c = i & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + c;
if (token < n_tokens && comp < n_comp) {
float out = acc_sh[i] * scale;
if (causal) {
const uint32_t visible = (pos0 + token + 1u) / ratio;
if (comp >= visible) out = -INFINITY;
}
scores[(uint64_t)token * n_comp + comp] = out;
}
}
#endif
}
__global__ static void indexer_scores_wmma32_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
int causal) {
#if __CUDA_ARCH__ >= 700
namespace wmma = nvcuda::wmma;
const uint32_t tile_c = blockIdx.x * 32u;
const uint32_t tile_t = blockIdx.y * 16u;
const uint32_t tid = threadIdx.x;
const uint32_t warp = tid >> 5u;
if (tid >= 64u || head_dim != 128u) return;
if (causal) {
const uint32_t last_token = min(tile_t + 16u, n_tokens);
const uint32_t max_visible = last_token > tile_t
? min((pos0 + last_token) / ratio, n_comp)
: 0u;
if (tile_c >= max_visible) {
for (uint32_t i = tid; i < 16u * 32u; i += 64u) {
const uint32_t r = i >> 5u;
const uint32_t c = i & 31u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + c;
if (token < n_tokens && comp < n_comp) {
scores[(uint64_t)token * n_comp + comp] = -INFINITY;
}
}
return;
}
}
__shared__ __half a_sh[16 * 128];
__shared__ __half b_sh[32 * 128];
__shared__ float c_sh[2 * 16 * 16];
__shared__ float acc_sh[2 * 16 * 16];
for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) acc_sh[i] = 0.0f;
for (uint32_t i = tid; i < 32u * 128u; i += 64u) {
const uint32_t c = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t comp = tile_c + c;
float v = 0.0f;
if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d];
b_sh[d + c * 128u] = __float2half(v);
}
__syncthreads();
for (uint32_t h = 0; h < n_head; h++) {
for (uint32_t i = tid; i < 16u * 128u; i += 64u) {
const uint32_t r = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t token = tile_t + r;
float v = 0.0f;
if (token < n_tokens) {
v = q[((uint64_t)token * n_head + h) * head_dim + d];
}
a_sh[i] = __float2half(v);
}
__syncthreads();
wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
wmma::fill_fragment(c_frag, 0.0f);
const uint32_t col0 = warp * 16u;
for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) {
wmma::load_matrix_sync(a_frag, a_sh + k0, 128);
wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128);
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
}
wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major);
__syncthreads();
for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) {
const uint32_t wtile = i >> 8u;
const uint32_t local = i & 255u;
const uint32_t r = local >> 4u;
const uint32_t c = local & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + wtile * 16u + c;
if (token < n_tokens && comp < n_comp) {
const float w = weights[(uint64_t)token * n_head + h];
acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w;
}
}
__syncthreads();
}
for (uint32_t i = tid; i < 2u * 16u * 16u; i += 64u) {
const uint32_t wtile = i >> 8u;
const uint32_t local = i & 255u;
const uint32_t r = local >> 4u;
const uint32_t c = local & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + wtile * 16u + c;
if (token < n_tokens && comp < n_comp) {
float out = acc_sh[i] * scale;
if (causal) {
const uint32_t visible = (pos0 + token + 1u) / ratio;
if (comp >= visible) out = -INFINITY;
}
scores[(uint64_t)token * n_comp + comp] = out;
}
}
#endif
}
__global__ static void indexer_scores_wmma64_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
int causal) {
#if __CUDA_ARCH__ >= 700
namespace wmma = nvcuda::wmma;
const uint32_t tile_c = blockIdx.x * 64u;
const uint32_t tile_t = blockIdx.y * 16u;
const uint32_t tid = threadIdx.x;
const uint32_t warp = tid >> 5u;
if (tid >= 128u || head_dim != 128u) return;
if (causal) {
const uint32_t last_token = min(tile_t + 16u, n_tokens);
const uint32_t max_visible = last_token > tile_t
? min((pos0 + last_token) / ratio, n_comp)
: 0u;
if (tile_c >= max_visible) {
for (uint32_t i = tid; i < 16u * 64u; i += 128u) {
const uint32_t r = i >> 6u;
const uint32_t c = i & 63u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + c;
if (token < n_tokens && comp < n_comp) {
scores[(uint64_t)token * n_comp + comp] = -INFINITY;
}
}
return;
}
}
__shared__ __half a_sh[16 * 128];
__shared__ __half b_sh[64 * 128];
__shared__ float c_sh[4 * 16 * 16];
__shared__ float acc_sh[4 * 16 * 16];
for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) acc_sh[i] = 0.0f;
for (uint32_t i = tid; i < 64u * 128u; i += 128u) {
const uint32_t c = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t comp = tile_c + c;
float v = 0.0f;
if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d];
b_sh[d + c * 128u] = __float2half(v);
}
__syncthreads();
for (uint32_t h = 0; h < n_head; h++) {
for (uint32_t i = tid; i < 16u * 128u; i += 128u) {
const uint32_t r = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t token = tile_t + r;
float v = 0.0f;
if (token < n_tokens) {
v = q[((uint64_t)token * n_head + h) * head_dim + d];
}
a_sh[i] = __float2half(v);
}
__syncthreads();
wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;
wmma::fill_fragment(c_frag, 0.0f);
const uint32_t col0 = warp * 16u;
for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) {
wmma::load_matrix_sync(a_frag, a_sh + k0, 128);
wmma::load_matrix_sync(b_frag, b_sh + col0 * 128u + k0, 128);
wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
}
wmma::store_matrix_sync(c_sh + warp * 16u * 16u, c_frag, 16, wmma::mem_row_major);
__syncthreads();
for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) {
const uint32_t wtile = i >> 8u;
const uint32_t local = i & 255u;
const uint32_t r = local >> 4u;
const uint32_t c = local & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + wtile * 16u + c;
if (token < n_tokens && comp < n_comp) {
const float w = weights[(uint64_t)token * n_head + h];
acc_sh[i] += fmaxf(c_sh[i], 0.0f) * w;
}
}
__syncthreads();
}
for (uint32_t i = tid; i < 4u * 16u * 16u; i += 128u) {
const uint32_t wtile = i >> 8u;
const uint32_t local = i & 255u;
const uint32_t r = local >> 4u;
const uint32_t c = local & 15u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + wtile * 16u + c;
if (token < n_tokens && comp < n_comp) {
float out = acc_sh[i] * scale;
if (causal) {
const uint32_t visible = (pos0 + token + 1u) / ratio;
if (comp >= visible) out = -INFINITY;
}
scores[(uint64_t)token * n_comp + comp] = out;
}
}
#endif
}
__global__ static void indexer_scores_wmma128_kernel(
float *scores,
const float *q,
const float *weights,
const float *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
int causal) {
#if __CUDA_ARCH__ >= 700
namespace wmma = nvcuda::wmma;
const uint32_t tile_c = blockIdx.x * 128u;
const uint32_t tile_t = blockIdx.y * 32u;
const uint32_t tid = threadIdx.x;
const uint32_t warp = tid >> 5u;
if (tid >= 256u || head_dim != 128u) return;
if (causal) {
const uint32_t last_token = min(tile_t + 32u, n_tokens);
const uint32_t max_visible = last_token > tile_t
? min((pos0 + last_token) / ratio, n_comp)
: 0u;
if (tile_c >= max_visible) {
for (uint32_t i = tid; i < 32u * 128u; i += 256u) {
const uint32_t r = i >> 7u;
const uint32_t c = i & 127u;
const uint32_t token = tile_t + r;
const uint32_t comp = tile_c + c;
if (token < n_tokens && comp < n_comp) {
scores[(uint64_t)token * n_comp + comp] = -INFINITY;
}
}
return;
}
}
/* The padded stride avoids ldmatrix bank conflicts. Two token tiles
* share each staged index fragment, while WMMA accumulators remain in
* registers through the head reduction. */
__shared__ __half a_sh[32 * 136];
__shared__ __half b_sh[128 * 136];
const uint32_t lane = tid & 31u;
const uint32_t quad = lane >> 2u;
const uint32_t tcol = (lane & 3u) * 2u;
float acc0[8], acc1[8];
#pragma unroll
for (uint32_t i = 0; i < 8u; i++) {
acc0[i] = 0.0f;
acc1[i] = 0.0f;
}
for (uint32_t i = tid; i < 128u * 128u; i += 256u) {
const uint32_t c = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t comp = tile_c + c;
float v = 0.0f;
if (comp < n_comp) v = index_comp[(uint64_t)comp * head_dim + d];
b_sh[d + c * 136u] = __float2half(v);
}
__syncthreads();
const uint32_t t0_lo = tile_t + quad;
const uint32_t t0_hi = tile_t + quad + 8u;
const uint32_t t1_lo = tile_t + 16u + quad;
const uint32_t t1_hi = tile_t + 24u + quad;
for (uint32_t h = 0; h < n_head; h++) {
for (uint32_t i = tid; i < 32u * 128u; i += 256u) {
const uint32_t r = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t token = tile_t + r;
float v = 0.0f;
if (token < n_tokens) {
v = q[((uint64_t)token * n_head + h) * head_dim + d];
}
a_sh[r * 136u + d] = __float2half(v);
}
__syncthreads();
wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a0, a1;
wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> c0, c1;
wmma::fill_fragment(c0, 0.0f);
wmma::fill_fragment(c1, 0.0f);
const uint32_t col0 = warp * 16u;
for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) {
wmma::load_matrix_sync(a0, a_sh + k0, 136);
wmma::load_matrix_sync(a1, a_sh + 16u * 136u + k0, 136);
wmma::load_matrix_sync(b_frag, b_sh + col0 * 136u + k0, 136);
wmma::mma_sync(c0, a0, b_frag, c0);
wmma::mma_sync(c1, a1, b_frag, c1);
}
const float w0_lo = t0_lo < n_tokens ? weights[(uint64_t)t0_lo * n_head + h] : 0.0f;
const float w0_hi = t0_hi < n_tokens ? weights[(uint64_t)t0_hi * n_head + h] : 0.0f;
const float w1_lo = t1_lo < n_tokens ? weights[(uint64_t)t1_lo * n_head + h] : 0.0f;
const float w1_hi = t1_hi < n_tokens ? weights[(uint64_t)t1_hi * n_head + h] : 0.0f;
#pragma unroll
for (int i = 0; i < 8; i++) {
acc0[i] += fmaxf(c0.x[i], 0.0f) * ((i & 2) ? w0_hi : w0_lo);
acc1[i] += fmaxf(c1.x[i], 0.0f) * ((i & 2) ? w1_hi : w1_lo);
}
__syncthreads();
}
#pragma unroll
for (int i = 0; i < 8; i++) {
const uint32_t row = quad + ((i & 2) ? 8u : 0u);
const uint32_t col = tcol + (i & 1) + ((i & 4) ? 8u : 0u);
const uint32_t comp = tile_c + warp * 16u + col;
for (uint32_t rt = 0; rt < 2u; rt++) {
const uint32_t token = tile_t + rt * 16u + row;
if (token < n_tokens && comp < n_comp) {
float out = (rt ? acc1[i] : acc0[i]) * scale;
if (causal) {
const uint32_t visible = (pos0 + token + 1u) / ratio;
if (comp >= visible) out = -INFINITY;
}
scores[(uint64_t)token * n_comp + comp] = out;
}
}
}
#endif
}
#ifdef DS4_CUDA_HAVE_MXF4
#define DS4_MXF4_INDEXER_HEADS 64u
#define DS4_MXF4_INDEXER_DIM 128u
__device__ static inline uint32_t indexer_mxf4_nearest_level(float v) {
if (v <= 0.25f) return 0u;
if (v <= 0.75f) return 1u;
if (v <= 1.25f) return 2u;
if (v <= 1.75f) return 3u;
if (v <= 2.5f) return 4u;
if (v <= 3.5f) return 5u;
if (v <= 5.0f) return 6u;
return 7u;
}
/* Encode rows that have already passed the model's E2M1 QAT. The packed
* exponent word is consumed directly by the SM121 block-scaled MMA. */
__global__ static void indexer_mxf4_encode_rows_kernel(
const float * __restrict__ src,
unsigned char * __restrict__ codes,
uint32_t * __restrict__ scales,
uint32_t rows) {
const uint32_t row =
(blockIdx.x * blockDim.x + threadIdx.x) >> 5u;
const uint32_t lane = threadIdx.x & 31u;
if (row >= rows) return;
const float *xr = src + (uint64_t)row * DS4_MXF4_INDEXER_DIM;
float values[4];
float amax = 0.0f;
#pragma unroll
for (int i = 0; i < 4; i++) {
values[i] = xr[lane * 4u + i];
amax = fmaxf(amax, fabsf(values[i]));
}
#pragma unroll
for (uint32_t off = 4u; off != 0u; off >>= 1u) {
amax = fmaxf(amax,
__shfl_xor_sync(0xffffffffu, amax, off, 8));
}
uint32_t exponent = 127u;
if (amax > 0.0f) {
const uint32_t bits = __float_as_uint(amax / 6.0f);
exponent = (bits >> 23u) & 0xffu;
if ((bits & 0x7fffffu) != 0u) exponent++;
exponent = min(max(exponent, 1u), 254u);
}
const float reciprocal =
__uint_as_float((254u - exponent) << 23u);
uint32_t packed = 0u;
#pragma unroll
for (int i = 0; i < 4; i++) {
const uint32_t level =
indexer_mxf4_nearest_level(fabsf(values[i]) * reciprocal);
packed |= (level | (values[i] < 0.0f ? 8u : 0u)) << (4 * i);
}
((uint16_t *)codes)[(uint64_t)row * 32u + lane] = (uint16_t)packed;
const uint32_t e1 = __shfl_sync(0xffffffffu, exponent, 8);
const uint32_t e2 = __shfl_sync(0xffffffffu, exponent, 16);
const uint32_t e3 = __shfl_sync(0xffffffffu, exponent, 24);
if (lane == 0u) {
scales[row] = exponent | (e1 << 8u) | (e2 << 16u) | (e3 << 24u);
}
}
__device__ static inline void indexer_mxf4_mma(
float out[4],
const uint32_t a[4],
uint32_t b0,
uint32_t b1,
uint32_t scale_a,
uint32_t scale_b) {
asm volatile(
"mma.sync.aligned.kind::mxf4.block_scale.scale_vec::2X"
".m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3}, "
"{%10}, {0, 0}, {%11}, {0, 0};\n"
: "+f"(out[0]), "+f"(out[1]), "+f"(out[2]), "+f"(out[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
"r"(b0), "r"(b1), "r"(scale_a), "r"(scale_b));
}
__device__ __forceinline__ static void indexer_cp_async_u32(
uint32_t *dst,
const void *src,
bool valid) {
const unsigned smem = (unsigned)__cvta_generic_to_shared(dst);
asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n"
:: "r"(smem), "l"(src), "r"(valid ? 4u : 0u));
}
__device__ __forceinline__ static void indexer_cp_commit(void) {
asm volatile("cp.async.commit_group;\n");
}
__device__ __forceinline__ static void indexer_cp_wait(void) {
asm volatile("cp.async.wait_group 0;\n");
}
/* CUDA SM121 implementation adapted from entrpi/batched-serving's v3
* scorer and Marco Palaferri's original 16-token staging geometry. */
__global__ static void __launch_bounds__(512) indexer_scores_mxf4_kernel(
float *scores,
const unsigned char * __restrict__ q_codes,
const uint32_t * __restrict__ q_scales,
const float * __restrict__ weights,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t ratio,
float scale,
const unsigned char * __restrict__ k_codes,
const uint32_t * __restrict__ k_scales) {
constexpr uint32_t HEAD_GROUP = 16u;
const uint32_t tid = threadIdx.x;
const uint32_t warp = tid >> 5u;
const uint32_t lane = tid & 31u;
const uint32_t lane_in_group = lane & 3u;
const uint32_t tile_t = blockIdx.y * 16u;
const uint32_t tile_c = blockIdx.x * 256u;
const uint32_t last_token = min(tile_t + 16u, n_tokens);
const uint32_t max_visible = last_token > tile_t
? min((pos0 + last_token) / ratio, n_comp) : 0u;
if (tile_c >= max_visible) {
for (uint32_t i = tid; i < 16u * 256u; i += 512u) {
const uint32_t row = i >> 8u;
const uint32_t col = i & 255u;
const uint32_t token = tile_t + row;
const uint32_t comp = tile_c + col;
if (token < n_tokens && comp < n_comp) {
scores[(uint64_t)token * n_comp + comp] = -INFINITY;
}
}
return;
}
__shared__ __align__(16) uint32_t q_shared[2][16][8][32];
__shared__ uint32_t scale_shared[2][16][16];
__shared__ float weight_shared[2][16][16];
uint32_t b_frag[2][4];
uint32_t scale_b[2][2];
#pragma unroll
for (uint32_t cg = 0; cg < 2u; cg++) {
const uint32_t col =
tile_c + (cg * 16u + warp) * 8u + (lane >> 2u);
if (col < n_comp) {
const uint32_t *words =
(const uint32_t *)(k_codes + (uint64_t)col * 64u);
b_frag[cg][0] = words[lane_in_group];
b_frag[cg][1] = words[4u + lane_in_group];
b_frag[cg][2] = words[8u + lane_in_group];
b_frag[cg][3] = words[12u + lane_in_group];
const uint32_t packed_scale = k_scales[col];
scale_b[cg][0] = packed_scale & 0xffffu;
scale_b[cg][1] = packed_scale >> 16u;
} else {
b_frag[cg][0] = b_frag[cg][1] = 0u;
b_frag[cg][2] = b_frag[cg][3] = 0u;
scale_b[cg][0] = scale_b[cg][1] = 127u | (127u << 8u);
}
}
auto stage = [&](uint32_t head0, uint32_t buffer) {
for (uint32_t i = tid; i < HEAD_GROUP * 16u * 17u; i += 512u) {
const uint32_t hh = i / (16u * 17u);
const uint32_t rem = i - hh * (16u * 17u);
const uint32_t row = rem / 17u;
const uint32_t word = rem - row * 17u;
const uint32_t token = tile_t + row;
const bool valid = token < n_tokens;
const uint64_t qrow =
(uint64_t)token * DS4_MXF4_INDEXER_HEADS + head0 + hh;
const void *src;
uint32_t *dst;
if (word < 16u) {
src = valid
? (const void *)(q_codes + qrow * 64u + word * 4u)
: (const void *)q_codes;
dst = &q_shared[buffer][hh][row >> 1u]
[(row & 1u) * 16u + (word & 3u) * 4u + (word >> 2u)];
} else {
src = valid ? (const void *)(q_scales + qrow)
: (const void *)q_scales;
dst = &scale_shared[buffer][hh][row];
}
indexer_cp_async_u32(dst, src, valid);
}
for (uint32_t i = tid; i < HEAD_GROUP * 16u; i += 512u) {
const uint32_t hh = i >> 4u;
const uint32_t row = i & 15u;
const uint32_t token = tile_t + row;
const bool valid = token < n_tokens;
const void *src = valid
? (const void *)(weights +
(uint64_t)token * DS4_MXF4_INDEXER_HEADS + head0 + hh)
: (const void *)weights;
indexer_cp_async_u32(
(uint32_t *)&weight_shared[buffer][hh][row], src, valid);
}
};
stage(0u, 0u);
indexer_cp_commit();
indexer_cp_wait();
__syncthreads();
float accum[2][4];
#pragma unroll
for (uint32_t cg = 0; cg < 2u; cg++) {
accum[cg][0] = accum[cg][1] = 0.0f;
accum[cg][2] = accum[cg][3] = 0.0f;
}
const uint32_t row0 = lane >> 2u;
const uint32_t scale_row = (lane & 1u) * 8u + row0;
for (uint32_t group = 0;
group < DS4_MXF4_INDEXER_HEADS / HEAD_GROUP;
group++) {
const uint32_t current = group & 1u;
if (group + 1u < DS4_MXF4_INDEXER_HEADS / HEAD_GROUP) {
stage((group + 1u) * HEAD_GROUP, current ^ 1u);
indexer_cp_commit();
}
#pragma unroll
for (uint32_t hh = 0; hh < HEAD_GROUP; hh++) {
const uint4 a0 = *reinterpret_cast<const uint4 *>(
&q_shared[current][hh][row0 >> 1u]
[(row0 & 1u) * 16u + lane_in_group * 4u]);
const uint4 a1 = *reinterpret_cast<const uint4 *>(
&q_shared[current][hh][(row0 + 8u) >> 1u]
[(row0 & 1u) * 16u + lane_in_group * 4u]);
const uint32_t packed_scale =
scale_shared[current][hh][scale_row];
const uint32_t scale_a0 = packed_scale & 0xffffu;
const uint32_t scale_a1 = packed_scale >> 16u;
const float w0 = weight_shared[current][hh][row0];
const float w1 = weight_shared[current][hh][row0 + 8u];
const uint32_t a_lo[4] = {a0.x, a1.x, a0.y, a1.y};
const uint32_t a_hi[4] = {a0.z, a1.z, a0.w, a1.w};
#pragma unroll
for (uint32_t cg = 0; cg < 2u; cg++) {
float product[4] = {0.0f, 0.0f, 0.0f, 0.0f};
indexer_mxf4_mma(product, a_lo,
b_frag[cg][0], b_frag[cg][1],
scale_a0, scale_b[cg][0]);
indexer_mxf4_mma(product, a_hi,
b_frag[cg][2], b_frag[cg][3],
scale_a1, scale_b[cg][1]);
accum[cg][0] += fmaxf(product[0], 0.0f) * w0;
accum[cg][1] += fmaxf(product[1], 0.0f) * w0;
accum[cg][2] += fmaxf(product[2], 0.0f) * w1;
accum[cg][3] += fmaxf(product[3], 0.0f) * w1;
}
}
if (group + 1u < DS4_MXF4_INDEXER_HEADS / HEAD_GROUP) {
indexer_cp_wait();
__syncthreads();
}
}
const uint32_t token0 = tile_t + row0;
const uint32_t token1 = token0 + 8u;
const uint32_t visible0 = (pos0 + token0 + 1u) / ratio;
const uint32_t visible1 = (pos0 + token1 + 1u) / ratio;
#pragma unroll
for (uint32_t cg = 0; cg < 2u; cg++) {
const uint32_t col0 = tile_c +
(cg * 16u + warp) * 8u + lane_in_group * 2u;
#pragma unroll
for (uint32_t j = 0; j < 2u; j++) {
const uint32_t col = col0 + j;
if (col >= n_comp) continue;
if (token0 < n_tokens) {
scores[(uint64_t)token0 * n_comp + col] =
col >= visible0 ? -INFINITY : accum[cg][j] * scale;
}
if (token1 < n_tokens) {
scores[(uint64_t)token1 * n_comp + col] =
col >= visible1 ? -INFINITY : accum[cg][2u + j] * scale;
}
}
}
}
#endif
#ifdef DS4_CUDA_HAVE_MXF4
static void *indexer_mxf4_scratch_alloc(uint64_t bytes) {
if (g_indexer_mxf4_scratch_bytes >= bytes) {
return g_indexer_mxf4_scratch;
}
int device = -1;
if (!cuda_ok(cudaGetDevice(&device),
"get device for indexer mxf4 scratch")) {
return NULL;
}
if (g_indexer_mxf4_scratch) {
if (!cuda_ok(cudaDeviceSynchronize(),
"synchronize indexer mxf4 scratch growth")) {
return NULL;
}
(void)cudaFree(g_indexer_mxf4_scratch);
g_indexer_mxf4_scratch = NULL;
g_indexer_mxf4_scratch_bytes = 0;
g_indexer_mxf4_scratch_device = -1;
}
if (!cuda_ok(cudaMalloc(&g_indexer_mxf4_scratch, (size_t)bytes),
"allocate indexer mxf4 scratch")) {
g_indexer_mxf4_scratch = NULL;
return NULL;
}
g_indexer_mxf4_scratch_bytes = bytes;
g_indexer_mxf4_scratch_device = device;
return g_indexer_mxf4_scratch;
}
#endif
/* Return 1 when the SM121 path handled the call, 0 to use the established
* scorer, and -1 on a launch failure. */
static int indexer_scores_mxf4_try(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
#ifndef DS4_CUDA_HAVE_MXF4
(void)scores; (void)q; (void)weights; (void)index_comp;
(void)n_comp; (void)n_tokens; (void)pos0; (void)n_head;
(void)head_dim; (void)ratio; (void)scale;
return 0;
#else
if (g_quality_mode || getenv("DS4_CUDA_NO_INDEXER_MXF4") != NULL ||
g_n_gpus != 1 || ds4_tensor_device_idx(scores) != 0 ||
n_head != DS4_MXF4_INDEXER_HEADS ||
head_dim != DS4_MXF4_INDEXER_DIM || ratio != 4u ||
n_comp < 1024u || n_tokens < 32u) {
return 0;
}
cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone;
if (cudaStreamIsCapturing(0, &capture) != cudaSuccess ||
capture != cudaStreamCaptureStatusNone) {
return 0;
}
const uint64_t q_rows64 =
(uint64_t)n_tokens * DS4_MXF4_INDEXER_HEADS;
if (q_rows64 > UINT32_MAX) return 0;
const uint32_t q_rows = (uint32_t)q_rows64;
const uint64_t k_codes_bytes = ((uint64_t)n_comp * 64u + 255u) & ~255ull;
const uint64_t k_scales_bytes = ((uint64_t)n_comp * 4u + 255u) & ~255ull;
const uint64_t q_codes_bytes = (q_rows64 * 64u + 255u) & ~255ull;
const uint64_t q_scales_bytes = (q_rows64 * 4u + 255u) & ~255ull;
const uint64_t scratch_bytes = k_codes_bytes + k_scales_bytes +
q_codes_bytes + q_scales_bytes;
unsigned char *scratch = (unsigned char *)
indexer_mxf4_scratch_alloc(scratch_bytes);
if (!scratch) return 0;
unsigned char *k_codes = scratch;
uint32_t *k_scales = (uint32_t *)(scratch + k_codes_bytes);
unsigned char *q_codes = scratch + k_codes_bytes + k_scales_bytes;
uint32_t *q_scales = (uint32_t *)(q_codes + q_codes_bytes);
indexer_mxf4_encode_rows_kernel<<<
(uint32_t)(((uint64_t)n_comp * 32u + 255u) / 256u), 256>>>(
(const float *)index_comp->ptr, k_codes, k_scales, n_comp);
if (!cuda_ok(cudaGetLastError(), "indexer mxf4 K encode launch")) {
return -1;
}
indexer_mxf4_encode_rows_kernel<<<
(uint32_t)(((uint64_t)q_rows * 32u + 255u) / 256u), 256>>>(
(const float *)q->ptr, q_codes, q_scales, q_rows);
if (!cuda_ok(cudaGetLastError(), "indexer mxf4 Q encode launch")) {
return -1;
}
dim3 grid((n_comp + 255u) / 256u,
(n_tokens + 15u) / 16u, 1u);
indexer_scores_mxf4_kernel<<<grid, 512>>>(
(float *)scores->ptr, q_codes, q_scales,
(const float *)weights->ptr, n_comp, n_tokens, pos0, ratio,
scale, k_codes, k_scales);
if (!cuda_ok(cudaGetLastError(), "indexer mxf4 score launch")) {
return -1;
}
return 1;
#endif
}
__global__ static void indexer_topk_kernel(uint32_t *selected, const float *scores, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) {
uint32_t t = blockIdx.x;
if (t >= n_tokens || threadIdx.x != 0) return;
const float *row = scores + (uint64_t)t * n_comp;
uint32_t *sel = selected + (uint64_t)t * top_k;
for (uint32_t k = 0; k < top_k; k++) sel[k] = 0;
for (uint32_t c = 0; c < n_comp; c++) {
float v = row[c];
for (uint32_t k = 0; k < top_k; k++) {
if ((k >= c) || v > row[sel[k]]) {
for (uint32_t j = top_k - 1; j > k; j--) sel[j] = sel[j - 1];
sel[k] = c;
break;
}
}
}
}
__device__ __forceinline__ static bool topk_score_better(float av, uint32_t ai, float bv, uint32_t bi) {
return av > bv || (av == bv && ai < bi);
}
__device__ __forceinline__ static void top2_insert_candidate(
float v,
uint32_t i,
float *v0,
uint32_t *i0,
float *v1,
uint32_t *i1) {
if (i == *i0 || i == *i1) return;
if (topk_score_better(v, i, *v0, *i0)) {
*v1 = *v0;
*i1 = *i0;
*v0 = v;
*i0 = i;
} else if (topk_score_better(v, i, *v1, *i1)) {
*v1 = v;
*i1 = i;
}
}
/* DSpark markov chain step: out = argmax_i(logits[i] + dot(w2[i], w1[prev]))
* over the vocab, entirely on-device (logits row stays resident; the chain
* loop only reads back 4 bytes per draft). w1/w2 are q8_0 with 272-byte rows
* (8 blocks of 32). Single block; ~35 MB w2 read per step. */
__global__ static void dspark_markov_argmax_kernel(
unsigned long long *out_key,
const float *logits,
const unsigned char *w1_row,
const unsigned char *w2,
uint32_t vocab,
uint32_t rank_blocks) {
__shared__ float state[256];
const uint32_t tid = threadIdx.x;
if (tid < rank_blocks * 32u) {
const uint32_t b = tid >> 5, k = tid & 31u;
const unsigned char *blk = w1_row + (uint64_t)b * 34u;
const float d = __half2float(*(const __half *)blk);
state[tid] = d * (float)((const int8_t *)(blk + 2))[k];
}
__syncthreads();
float best_v = -INFINITY;
uint32_t best_i = 0;
for (uint32_t i = blockIdx.x * blockDim.x + tid; i < vocab;
i += gridDim.x * blockDim.x) {
const unsigned char *row = w2 + (uint64_t)i * rank_blocks * 34u;
float acc = 0.0f;
for (uint32_t b = 0; b < rank_blocks; b++) {
const unsigned char *blk = row + (uint64_t)b * 34u;
const float d = __half2float(*(const __half *)blk);
const int8_t *q = (const int8_t *)(blk + 2);
float s = 0.0f;
#pragma unroll
for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * state[b * 32u + k];
acc += d * s;
}
const float v = logits[i] + acc;
if (topk_score_better(v, i, best_v, best_i)) {
best_v = v;
best_i = i;
}
}
__shared__ float vals[256];
__shared__ uint32_t idxs[256];
vals[tid] = best_v;
idxs[tid] = best_i;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0u; stride >>= 1u) {
if (tid < stride) {
if (topk_score_better(vals[tid + stride], idxs[tid + stride],
vals[tid], idxs[tid])) {
vals[tid] = vals[tid + stride];
idxs[tid] = idxs[tid + stride];
}
}
__syncthreads();
}
if (tid == 0u) {
/* Monotonic float key; ~idx in the low bits makes ties resolve to
* the smaller index under atomicMax (matches topk_score_better). */
const unsigned int f = __float_as_uint(vals[0]);
const unsigned int fkey = (f & 0x80000000u) ? ~f : (f | 0x80000000u);
const unsigned long long key =
((unsigned long long)fkey << 32) | (unsigned int)(~idxs[0]);
atomicMax(out_key, key);
}
}
__global__ static void indexer_top1_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens) {
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= 1024u) return;
const float *row = scores + (uint64_t)t * n_comp;
float best_v = -INFINITY;
uint32_t best_i = 0;
for (uint32_t i = tid; i < n_comp; i += 1024u) {
const float v = row[i];
if (topk_score_better(v, i, best_v, best_i)) {
best_v = v;
best_i = i;
}
}
__shared__ float vals[1024];
__shared__ uint32_t idxs[1024];
vals[tid] = best_v;
idxs[tid] = best_i;
__syncthreads();
for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) {
if (tid < stride) {
const float ov = vals[tid + stride];
const uint32_t oi = idxs[tid + stride];
if (topk_score_better(ov, oi, vals[tid], idxs[tid])) {
vals[tid] = ov;
idxs[tid] = oi;
}
}
__syncthreads();
}
if (tid == 0u) selected[t] = idxs[0];
}
__global__ static void indexer_top1_value_kernel(
uint32_t *selected,
float *values,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t index_offset) {
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= 1024u) return;
const float *row = scores + (uint64_t)t * n_comp;
float best_v = -INFINITY;
uint32_t best_i = 0;
for (uint32_t i = tid; i < n_comp; i += 1024u) {
const float v = row[i];
const uint32_t gi = index_offset + i;
const uint32_t best_gi = index_offset + best_i;
if (topk_score_better(v, gi, best_v, best_gi)) {
best_v = v;
best_i = i;
}
}
__shared__ float vals[1024];
__shared__ uint32_t idxs[1024];
vals[tid] = best_v;
idxs[tid] = best_i;
__syncthreads();
for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) {
if (tid < stride) {
const float ov = vals[tid + stride];
const uint32_t oi = idxs[tid + stride];
const uint32_t ogi = index_offset + oi;
const uint32_t gi = index_offset + idxs[tid];
if (topk_score_better(ov, ogi, vals[tid], gi)) {
vals[tid] = ov;
idxs[tid] = oi;
}
}
__syncthreads();
}
if (tid == 0u) {
selected[t] = index_offset + idxs[0];
values[t] = vals[0];
}
}
__global__ static void indexer_top2_value_kernel(
uint32_t *selected,
float *values,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t index_offset) {
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= 1024u) return;
const float *row = scores + (uint64_t)t * n_comp;
float best0_v = -INFINITY;
float best1_v = -INFINITY;
uint32_t best0_i = UINT32_MAX;
uint32_t best1_i = UINT32_MAX;
for (uint32_t i = tid; i < n_comp; i += 1024u) {
const uint32_t gi = index_offset + i;
top2_insert_candidate(row[i], gi,
&best0_v, &best0_i,
&best1_v, &best1_i);
}
__shared__ float vals0[1024];
__shared__ float vals1[1024];
__shared__ uint32_t idxs0[1024];
__shared__ uint32_t idxs1[1024];
vals0[tid] = best0_v;
vals1[tid] = best1_v;
idxs0[tid] = best0_i;
idxs1[tid] = best1_i;
__syncthreads();
for (uint32_t stride = 512u; stride > 0u; stride >>= 1u) {
if (tid < stride) {
top2_insert_candidate(vals0[tid + stride], idxs0[tid + stride],
&vals0[tid], &idxs0[tid],
&vals1[tid], &idxs1[tid]);
top2_insert_candidate(vals1[tid + stride], idxs1[tid + stride],
&vals0[tid], &idxs0[tid],
&vals1[tid], &idxs1[tid]);
}
__syncthreads();
}
if (tid == 0u) {
selected[(uint64_t)t * 2u + 0u] = idxs0[0];
selected[(uint64_t)t * 2u + 1u] = idxs1[0];
values[(uint64_t)t * 2u + 0u] = vals0[0];
values[(uint64_t)t * 2u + 1u] = vals1[0];
}
}
__device__ __forceinline__ static uint32_t topk_float_ordered_key(float v) {
uint32_t u = __float_as_uint(v);
/* The comparison sort treats both signed zeros as equal. */
if ((u & 0x7fffffffu) == 0u) u = 0u;
return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u);
}
__device__ __forceinline__ static uint64_t topk_pack_key(float v, uint32_t idx) {
return ((uint64_t)topk_float_ordered_key(v) << 32u) | (uint64_t)(0xffffffffu - idx);
}
__global__ static void indexer_topk_8192_cub_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
constexpr uint32_t BLOCK_THREADS = 512u;
constexpr uint32_t ITEMS_PER_THREAD = 16u;
using BlockSort = cub::BlockRadixSort<uint64_t, BLOCK_THREADS, ITEMS_PER_THREAD>;
extern __shared__ __align__(16) unsigned char sort_smem[];
typename BlockSort::TempStorage &sort_storage =
*reinterpret_cast<typename BlockSort::TempStorage *>(sort_smem);
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= BLOCK_THREADS) return;
const float *row = scores + (uint64_t)t * n_comp;
uint64_t keys[ITEMS_PER_THREAD];
#pragma unroll
for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) {
const uint32_t i = tid * ITEMS_PER_THREAD + item;
if (i < n_comp) {
keys[item] = topk_pack_key(row[i], i);
} else {
keys[item] = topk_pack_key(-INFINITY, UINT32_MAX);
}
}
BlockSort(sort_storage).SortDescending(keys);
#pragma unroll
for (uint32_t item = 0; item < ITEMS_PER_THREAD; item++) {
const uint32_t i = tid * ITEMS_PER_THREAD + item;
if (i < top_k) {
selected[(uint64_t)t * top_k + i] = 0xffffffffu - (uint32_t)keys[item];
}
}
}
__global__ static void indexer_topk_1024_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
uint32_t t = blockIdx.x;
uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= 1024u) return;
__shared__ float vals[1024];
__shared__ uint32_t idxs[1024];
const float *row = scores + (uint64_t)t * n_comp;
if (tid < n_comp) {
vals[tid] = row[tid];
idxs[tid] = tid;
} else {
vals[tid] = -INFINITY;
idxs[tid] = UINT32_MAX;
}
__syncthreads();
for (uint32_t k = 2u; k <= 1024u; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
uint32_t other = tid ^ j;
if (other > tid && other < 1024u) {
const float av = vals[tid];
const float bv = vals[other];
const uint32_t ai = idxs[tid];
const uint32_t bi = idxs[other];
const bool desc_half = (tid & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[tid] = bv;
idxs[tid] = bi;
vals[other] = av;
idxs[other] = ai;
}
}
__syncthreads();
}
}
if (tid < top_k) selected[(uint64_t)t * top_k + tid] = idxs[tid];
}
template <uint32_t SORT_N>
__global__ static void indexer_topk_pow2_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
uint32_t t = blockIdx.x;
uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
__shared__ float vals[SORT_N];
__shared__ uint32_t idxs[SORT_N];
const float *row = scores + (uint64_t)t * n_comp;
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
if (i < n_comp) {
vals[i] = row[i];
idxs[i] = i;
} else {
vals[i] = -INFINITY;
idxs[i] = UINT32_MAX;
}
}
__syncthreads();
for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t other = i ^ j;
if (other > i && other < SORT_N) {
const float av = vals[i];
const float bv = vals[other];
const uint32_t ai = idxs[i];
const uint32_t bi = idxs[other];
const bool desc_half = (i & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[i] = bv;
idxs[i] = bi;
vals[other] = av;
idxs[other] = ai;
}
}
}
__syncthreads();
}
}
for (uint32_t i = tid; i < top_k; i += blockDim.x) {
selected[(uint64_t)t * top_k + i] = idxs[i];
}
}
template <uint32_t SORT_N>
__global__ static void indexer_topk_pow2_u16_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
uint32_t t = blockIdx.x;
uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
__shared__ float vals[SORT_N];
__shared__ uint16_t idxs[SORT_N];
const float *row = scores + (uint64_t)t * n_comp;
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
if (i < n_comp) {
vals[i] = row[i];
idxs[i] = (uint16_t)i;
} else {
vals[i] = -INFINITY;
idxs[i] = UINT16_MAX;
}
}
__syncthreads();
for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t other = i ^ j;
if (other > i && other < SORT_N) {
const float av = vals[i];
const float bv = vals[other];
const uint32_t ai = idxs[i];
const uint32_t bi = idxs[other];
const bool desc_half = (i & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[i] = bv;
idxs[i] = (uint16_t)bi;
vals[other] = av;
idxs[other] = (uint16_t)ai;
}
}
}
__syncthreads();
}
}
for (uint32_t i = tid; i < top_k; i += blockDim.x) {
selected[(uint64_t)t * top_k + i] = idxs[i];
}
}
template <uint32_t SORT_N>
__global__ static void indexer_topk_chunk_pow2_kernel(
uint32_t *candidates,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k,
uint32_t candidate_stride) {
uint32_t t = blockIdx.x;
uint32_t chunk = blockIdx.y;
uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
const uint32_t chunk_start = chunk * SORT_N;
if (chunk_start >= n_comp) return;
const uint32_t chunk_n = n_comp - chunk_start < SORT_N ? n_comp - chunk_start : SORT_N;
__shared__ float vals[SORT_N];
__shared__ uint32_t idxs[SORT_N];
const float *row = scores + (uint64_t)t * n_comp;
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
if (i < chunk_n) {
vals[i] = row[chunk_start + i];
idxs[i] = chunk_start + i;
} else {
vals[i] = -INFINITY;
idxs[i] = UINT32_MAX;
}
}
__syncthreads();
for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t other = i ^ j;
if (other > i && other < SORT_N) {
const float av = vals[i];
const float bv = vals[other];
const uint32_t ai = idxs[i];
const uint32_t bi = idxs[other];
const bool desc_half = (i & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[i] = bv;
idxs[i] = bi;
vals[other] = av;
idxs[other] = ai;
}
}
}
__syncthreads();
}
}
uint32_t *out = candidates + (uint64_t)t * candidate_stride + chunk * top_k;
for (uint32_t i = tid; i < top_k; i += blockDim.x) {
out[i] = idxs[i];
}
}
template <uint32_t SORT_N>
__global__ static void indexer_topk_merge_pow2_kernel(
uint32_t *selected,
const uint32_t *candidates,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k,
uint32_t candidate_count,
uint32_t candidate_stride) {
uint32_t t = blockIdx.x;
uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
__shared__ float vals[SORT_N];
__shared__ uint32_t idxs[SORT_N];
const float *row = scores + (uint64_t)t * n_comp;
const uint32_t *cand = candidates + (uint64_t)t * candidate_stride;
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t idx = UINT32_MAX;
float v = -INFINITY;
if (i < candidate_count) {
idx = cand[i];
if (idx < n_comp) v = row[idx];
}
vals[i] = v;
idxs[i] = idx;
}
__syncthreads();
for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t other = i ^ j;
if (other > i && other < SORT_N) {
const float av = vals[i];
const float bv = vals[other];
const uint32_t ai = idxs[i];
const uint32_t bi = idxs[other];
const bool desc_half = (i & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[i] = bv;
idxs[i] = bi;
vals[other] = av;
idxs[other] = ai;
}
}
}
__syncthreads();
}
}
for (uint32_t i = tid; i < top_k; i += blockDim.x) {
selected[(uint64_t)t * top_k + i] = idxs[i];
}
}
template <uint32_t SORT_N>
__global__ static void indexer_topk_tree_merge_pow2_kernel(
uint32_t *out,
const uint32_t *candidates,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k,
uint32_t n_sets,
uint32_t merge_group,
uint32_t candidate_stride,
uint32_t out_stride) {
uint32_t t = blockIdx.x;
uint32_t group = blockIdx.y;
uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
const uint32_t set0 = group * merge_group;
if (set0 >= n_sets) return;
uint32_t set_count = n_sets - set0;
if (set_count > merge_group) set_count = merge_group;
const uint32_t candidate_count = set_count * top_k;
__shared__ float vals[SORT_N];
__shared__ uint32_t idxs[SORT_N];
const float *row = scores + (uint64_t)t * n_comp;
const uint32_t *cand = candidates + (uint64_t)t * candidate_stride + set0 * top_k;
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t idx = UINT32_MAX;
float v = -INFINITY;
if (i < candidate_count) {
idx = cand[i];
if (idx < n_comp) v = row[idx];
}
vals[i] = v;
idxs[i] = idx;
}
__syncthreads();
for (uint32_t k = 2u; k <= SORT_N; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
for (uint32_t i = tid; i < SORT_N; i += blockDim.x) {
uint32_t other = i ^ j;
if (other > i && other < SORT_N) {
const float av = vals[i];
const float bv = vals[other];
const uint32_t ai = idxs[i];
const uint32_t bi = idxs[other];
const bool desc_half = (i & k) == 0u;
const bool swap = desc_half
? topk_score_better(bv, bi, av, ai)
: topk_score_better(av, ai, bv, bi);
if (swap) {
vals[i] = bv;
idxs[i] = bi;
vals[other] = av;
idxs[other] = ai;
}
}
}
__syncthreads();
}
}
uint32_t *dst = out + (uint64_t)t * out_stride + group * top_k;
for (uint32_t i = tid; i < top_k; i += blockDim.x) {
dst[i] = idxs[i];
}
}
/* Exact streaming top-512 for wide prefill rows. The threshold is the
* 512th-best key seen so far, so it can only discard candidates that cannot
* belong to the final top set. Packing supplies the same value/index total
* order as the existing bitonic and CUB tiers. */
__global__ static void __launch_bounds__(512) indexer_topk_stream512_kernel(
uint32_t *selected,
const float *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
constexpr uint32_t STREAM_THREADS = 512u;
constexpr uint32_t STREAM_ITEMS = 4u;
constexpr uint32_t STREAM_CAP = STREAM_THREADS * STREAM_ITEMS;
using StreamSort = cub::BlockRadixSort<uint64_t, STREAM_THREADS, STREAM_ITEMS>;
__shared__ uint64_t buf[STREAM_CAP];
__shared__ typename StreamSort::TempStorage sort_tmp;
__shared__ uint32_t s_cnt;
__shared__ uint64_t s_thr;
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens) return;
const float *row = scores + (uint64_t)t * n_comp;
if (tid == 0u) {
s_cnt = 0u;
s_thr = 0u;
}
__syncthreads();
const uint32_t start =
(uint32_t)(((uint64_t)(t + 1u) * 0x9E3779B9u) % n_comp);
const uint32_t tile = blockDim.x;
for (uint32_t base = 0; base < n_comp; base += tile) {
const uint64_t thr = s_thr;
const uint32_t i = base + tid;
uint64_t key = 0u;
bool take = false;
if (i < n_comp) {
uint32_t c = start + i;
if (c >= n_comp) c -= n_comp;
key = topk_pack_key(row[c], c);
take = key > thr;
}
const uint32_t ballot = __ballot_sync(0xffffffffu, take);
if (take) {
const uint32_t lane = tid & 31u;
const uint32_t rank = __popc(ballot & ((1u << lane) - 1u));
uint32_t pos = 0u;
if (rank == 0u) pos = atomicAdd(&s_cnt, __popc(ballot));
pos = __shfl_sync(ballot, pos, __ffs(ballot) - 1);
buf[pos + rank] = key;
}
__syncthreads();
if (s_cnt > STREAM_CAP - tile) {
const uint32_t cnt = s_cnt;
uint64_t keys[STREAM_ITEMS];
#pragma unroll
for (uint32_t k = 0; k < STREAM_ITEMS; k++) {
const uint32_t j = tid * STREAM_ITEMS + k;
keys[k] = j < cnt ? buf[j] : 0u;
}
__syncthreads();
StreamSort(sort_tmp).SortDescending(keys);
__syncthreads();
if (tid < top_k / STREAM_ITEMS) {
#pragma unroll
for (uint32_t k = 0; k < STREAM_ITEMS; k++) {
buf[tid * STREAM_ITEMS + k] = keys[k];
}
}
if (tid == top_k / STREAM_ITEMS - 1u) {
s_thr = keys[STREAM_ITEMS - 1u];
}
if (tid == 0u) s_cnt = top_k;
__syncthreads();
}
}
const uint32_t cnt = s_cnt;
uint64_t keys[STREAM_ITEMS];
#pragma unroll
for (uint32_t k = 0; k < STREAM_ITEMS; k++) {
const uint32_t i = tid * STREAM_ITEMS + k;
keys[k] = i < cnt ? buf[i] : 0u;
}
__syncthreads();
StreamSort(sort_tmp).SortDescending(keys);
__syncthreads();
if (tid < top_k / STREAM_ITEMS) {
#pragma unroll
for (uint32_t k = 0; k < STREAM_ITEMS; k++) {
selected[(uint64_t)t * top_k + tid * STREAM_ITEMS + k] =
0xffffffffu - (uint32_t)keys[k];
}
}
}
__global__ static void indexed_topk_sort_512_asc_kernel(
int32_t *dst,
const int32_t *src,
uint32_t n_tokens) {
const uint32_t t = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (t >= n_tokens || tid >= 512u) return;
__shared__ int32_t rows[512];
const int32_t *src_row = src + (uint64_t)t * 512u;
int32_t *dst_row = dst + (uint64_t)t * 512u;
rows[tid] = src_row[tid];
__syncthreads();
for (uint32_t k = 2u; k <= 512u; k <<= 1u) {
for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) {
const uint32_t other = tid ^ j;
if (other > tid && other < 512u) {
const int32_t a = rows[tid];
const int32_t b = rows[other];
const bool up = (tid & k) == 0u;
if ((up && a > b) || (!up && a < b)) {
rows[tid] = b;
rows[other] = a;
}
}
__syncthreads();
}
}
dst_row[tid] = rows[tid];
}
__global__ static void topk_mask_kernel(float *mask, const uint32_t *topk, uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_comp;
if (gid >= n) return;
uint32_t t = gid / n_comp;
uint32_t c = gid - (uint64_t)t * n_comp;
float v = -INFINITY;
for (uint32_t k = 0; k < top_k; k++) {
if (topk[(uint64_t)t * top_k + k] == c) {
v = 0.0f;
break;
}
}
mask[gid] = v;
}
extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n_vocab, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
(void)n_vocab;
if (!out_hc || !model_map || weight_offset >= model_size) return 0;
uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t);
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0;
const int logical_tier = ds4_tensor_device_idx(out_hc);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "token_embd");
if (!wptr) return 0;
uint32_t n = n_embd * n_hc;
embed_token_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc);
return cuda_ok(cudaGetLastError(), "embed token launch");
}
extern "C" int ds4_gpu_embed_tokens_hc_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *tokens_t,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc) {
if (!out_hc || !tokens_t || !model_map ||
weight_offset > model_size ||
(uint64_t)n_vocab * n_embd * sizeof(uint16_t) > model_size - weight_offset ||
tokens_t->bytes < (uint64_t)n_tokens * sizeof(int32_t) ||
out_hc->bytes < (uint64_t)n_tokens * n_hc * n_embd * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out_hc);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset,
(uint64_t)n_vocab * n_embd * sizeof(uint16_t),
logical_tier,
"token_embd");
if (!wptr) return 0;
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>(
(float *)out_hc->ptr,
(const int32_t *)tokens_t->ptr,
(const __half *)wptr,
n_vocab, n_tokens, n_embd, n_hc);
return cuda_ok(cudaGetLastError(), "embed tokens launch");
}
static int indexer_scores_launch(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale,
uint32_t causal) {
if (!scores || !q || !weights || !index_comp ||
n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) ||
index_comp->bytes < (uint64_t)n_comp * head_dim * sizeof(float) ||
scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float)) {
return 0;
}
if (causal && ratio == 0) return 0;
const int mxf4 = indexer_scores_mxf4_try(
scores, q, weights, index_comp, n_comp, n_tokens, pos0,
n_head, head_dim, ratio, scale);
if (mxf4 != 0) return mxf4 > 0;
if (n_tokens == 1u && head_dim == 128u && n_head == 64u &&
getenv("DS4_CUDA_NO_INDEXER_DIRECT_ONE") == NULL) {
indexer_score_one_direct_kernel<<<n_comp, 128>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, pos0, ratio,
scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer score one direct launch");
}
if (!g_quality_mode && head_dim == 128u && n_head == 64u &&
getenv("DS4_CUDA_NO_INDEXER_WMMA") == NULL) {
if (getenv("DS4_CUDA_NO_INDEXER_WMMA128") == NULL) {
dim3 grid((n_comp + 127u) / 128u, (n_tokens + 31u) / 32u, 1);
indexer_scores_wmma128_kernel<<<grid, 256>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, n_tokens, pos0, n_head,
head_dim, ratio, scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer scores wmma128 launch");
} else if (getenv("DS4_CUDA_NO_INDEXER_WMMA64") == NULL) {
dim3 grid((n_comp + 63u) / 64u, (n_tokens + 15u) / 16u, 1);
indexer_scores_wmma64_kernel<<<grid, 128>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, n_tokens, pos0, n_head,
head_dim, ratio, scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer scores wmma64 launch");
} else if (getenv("DS4_CUDA_NO_INDEXER_WMMA32") == NULL) {
dim3 grid((n_comp + 31u) / 32u, (n_tokens + 15u) / 16u, 1);
indexer_scores_wmma32_kernel<<<grid, 64>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, n_tokens, pos0, n_head,
head_dim, ratio, scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer scores wmma32 launch");
} else {
dim3 grid((n_comp + 15u) / 16u, (n_tokens + 15u) / 16u, 1);
indexer_scores_wmma_kernel<<<grid, 32>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, n_tokens, pos0, n_head,
head_dim, ratio, scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer scores wmma launch");
}
}
dim3 grid(n_comp, n_tokens, 1);
indexer_scores_kernel<<<grid, 256>>>((float *)scores->ptr,
(const float *)q->ptr,
(const float *)weights->ptr,
(const float *)index_comp->ptr,
n_comp, n_tokens, pos0, n_head,
head_dim, ratio, scale, causal ? 1 : 0);
return cuda_ok(cudaGetLastError(), "indexer scores launch");
}
extern "C" int ds4_gpu_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_head,
uint32_t head_dim,
float scale) {
return indexer_scores_launch(scores, q, weights, index_comp, n_comp, 1, 0,
n_head, head_dim, 1, scale, 0);
}
extern "C" int ds4_gpu_indexer_scores_prefill_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, 0,
n_head, head_dim, ratio, scale, 1);
}
extern "C" int ds4_gpu_indexer_scores_decode_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
return indexer_scores_launch(scores, q, weights, index_comp, n_comp, n_tokens, pos0,
n_head, head_dim, ratio, scale, 1);
}
extern "C" int ds4_gpu_dspark_markov_argmax_tensor(
ds4_gpu_tensor *out_idx,
const ds4_gpu_tensor *logits_row,
const void *model_map,
uint64_t model_size,
uint64_t w1_offset,
uint64_t w2_offset,
uint32_t prev_token,
uint32_t vocab,
uint32_t rank) {
if (!out_idx || !logits_row || !model_map || vocab == 0 ||
rank == 0 || (rank & 31u) != 0u || rank > 256u ||
out_idx->bytes < sizeof(unsigned long long) ||
logits_row->bytes < (uint64_t)vocab * sizeof(float)) {
return 0;
}
const uint32_t rank_blocks = rank / 32u;
const uint64_t row_bytes = (uint64_t)rank_blocks * 34u;
if (w1_offset > model_size ||
(uint64_t)prev_token * row_bytes + row_bytes > model_size - w1_offset ||
w2_offset > model_size ||
(uint64_t)vocab * row_bytes > model_size - w2_offset) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(logits_row);
const unsigned char *w1_row = (const unsigned char *)cuda_resolve_weight_ptr(
model_map, w1_offset + (uint64_t)prev_token * row_bytes,
row_bytes, logical_tier, "markov_w1_row");
const unsigned char *w2 = (const unsigned char *)cuda_resolve_weight_ptr(
model_map, w2_offset, (uint64_t)vocab * row_bytes,
logical_tier, "markov_w2");
if (!w1_row || !w2) return 0;
int dev_save = 0;
if (cudaGetDevice(&dev_save) != cudaSuccess) return 0;
if (logical_tier != dev_save && cudaSetDevice(logical_tier) != cudaSuccess) {
return 0;
}
int rc = cudaMemsetAsync(out_idx->ptr, 0,
sizeof(unsigned long long)) == cudaSuccess;
if (rc) {
dspark_markov_argmax_kernel<<<128, 256>>>(
(unsigned long long *)out_idx->ptr,
(const float *)logits_row->ptr,
w1_row, w2, vocab, rank_blocks);
rc = cuda_ok(cudaGetLastError(), "dspark markov argmax launch");
}
if (logical_tier != dev_save) (void)cudaSetDevice(dev_save);
return rc;
}
extern "C" int ds4_gpu_indexer_topk_tensor(
ds4_gpu_tensor *selected,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 ||
top_k > n_comp ||
scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) {
return 0;
}
if (top_k == 1u && !g_cuda_no_top1) {
indexer_top1_kernel<<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp,
n_tokens);
return cuda_ok(cudaGetLastError(), "indexer top1 launch");
}
if (top_k == 2048u && n_comp <= 4096u &&
getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) {
indexer_topk_pow2_kernel<4096><<<n_tokens, 1024>>>(
(uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 2048-wide launch");
}
if (top_k == 2048u && n_comp > 4096u &&
getenv("DS4_CUDA_NO_TOPK2048_WIDE") == NULL) {
const uint32_t chunk_n = 4096u;
const uint32_t merge_group = 2u;
const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n;
const uint32_t candidate_stride = n_chunks * top_k;
uint32_t n_sets = n_chunks;
uint64_t scratch_u32_per_token = candidate_stride;
while (n_sets > merge_group) {
n_sets = (n_sets + merge_group - 1u) / merge_group;
scratch_u32_per_token += (uint64_t)n_sets * top_k;
}
if (scratch_u32_per_token >
UINT64_MAX / n_tokens / sizeof(uint32_t)) {
return 0;
}
int exec_tier = ds4_tensor_device_idx(selected);
int current_device = -1;
if (cudaGetDevice(&current_device) == cudaSuccess) {
for (int t = 0; t < g_n_gpus; t++) {
if (g_gpu[t].device_id == current_device) {
exec_tier = t;
break;
}
}
}
const uint64_t tmp_bytes =
(uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t);
uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on(
exec_tier, tmp_bytes, "indexer topk 2048-wide tree");
if (!scratch) return 0;
uint32_t *cur = scratch;
n_sets = n_chunks;
uint32_t cur_stride = candidate_stride;
dim3 grid_chunks(n_tokens, n_chunks, 1);
indexer_topk_chunk_pow2_kernel<4096><<<grid_chunks, 1024>>>(
cur, (const float *)scores->ptr,
n_comp, n_tokens, top_k, candidate_stride);
if (!cuda_ok(cudaGetLastError(),
"indexer topk 2048-wide chunk launch")) {
return 0;
}
while (n_sets > merge_group) {
const uint32_t next_sets =
(n_sets + merge_group - 1u) / merge_group;
const uint32_t next_stride = next_sets * top_k;
uint32_t *next = cur + (uint64_t)n_tokens * cur_stride;
dim3 grid_merge(n_tokens, next_sets, 1);
indexer_topk_tree_merge_pow2_kernel<4096><<<grid_merge, 1024>>>(
next, cur, (const float *)scores->ptr,
n_comp, n_tokens, top_k, n_sets, merge_group,
cur_stride, next_stride);
if (!cuda_ok(cudaGetLastError(),
"indexer topk 2048-wide merge launch")) {
return 0;
}
cur = next;
n_sets = next_sets;
cur_stride = next_stride;
}
indexer_topk_merge_pow2_kernel<4096><<<n_tokens, 1024>>>(
(uint32_t *)selected->ptr,
cur, (const float *)scores->ptr,
n_comp, n_tokens, top_k, n_sets * top_k, cur_stride);
return cuda_ok(cudaGetLastError(),
"indexer topk 2048-wide final launch");
}
if (top_k == 512u && n_comp <= 1024u &&
getenv("DS4_CUDA_NO_TOPK1024") == NULL) {
indexer_topk_1024_kernel<<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 1024 launch");
}
if (top_k == 512u && n_comp <= 2048u &&
getenv("DS4_CUDA_NO_TOPK2048") == NULL) {
indexer_topk_pow2_kernel<2048><<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 2048 launch");
}
if (top_k == 512u && n_comp <= 4096u &&
getenv("DS4_CUDA_NO_TOPK2048") == NULL) {
if (n_comp == 4096u) {
using TopkCubSort = cub::BlockRadixSort<uint64_t, 512, 16>;
const int smem = (int)sizeof(typename TopkCubSort::TempStorage);
int dev = 0;
int max_optin_smem = 0;
cudaError_t attr_err = cudaGetDevice(&dev);
if (attr_err == cudaSuccess) {
attr_err = cudaDeviceGetAttribute(&max_optin_smem,
cudaDevAttrMaxSharedMemoryPerBlockOptin,
dev);
}
if (attr_err == cudaSuccess && max_optin_smem >= smem) {
attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem);
if (attr_err == cudaSuccess) {
indexer_topk_8192_cub_kernel<<<n_tokens, 512, (size_t)smem>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 4096 cub launch");
}
}
}
indexer_topk_pow2_kernel<4096><<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 4096 launch");
}
if (top_k == 512u && n_comp <= 8192u &&
getenv("DS4_CUDA_NO_TOPK2048") == NULL &&
getenv("DS4_CUDA_NO_TOPK8192") == NULL) {
if (n_comp > 4096u) {
using TopkCubSort = cub::BlockRadixSort<uint64_t, 512, 16>;
const int smem = (int)sizeof(typename TopkCubSort::TempStorage);
int dev = 0;
int max_optin_smem = 0;
cudaError_t attr_err = cudaGetDevice(&dev);
if (attr_err == cudaSuccess) {
attr_err = cudaDeviceGetAttribute(&max_optin_smem,
cudaDevAttrMaxSharedMemoryPerBlockOptin,
dev);
}
if (attr_err == cudaSuccess && max_optin_smem >= smem) {
attr_err = cudaFuncSetAttribute(indexer_topk_8192_cub_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem);
if (attr_err == cudaSuccess) {
indexer_topk_8192_cub_kernel<<<n_tokens, 512, (size_t)smem>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 8192 cub launch");
}
}
}
indexer_topk_pow2_u16_kernel<8192><<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk 8192 launch");
}
if (top_k == 512u && n_tokens >= 32u &&
getenv("DS4_CUDA_NO_TOPK2048") == NULL &&
getenv("DS4_CUDA_NO_TOPK_STREAM") == NULL) {
indexer_topk_stream512_kernel<<<n_tokens, 512>>>(
(uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk stream launch");
}
if (top_k == 512u && getenv("DS4_CUDA_NO_TOPK2048") == NULL &&
getenv("DS4_CUDA_NO_TOPK_CHUNKED") == NULL) {
const uint32_t chunk_n = 4096u;
const uint32_t n_chunks = (n_comp + chunk_n - 1u) / chunk_n;
const uint32_t candidate_stride = n_chunks * top_k;
uint32_t n_sets = n_chunks;
uint64_t scratch_u32_per_token = candidate_stride;
while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) {
n_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP;
scratch_u32_per_token += (uint64_t)n_sets * top_k;
}
if (scratch_u32_per_token > UINT64_MAX / n_tokens / sizeof(uint32_t)) return 0;
const uint64_t tmp_bytes = (uint64_t)n_tokens * scratch_u32_per_token * sizeof(uint32_t);
const int logical_tier = ds4_tensor_device_idx(selected);
uint32_t *scratch = (uint32_t *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "indexer topk tree");
if (!scratch) return 0;
uint32_t *cur = scratch;
n_sets = n_chunks;
uint32_t cur_stride = candidate_stride;
dim3 grid_chunks(n_tokens, n_chunks, 1);
indexer_topk_chunk_pow2_kernel<4096><<<grid_chunks, 1024>>>(cur,
(const float *)scores->ptr,
n_comp,
n_tokens,
top_k,
candidate_stride);
if (!cuda_ok(cudaGetLastError(), "indexer topk chunk launch")) return 0;
while (n_sets > DS4_CUDA_TOPK_MERGE_GROUP) {
const uint32_t next_sets = (n_sets + DS4_CUDA_TOPK_MERGE_GROUP - 1u) / DS4_CUDA_TOPK_MERGE_GROUP;
const uint32_t next_stride = next_sets * top_k;
uint32_t *next = cur + (uint64_t)n_tokens * cur_stride;
dim3 grid_merge(n_tokens, next_sets, 1);
indexer_topk_tree_merge_pow2_kernel<4096><<<grid_merge, 1024>>>(
next,
cur,
(const float *)scores->ptr,
n_comp,
n_tokens,
top_k,
n_sets,
DS4_CUDA_TOPK_MERGE_GROUP,
cur_stride,
next_stride);
if (!cuda_ok(cudaGetLastError(), "indexer topk tree merge launch")) return 0;
cur = next;
n_sets = next_sets;
cur_stride = next_stride;
}
indexer_topk_merge_pow2_kernel<4096><<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
cur,
(const float *)scores->ptr,
n_comp,
n_tokens,
top_k,
n_sets * top_k,
cur_stride);
return cuda_ok(cudaGetLastError(), "indexer topk tree final launch");
}
/* Candidate-block selection and short prefixes need other k values too.
* Do not send a full small-row sort through serial insertion sort. */
if (n_comp <= 8192u) {
uint32_t *out = (uint32_t *)selected->ptr;
const float *in = (const float *)scores->ptr;
if (n_comp <= 256u)
indexer_topk_pow2_kernel<256><<<n_tokens, 256>>>(out, in, n_comp, n_tokens, top_k);
else if (n_comp <= 1024u)
indexer_topk_pow2_kernel<1024><<<n_tokens, 256>>>(out, in, n_comp, n_tokens, top_k);
else if (n_comp <= 2048u)
indexer_topk_pow2_kernel<2048><<<n_tokens, 512>>>(out, in, n_comp, n_tokens, top_k);
else if (n_comp <= 4096u)
indexer_topk_pow2_kernel<4096><<<n_tokens, 1024>>>(out, in, n_comp, n_tokens, top_k);
else
indexer_topk_pow2_u16_kernel<8192><<<n_tokens, 1024>>>(out, in, n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer general small topk launch");
}
indexer_topk_kernel<<<n_tokens, 1>>>((uint32_t *)selected->ptr,
(const float *)scores->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "indexer topk launch");
}
extern "C" int ds4_gpu_indexer_top1_value_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *values,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t index_offset) {
if (!selected || !values || !scores || n_comp == 0 || n_tokens == 0 ||
scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * sizeof(uint32_t) ||
values->bytes < (uint64_t)n_tokens * sizeof(float)) {
return 0;
}
indexer_top1_value_kernel<<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(float *)values->ptr,
(const float *)scores->ptr,
n_comp,
n_tokens,
index_offset);
return cuda_ok(cudaGetLastError(), "indexer top1 value launch");
}
extern "C" int ds4_gpu_indexer_top2_value_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *values,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t index_offset) {
if (!selected || !values || !scores || n_comp < 2u || n_tokens == 0 ||
scores->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * 2u * sizeof(uint32_t) ||
values->bytes < (uint64_t)n_tokens * 2u * sizeof(float)) {
return 0;
}
indexer_top2_value_kernel<<<n_tokens, 1024>>>((uint32_t *)selected->ptr,
(float *)values->ptr,
(const float *)scores->ptr,
n_comp,
n_tokens,
index_offset);
return cuda_ok(cudaGetLastError(), "indexer top2 value launch");
}
extern "C" int ds4_gpu_dsv4_topk_mask_tensor(
ds4_gpu_tensor *mask,
const ds4_gpu_tensor *topk,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0 ||
mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float) ||
topk->bytes < (uint64_t)n_tokens * top_k * sizeof(uint32_t)) {
return 0;
}
uint64_t n = (uint64_t)n_tokens * n_comp;
uint64_t nk = (uint64_t)n_tokens * top_k;
uint64_t blocks = ((n > nk ? n : nk) + 255) / 256;
topk_mask_kernel<<<blocks, 256>>>((float *)mask->ptr,
(const uint32_t *)topk->ptr,
n_comp, n_tokens, top_k);
return cuda_ok(cudaGetLastError(), "topk mask launch");
}
/* GLM opt-in: batched q8_0 matmuls with blocks > 32 may run as a
* streaming dequant-to-f16 GEMM (exact-q8 native kernels only cover
* blocks <= 32). Never enabled on DS4 paths, keeping them byte-stable. */
static int g_q8_dequant_gemm_enabled = 0;
extern "C" void ds4_gpu_enable_q8_dequant_gemm(void) {
g_q8_dequant_gemm_enabled = 1;
}
__global__ static void q8_0_dequant_f16_kernel(
__half *out,
const unsigned char *w,
uint64_t total_blocks,
uint32_t blocks_per_row,
uint32_t in_dim) {
/* Two threads per q8_0 block; each converts 16 values with half2
* stores so a warp writes 512B contiguously per block pair. */
const uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t b = tid >> 1;
if (b >= total_blocks) return;
const uint32_t half_idx = (uint32_t)tid & 1u;
const unsigned char *blk = w + b * 34u;
const float d = __half2float(*(const __half *)blk);
const int8_t *q = (const int8_t *)(blk + 2) + half_idx * 16u;
const uint64_t row = b / blocks_per_row;
const uint32_t col = (uint32_t)(b - row * blocks_per_row) * 32u +
half_idx * 16u;
__half2 *dst = (__half2 *)(out + row * in_dim + col);
#pragma unroll
for (int k = 0; k < 8; k++) {
dst[k] = __floats2half2_rn(d * (float)q[2 * k],
d * (float)q[2 * k + 1]);
}
}
static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok, const char *label) {
if (!out || !x || !model_map) return 0;
uint64_t blocks = (in_dim + 31) / 32;
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0;
uint64_t weight_bytes = out_dim * blocks * 34;
if (weight_bytes > model_size - weight_offset) return 0;
if (x->bytes < n_tok * in_dim * sizeof(float) ||
out->bytes < n_tok * out_dim * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const int physical_device =
(g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus)
? g_gpu[logical_tier].device_id : 0;
const uint64_t aligned_bytes =
(in_dim % 1024u) == 0 && (out_dim % 128u) == 0
? ds4_mmq_q8_0_aligned_bytes((int)out_dim, (int)in_dim)
: 0;
const char *aligned = aligned_bytes && cuda_aligned_q8_enabled()
? cuda_derived_weight_ptr(
model_map, weight_offset, weight_bytes,
CUDA_DERIVED_Q8_0_ALIGNED_DENSE,
in_dim, out_dim, 1u, aligned_bytes)
: NULL;
if (aligned && (n_tok == 1u ||
(n_tok <= 8u && ds4_gpu_device_is_spark() &&
!g_q8_dequant_gemm_enabled))) {
const int rc = ds4_mmq_q8_0_aligned_dense_vec(
aligned, (const float *)x->ptr, (float *)out->ptr,
(int)out_dim, (int)n_tok, (int)in_dim, cuda_decode_stream());
if (rc == 0) return 1;
}
if (aligned && n_tok >= 512u && out_dim >= 2048u &&
in_dim <= 4096u && cuda_use_mmq()) {
const char *d2r = getenv("DS4_MMQ_DENSE_D2R");
if (!d2r || strcmp(d2r, "0") != 0) {
const int rc = ds4_mmq_q8_0_dense_d2r(
aligned, (const float *)x->ptr, (float *)out->ptr,
(int)out_dim, (int)n_tok, (int)in_dim, (cudaStream_t)0);
if (rc == 0) {
static int logged = 0;
if (!logged) {
logged = 1;
fprintf(stderr,
"ds4: dense Q8 prefill using aligned D2R\n");
}
return 1;
}
fprintf(stderr,
"ds4: aligned dense Q8 D2R returned %d "
"(label='%s' M=%llu N=%llu K=%llu); falling back\n",
rc, label ? label : "",
(unsigned long long)out_dim,
(unsigned long long)n_tok,
(unsigned long long)in_dim);
}
}
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "q8_0");
if (!wptr) return 0;
/* mmq fused-dequant-matmul prefill tier (ported from the Entrpi/ds4
* fork). Layout-compatible drop-in for the cuBLAS+dequant pipeline
* below: mmq's [out_dim, n_tok] column-major output flattens to
* [n_tok, out_dim] row-major, exactly what ds4 stores in out->ptr,
* and the Q8_0 weight is already in mmq's expected row-major block
* layout (the GGUF on-disk format). K must be a multiple of 256;
* every Q8_0 weight in V4 Flash satisfies this, odd shapes fall
* through to the legacy paths. */
/* GLM uses aligned D2R above for its profitable large projections. The
* raw dense MMQ fallback is allocator-dependent for other GLM shapes on
* GB10, including short prefills, so retain the deterministic kernels
* below. DeepSeek keeps its established MMQ prefill path. */
if (n_tok > 1 && (in_dim % 256u) == 0 &&
!g_q8_dequant_gemm_enabled && cuda_use_mmq()) {
int rc = ds4_mmq_q8_0_dense(wptr, (const float *)x->ptr, (float *)out->ptr,
(int)out_dim, (int)n_tok, (int)in_dim,
(cudaStream_t)0);
if (rc == 0) return 1;
fprintf(stderr, "ds4: ds4_mmq_q8_0_dense returned %d (label='%s' in=%llu out=%llu n_tok=%llu); falling back\n",
rc, label ? label : "", (unsigned long long)in_dim,
(unsigned long long)out_dim, (unsigned long long)n_tok);
}
if (g_cublas_ready && n_tok > 1) {
const float *w_f32 = cuda_q8_f32_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label);
if (w_f32) {
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w_f32,
(int)in_dim,
(const float *)x->ptr,
(int)in_dim,
&beta,
(float *)out->ptr,
(int)out_dim);
return cublas_ok(st, "q8 fp32 matmul");
}
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, physical_device, label);
if (w_f16) {
const uint64_t xh_count = n_tok * in_dim;
__half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "q8 f16 gemm activations");
if (!xh) return 0;
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256, 0, cuda_decode_stream()>>>(xh, (const float *)x->ptr, xh_count);
if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0;
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w_f16,
CUDA_R_16F,
(int)in_dim,
xh,
CUDA_R_16F,
(int)in_dim,
&beta,
out->ptr,
CUDA_R_32F,
(int)out_dim,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
if (st == CUBLAS_STATUS_SUCCESS) return 1;
fprintf(stderr, "ds4: cuBLAS q8 f16 matmul failed: status %d\n", (int)st);
cuda_q8_f16_cache_disable_after_failure("cuBLAS f16 matmul failure",
in_dim * out_dim * sizeof(__half));
/* The F16 expansion cache is only an optimization. If cuBLAS
* rejects the cached path under memory pressure, retry the same
* operation through the native Q8 kernels below. */
}
}
if (g_q8_dequant_gemm_enabled && g_cublas_ready &&
n_tok >= 128u && blocks > 32u && (in_dim & 31u) == 0u) {
/* Streaming dequant + f16 GEMM: the exact-q8 batched kernels only
* cover blocks <= 32 (DS4 TP shard widths); the per-token fallback
* re-reads the full weight per token (~30x the bytes at GLM dims).
* Scratch layout: [w_f16][x_f16] in one arena grab. */
const uint64_t wh_bytes = in_dim * out_dim * sizeof(__half);
const uint64_t xh_off = (wh_bytes + 255u) & ~255ull;
const uint64_t oo_off =
(xh_off + n_tok * in_dim * sizeof(__half) + 255u) & ~255ull;
const uint64_t gemm_tmp = oo_off + n_tok * out_dim * sizeof(float);
/* Scratch must live on the EXECUTING device: logical_tier is the
* out tensor's tier (0 for GLM graph buffers), and a GEMM reading
* its staged weights across PCIe costs ~20ms instead of ~0.1ms. */
int exec_tier = logical_tier;
{
int cur_dev = -1;
if (cudaGetDevice(&cur_dev) == cudaSuccess) {
for (int t = 0; t < g_n_gpus; t++) {
if (g_gpu[t].device_id == cur_dev) { exec_tier = t; break; }
}
}
}
void *tmp16 = cuda_tmp_alloc_on(exec_tier, gemm_tmp, "q8 dequant gemm");
if (tmp16) {
__half *wh = (__half *)tmp16;
__half *xh = (__half *)((char *)tmp16 + xh_off);
/* GEMM into device-local scratch, then one bulk D2D to the
* (possibly peer-mapped) out tensor: scattered peer stores
* from GEMM kernels run at <1GB/s over PCIe. */
float *olocal = (float *)((char *)tmp16 + oo_off);
const uint64_t total_blocks = out_dim * blocks;
q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256, 0, cuda_decode_stream()>>>(
wh, reinterpret_cast<const unsigned char *>(wptr),
total_blocks, (uint32_t)blocks, (uint32_t)in_dim);
const uint64_t xh_count = n_tok * in_dim;
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256, 0, cuda_decode_stream()>>>(
xh, (const float *)x->ptr, xh_count);
if (cuda_ok(cudaGetLastError(), "q8 dequant gemm staging")) {
const int gemm_trace = getenv("DS4_GLM_GEMM_TRACE") != NULL;
cudaEvent_t ev0, ev1, ev2;
if (gemm_trace) {
cudaEventCreate(&ev0); cudaEventCreate(&ev1); cudaEventCreate(&ev2);
cudaEventRecord(ev0);
}
const float alpha = 1.0f;
const float beta = 0.0f;
if (gemm_trace) cudaEventRecord(ev1);
cublasStatus_t st = cublasGemmEx(
cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T, CUBLAS_OP_N,
(int)out_dim, (int)n_tok, (int)in_dim,
&alpha,
wh, CUDA_R_16F, (int)in_dim,
xh, CUDA_R_16F, (int)in_dim,
&beta,
olocal, CUDA_R_32F, (int)out_dim,
CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
if (st == CUBLAS_STATUS_SUCCESS) {
if (!cuda_ok(cudaMemcpyAsync(out->ptr, olocal,
n_tok * out_dim * sizeof(float),
cudaMemcpyDeviceToDevice, 0),
"q8 dequant gemm out copy")) {
st = CUBLAS_STATUS_INTERNAL_ERROR;
}
}
if (gemm_trace) {
cudaEventRecord(ev2);
cudaEventSynchronize(ev2);
float stage_ms = 0, gemm_ms = 0;
cudaEventElapsedTime(&stage_ms, ev0, ev1);
cudaEventElapsedTime(&gemm_ms, ev1, ev2);
fprintf(stderr,
"ds4: gemm trace in=%llu out=%llu n=%llu stage(before)=%.2f gemm=%.2f ms\n",
(unsigned long long)in_dim, (unsigned long long)out_dim,
(unsigned long long)n_tok, stage_ms, gemm_ms);
cudaEventDestroy(ev0); cudaEventDestroy(ev1); cudaEventDestroy(ev2);
}
if (st == CUBLAS_STATUS_SUCCESS) return 1;
fprintf(stderr,
"ds4: q8 dequant gemm failed: status %d; using native path\n",
(int)st);
}
}
}
const uint64_t xq_bytes = n_tok * blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1);
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0;
if (n_tok == 1) {
matmul_q8_0_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 warp launch");
}
const bool force_decode_warp =
n_tok == 2u && g_glm_mtp_verify_mode;
if (n_tok > 1u && !force_decode_warp) {
/* T matches the reduction width of whichever reference kernel would
* have run: warp tree (32) for blocks <= 32, exact-thread tree
* otherwise. */
const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks);
const int mma_rc = cuda_q8_mma_try_launch(
(float *)out->ptr, reinterpret_cast<const unsigned char *>(wptr),
xq, xscale, in_dim, out_dim, n_tok, blocks, blocks, out_dim, mma_T);
if (mma_rc) return mma_rc > 0;
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL &&
blocks <= 32u &&
n_tok >= 8u) {
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1);
matmul_q8_0_preq_batch_warp8_tok8_kernel<<<bgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok8 warp launch");
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL &&
blocks <= 32u &&
n_tok >= 4u) {
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1);
matmul_q8_0_preq_batch_warp8_tok4_kernel<<<bgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch tok4 warp launch");
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
(blocks <= 32u || force_decode_warp)) {
if (force_decode_warp &&
getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) {
matmul_q8_0_preq_batch_warp8_tok2_kernel
<<<((unsigned)out_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(),
"matmul_q8_0 batch tok2 warp launch");
}
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1);
matmul_q8_0_preq_batch_warp8_kernel<<<bgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch warp launch");
}
const unsigned exact_threads = cuda_q8_exact_threads(blocks);
if (getenv("DS4_CUDA_NO_Q8_BATCH_EXACT_TOK2") == NULL &&
n_tok >= 2u) {
dim3 bgrid((unsigned)out_dim, ((unsigned)n_tok + 1u) / 2u, 1);
matmul_q8_0_preq_batch_tok2_exact_kernel<<<bgrid, exact_threads, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 exact tok2 launch");
}
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
matmul_q8_0_preq_kernel<<<grid, exact_threads, 0, cuda_decode_stream()>>>((float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim, out_dim, n_tok, blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 launch");
}
extern "C" int ds4_gpu_matmul_q8_0_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
return cuda_matmul_q8_0_tensor_labeled(out, model_map, model_size, weight_offset,
in_dim, out_dim, x, n_tok, "q8_0");
}
extern "C" int ds4_gpu_matmul_q8_0_top1_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *values,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t index_offset) {
if (!selected || !values || !x || !model_map ||
in_dim == 0 || out_dim == 0 || out_dim > UINT32_MAX) {
return 0;
}
const uint64_t blocks = (in_dim + 31u) / 32u;
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34u)) {
return 0;
}
const uint64_t weight_bytes = out_dim * blocks * 34u;
if (weight_bytes > model_size - weight_offset ||
x->bytes < in_dim * sizeof(float) ||
selected->bytes < sizeof(uint32_t) ||
values->bytes < sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(selected);
const char *wptr = cuda_resolve_weight_ptr(model_map,
weight_offset,
weight_bytes,
logical_tier,
"q8_0_top1");
if (!wptr) return 0;
const uint64_t xq_bytes = blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t key_offset =
(scale_offset + blocks * sizeof(float) + 7u) & ~7ull;
const uint64_t tmp_bytes = key_offset + sizeof(unsigned long long);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 top1 prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
unsigned long long *best_key =
(unsigned long long *)((char *)tmp + key_offset);
const int use_dp4a = cuda_q8_use_dp4a();
if (!cuda_ok(cudaMemsetAsync(best_key, 0, sizeof(*best_key)),
"matmul_q8_0_top1 clear")) {
return 0;
}
quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>(
xq,
xscale,
(const float *)x->ptr,
in_dim,
blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 quantize launch")) return 0;
matmul_q8_0_top1_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
best_key,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
blocks,
index_offset,
use_dp4a);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 launch")) return 0;
matmul_q8_0_top1_unpack_kernel<<<1, 1>>>(
(uint32_t *)selected->ptr,
(float *)values->ptr,
best_key);
return cuda_ok(cudaGetLastError(), "matmul_q8_0_top1 unpack launch");
}
extern "C" int ds4_gpu_matmul_q8_0_kslice_rows_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
uint64_t in_start,
uint64_t in_count,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 ||
in_count == 0 || n_tok == 0 || n_tok > 65535u) return 0;
if ((in_start % 32u) != 0 || (in_count % 32u) != 0 ||
in_start > in_dim || in_count > in_dim - in_start) return 0;
const uint64_t full_blocks = (in_dim + 31u) / 32u;
const uint64_t block_start = in_start / 32u;
const uint64_t slice_blocks = in_count / 32u;
if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0;
const uint64_t weight_bytes = out_dim * full_blocks * 34u;
if (in_count > UINT64_MAX / n_tok || out_dim > UINT64_MAX / n_tok) {
return 0;
}
if (weight_bytes > model_size - weight_offset ||
x->bytes < n_tok * in_count * sizeof(float) ||
out->bytes < n_tok * out_dim * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const unsigned char *wptr = reinterpret_cast<const unsigned char *>(
cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes,
logical_tier, "q8_0_kslice"));
if (!wptr) return 0;
const uint64_t xq_bytes = n_tok * slice_blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes =
scale_offset + n_tok * slice_blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
const dim3 qgrid((unsigned)slice_blocks, (unsigned)n_tok, 1u);
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(
xq,
xscale,
(const float *)x->ptr,
in_count,
slice_blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice quantize launch")) return 0;
const dim3 grid(((unsigned)out_dim + 7u) / 8u,
(unsigned)n_tok, 1u);
matmul_q8_0_kslice_preq_warp8_kernel<<<grid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
wptr,
xq,
xscale,
in_count,
out_dim,
full_blocks,
block_start,
slice_blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice launch");
}
extern "C" int ds4_gpu_matmul_q8_0_kslice_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
uint64_t in_start,
uint64_t in_count,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (!out_hc || !block_out || !x || !block_add || !residual_hc || !split ||
!model_map || in_dim == 0 || out_dim == 0 || in_count == 0 ||
n_embd == 0 || n_hc == 0 || out_dim != (uint64_t)n_embd) {
return 0;
}
if ((in_start % 32u) != 0 || (in_count % 32u) != 0 ||
in_start > in_dim || in_count > in_dim - in_start) return 0;
const uint64_t full_blocks = (in_dim + 31u) / 32u;
const uint64_t block_start = in_start / 32u;
const uint64_t slice_blocks = in_count / 32u;
if (weight_offset > model_size || out_dim > UINT64_MAX / (full_blocks * 34u)) return 0;
const uint64_t weight_bytes = out_dim * full_blocks * 34u;
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t split_bytes =
(uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float);
if (weight_bytes > model_size - weight_offset ||
x->bytes < in_count * sizeof(float) ||
block_out->bytes < out_dim * sizeof(float) ||
block_add->bytes < out_dim * sizeof(float) ||
residual_hc->bytes < hc_bytes ||
split->bytes < split_bytes ||
out_hc->bytes < hc_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out_hc);
const unsigned char *wptr = reinterpret_cast<const unsigned char *>(
cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes,
logical_tier, "q8_0_kslice_hc_expand_add"));
if (!wptr) return 0;
const uint64_t xq_bytes = slice_blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + slice_blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 kslice hc expand prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
quantize_q8_0_f32_kernel<<<(unsigned)slice_blocks, 32>>>(
xq,
xscale,
(const float *)x->ptr,
in_count,
slice_blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add quantize launch")) return 0;
matmul_q8_0_kslice_hc_expand_add_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
(float *)out_hc->ptr,
(float *)block_out->ptr,
(const float *)block_add->ptr,
(const float *)residual_hc->ptr,
(const float *)split->ptr,
wptr,
xq,
xscale,
in_count,
out_dim,
full_blocks,
block_start,
slice_blocks,
n_embd,
n_hc,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0_kslice_hc_expand_add launch");
}
static int cuda_matmul_q8_0_aligned_pair_try(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x) {
if (!out0 || !out1 || !model_map || !x || in_dim == 0u ||
out0_dim == 0u || out1_dim == 0u ||
(in_dim % 1024u) != 0 || (out0_dim % 128u) != 0 ||
(out1_dim % 128u) != 0 || !cuda_aligned_q8_enabled() ||
getenv("DS4_CUDA_NO_FUSED_ALIGNED_Q8_PAIR") != NULL) {
return 0;
}
const uint64_t blocks = in_dim / 32u;
if (out0_dim > UINT64_MAX / (blocks * 34u) ||
out1_dim > UINT64_MAX / (blocks * 34u)) {
return 0;
}
const uint64_t weight0_bytes = out0_dim * blocks * 34u;
const uint64_t weight1_bytes = out1_dim * blocks * 34u;
if (weight0_offset > model_size ||
weight0_bytes > model_size - weight0_offset ||
weight1_offset > model_size ||
weight1_bytes > model_size - weight1_offset ||
x->bytes < in_dim * sizeof(float) ||
out0->bytes < out0_dim * sizeof(float) ||
out1->bytes < out1_dim * sizeof(float)) {
return 0;
}
const uint64_t aligned0_bytes =
ds4_mmq_q8_0_aligned_bytes((int)out0_dim, (int)in_dim);
const uint64_t aligned1_bytes =
ds4_mmq_q8_0_aligned_bytes((int)out1_dim, (int)in_dim);
const char *aligned0 = cuda_derived_weight_ptr(
model_map, weight0_offset, weight0_bytes,
CUDA_DERIVED_Q8_0_ALIGNED_DENSE,
in_dim, out0_dim, 1u, aligned0_bytes);
const char *aligned1 = cuda_derived_weight_ptr(
model_map, weight1_offset, weight1_bytes,
CUDA_DERIVED_Q8_0_ALIGNED_DENSE,
in_dim, out1_dim, 1u, aligned1_bytes);
if (!aligned0 || !aligned1) return 0;
return ds4_mmq_q8_0_aligned_dense_vec_pair(
aligned0, aligned1, (const float *)x->ptr,
(float *)out0->ptr, (float *)out1->ptr,
(int)out0_dim, (int)out1_dim, (int)in_dim,
cuda_decode_stream()) == 0 ? 1 : -1;
}
extern "C" int ds4_gpu_matmul_q4_K_pair_decode_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x) {
if (!out0 || !out1 || !model_map || !x || in_dim == 0u ||
out_dim == 0u || in_dim % 256u != 0u || in_dim > INT_MAX ||
out_dim > INT_MAX) {
return 0;
}
const uint64_t row_bytes = (in_dim / 256u) * 144u;
if (out_dim > UINT64_MAX / row_bytes) return 0;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight0_offset > model_size ||
weight_bytes > model_size - weight0_offset ||
weight1_offset > model_size ||
weight_bytes > model_size - weight1_offset ||
x->bytes < in_dim * sizeof(float) ||
out0->bytes < out_dim * sizeof(float) ||
out1->bytes < out_dim * sizeof(float)) {
return 0;
}
const int tier = ds4_tensor_device_idx(out0);
if (tier != ds4_tensor_device_idx(out1)) return 0;
const void *w0 = cuda_resolve_weight_ptr(
model_map, weight0_offset, weight_bytes, tier, "q4_K_pair_0");
const void *w1 = cuda_resolve_weight_ptr(
model_map, weight1_offset, weight_bytes, tier, "q4_K_pair_1");
if (!w0 || !w1) return 0;
return ds4_mmq_q4_K_dense_pair_vec(
w0, w1, (const float *)x->ptr,
(float *)out0->ptr, (float *)out1->ptr,
(int)out_dim, (int)in_dim, cuda_decode_stream()) == 0;
}
extern "C" int ds4_gpu_matmul_q8_0_pair_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out0_dim == 0 || out1_dim == 0 || n_tok == 0) {
return 0;
}
const uint64_t blocks = (in_dim + 31) / 32;
if (weight0_offset > model_size || weight1_offset > model_size ||
out0_dim > UINT64_MAX / (blocks * 34) ||
out1_dim > UINT64_MAX / (blocks * 34)) {
return 0;
}
const uint64_t weight0_bytes = out0_dim * blocks * 34;
const uint64_t weight1_bytes = out1_dim * blocks * 34;
if (weight0_bytes > model_size - weight0_offset ||
weight1_bytes > model_size - weight1_offset ||
x->bytes < in_dim * sizeof(float) ||
out0->bytes < out0_dim * sizeof(float) ||
out1->bytes < out1_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out0);
const int aligned_pair = cuda_matmul_q8_0_aligned_pair_try(
out0, out1, model_map, model_size,
weight0_offset, weight1_offset,
in_dim, out0_dim, out1_dim, x);
if (aligned_pair != 0) return aligned_pair > 0;
const char *w0 = cuda_resolve_weight_ptr(model_map, weight0_offset, weight0_bytes, logical_tier, "q8_0_pair0");
const char *w1 = cuda_resolve_weight_ptr(model_map, weight1_offset, weight1_bytes, logical_tier, "q8_0_pair1");
if (!w0 || !w1) return 0;
const bool force_decode_warp =
n_tok == 2u && g_glm_mtp_verify_mode;
if (n_tok != 1 && !force_decode_warp && !g_q8_cache_suppressed &&
getenv("DS4_CUDA_Q8_PAIR_BATCH") == NULL) {
return cuda_matmul_q8_0_tensor_labeled(out0, model_map, model_size, weight0_offset,
in_dim, out0_dim, x, n_tok, "q8_0_pair0") &&
cuda_matmul_q8_0_tensor_labeled(out1, model_map, model_size, weight1_offset,
in_dim, out1_dim, x, n_tok, "q8_0_pair1");
}
const uint64_t xq_bytes = n_tok * blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 pair prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1);
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) return 0;
if (n_tok != 1) {
if (force_decode_warp &&
getenv("DS4_CUDA_GLM_VERIFY_NO_Q8_TOK2") == NULL) {
matmul_q8_0_preq_batch_warp8_tok2_kernel
<<<((unsigned)out0_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
reinterpret_cast<const unsigned char *>(w0),
xq, xscale, in_dim, out0_dim, blocks, use_dp4a);
if (!cuda_ok(cudaGetLastError(),
"matmul_q8_0 pair0 tok2 warp launch")) {
return 0;
}
matmul_q8_0_preq_batch_warp8_tok2_kernel
<<<((unsigned)out1_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w1),
xq, xscale, in_dim, out1_dim, blocks, use_dp4a);
return cuda_ok(cudaGetLastError(),
"matmul_q8_0 pair1 tok2 warp launch");
}
const uint32_t mma_T = blocks <= 32u ? 32u : cuda_q8_exact_threads(blocks);
int mma_rc = cuda_q8_mma_try_launch(
(float *)out0->ptr, reinterpret_cast<const unsigned char *>(w0),
xq, xscale, in_dim, out0_dim, n_tok, blocks, blocks, out0_dim, mma_T);
if (mma_rc < 0) return 0;
if (mma_rc > 0) {
mma_rc = cuda_q8_mma_try_launch(
(float *)out1->ptr, reinterpret_cast<const unsigned char *>(w1),
xq, xscale, in_dim, out1_dim, n_tok, blocks, blocks, out1_dim, mma_T);
if (mma_rc > 0) return 1;
return 0;
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
getenv("DS4_CUDA_NO_Q8_BATCH_TOK8") == NULL &&
blocks <= 32u &&
n_tok >= 8u) {
dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1);
matmul_q8_0_preq_batch_warp8_tok8_kernel<<<grid0, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
reinterpret_cast<const unsigned char *>(w0),
xq,
xscale,
in_dim,
out0_dim,
n_tok,
blocks,
use_dp4a);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok8 launch")) return 0;
dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 7u) / 8u, 1);
matmul_q8_0_preq_batch_warp8_tok8_kernel<<<grid1, 256, 0, cuda_decode_stream()>>>(
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out1_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok8 launch");
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
getenv("DS4_CUDA_NO_Q8_BATCH_TOK4") == NULL &&
blocks <= 32u &&
n_tok >= 4u) {
dim3 grid0(((unsigned)out0_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1);
matmul_q8_0_preq_batch_warp8_tok4_kernel<<<grid0, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
reinterpret_cast<const unsigned char *>(w0),
xq,
xscale,
in_dim,
out0_dim,
n_tok,
blocks,
use_dp4a);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch tok4 launch")) return 0;
dim3 grid1(((unsigned)out1_dim + 7u) / 8u, ((unsigned)n_tok + 3u) / 4u, 1);
matmul_q8_0_preq_batch_warp8_tok4_kernel<<<grid1, 256, 0, cuda_decode_stream()>>>(
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out1_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch tok4 launch");
}
if (getenv("DS4_CUDA_NO_Q8_BATCH_WARP") == NULL &&
blocks <= 32u) {
dim3 grid0(((unsigned)out0_dim + 7u) / 8u, (unsigned)n_tok, 1);
matmul_q8_0_preq_batch_warp8_kernel<<<grid0, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
reinterpret_cast<const unsigned char *>(w0),
xq,
xscale,
in_dim,
out0_dim,
n_tok,
blocks,
use_dp4a);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch warp launch")) return 0;
dim3 grid1(((unsigned)out1_dim + 7u) / 8u, (unsigned)n_tok, 1);
matmul_q8_0_preq_batch_warp8_kernel<<<grid1, 256, 0, cuda_decode_stream()>>>(
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out1_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch warp launch");
}
if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT") == NULL) {
const uint64_t max_out_dim = out0_dim > out1_dim ? out0_dim : out1_dim;
const unsigned exact_threads = cuda_q8_exact_threads(blocks);
if (getenv("DS4_CUDA_NO_Q8_PAIR_BATCH_EXACT_TOK2") == NULL &&
n_tok >= 2u) {
dim3 grid((unsigned)max_out_dim, ((unsigned)n_tok + 1u) / 2u, 1);
matmul_q8_0_pair_preq_batch_tok2_exact_kernel<<<grid, exact_threads, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w0),
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out0_dim,
out1_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair exact tok2 launch");
}
dim3 grid((unsigned)max_out_dim, (unsigned)n_tok, 1);
matmul_q8_0_pair_preq_batch_kernel<<<grid, exact_threads, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w0),
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out0_dim,
out1_dim,
n_tok,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair batch exact launch");
}
const unsigned exact_threads = cuda_q8_exact_threads(blocks);
dim3 grid0((unsigned)out0_dim, (unsigned)n_tok, 1);
matmul_q8_0_preq_kernel<<<grid0, exact_threads, 0, cuda_decode_stream()>>>((float *)out0->ptr,
reinterpret_cast<const unsigned char *>(w0),
xq,
xscale,
in_dim, out0_dim, n_tok, blocks,
use_dp4a);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair0 batch launch")) return 0;
dim3 grid1((unsigned)out1_dim, (unsigned)n_tok, 1);
matmul_q8_0_preq_kernel<<<grid1, exact_threads, 0, cuda_decode_stream()>>>((float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim, out1_dim, n_tok, blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair1 batch launch");
}
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w0),
reinterpret_cast<const unsigned char *>(w1),
xq,
xscale,
in_dim,
out0_dim,
out1_dim,
blocks,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch");
}
extern "C" int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
if (!out || !x || !model_map || in_dim == 0u || out_dim == 0u ||
n_rows == 0u ||
x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) ||
out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) {
return 0;
}
const uint64_t blocks = (in_dim + 31u) / 32u;
if (weight_offset > model_size ||
out_dim > UINT64_MAX / (blocks * 34u)) {
return 0;
}
const uint64_t weight_bytes = out_dim * blocks * 34u;
if (weight_bytes > model_size - weight_offset) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
if (logical_tier < 0 || logical_tier >= g_n_gpus ||
ds4_tensor_device_idx(x) != logical_tier) {
return 0;
}
const char *wptr = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, logical_tier,
"q8_0 decode rows exact");
if (!wptr) return 0;
const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes =
scale_offset + (uint64_t)n_rows * blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(
logical_tier, tmp_bytes, "q8_0 decode rows exact prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
dim3 qgrid((unsigned)blocks, n_rows, 1u);
quantize_q8_0_f32_kernel<<<qgrid, 32>>>(
xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(),
"q8_0 decode rows exact quantize launch")) {
return 0;
}
const int mma_rc = cuda_q8_mma_try_launch(
(float *)out->ptr, (const unsigned char *)wptr, xq, xscale,
in_dim, out_dim, n_rows, blocks, blocks, out_dim, 32u);
if (mma_rc) return mma_rc > 0;
dim3 grid(((unsigned)out_dim + 7u) / 8u, n_rows, 1u);
matmul_q8_0_preq_warp8_kernel<<<grid, 256>>>(
(float *)out->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq, xscale, in_dim, out_dim, blocks, cuda_q8_use_dp4a());
return cuda_ok(cudaGetLastError(),
"q8_0 decode rows exact warp launch");
}
extern "C" int ds4_gpu_matmul_q8_0_pair_decode_rows_exact_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
if (!out0 || !out1 || !x || !model_map || in_dim == 0u ||
out0_dim == 0u || out1_dim == 0u || n_rows == 0u ||
x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) ||
out0->bytes < (uint64_t)n_rows * out0_dim * sizeof(float) ||
out1->bytes < (uint64_t)n_rows * out1_dim * sizeof(float)) {
return 0;
}
const uint64_t blocks = (in_dim + 31u) / 32u;
if (weight0_offset > model_size || weight1_offset > model_size ||
out0_dim > UINT64_MAX / (blocks * 34u) ||
out1_dim > UINT64_MAX / (blocks * 34u)) {
return 0;
}
const uint64_t weight0_bytes = out0_dim * blocks * 34u;
const uint64_t weight1_bytes = out1_dim * blocks * 34u;
if (weight0_bytes > model_size - weight0_offset ||
weight1_bytes > model_size - weight1_offset) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out0);
if (logical_tier < 0 || logical_tier >= g_n_gpus ||
ds4_tensor_device_idx(out1) != logical_tier ||
ds4_tensor_device_idx(x) != logical_tier) {
return 0;
}
const char *w0 = cuda_resolve_weight_ptr(
model_map, weight0_offset, weight0_bytes, logical_tier,
"q8_0 pair decode rows exact gate");
const char *w1 = cuda_resolve_weight_ptr(
model_map, weight1_offset, weight1_bytes, logical_tier,
"q8_0 pair decode rows exact up");
if (!w0 || !w1) return 0;
const uint64_t xq_bytes = (uint64_t)n_rows * blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes =
scale_offset + (uint64_t)n_rows * blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(
logical_tier, tmp_bytes, "q8_0 pair decode rows exact prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
dim3 qgrid((unsigned)blocks, n_rows, 1u);
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(
xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(),
"q8_0 pair decode rows exact quantize launch")) {
return 0;
}
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
dim3 grid(((unsigned)max_out + 7u) / 8u, n_rows, 1u);
matmul_q8_0_pair_preq_warp8_kernel<<<grid, 256, 0, cuda_decode_stream()>>>(
(float *)out0->ptr,
(float *)out1->ptr,
reinterpret_cast<const unsigned char *>(w0),
reinterpret_cast<const unsigned char *>(w1),
xq, xscale, in_dim, out0_dim, out1_dim, blocks,
cuda_q8_use_dp4a());
return cuda_ok(cudaGetLastError(),
"q8_0 pair decode rows exact warp launch");
}
static int cuda_matmul_q8_0_hc_expand_tensor_labeled(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *block_add2,
const ds4_gpu_tensor *owned_home_slots,
const ds4_gpu_tensor *owned_peer_packed,
const ds4_gpu_tensor *owned_selected,
uint32_t owned_expert_split,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc,
const char *label) {
if (!out_hc || !block_out || !x || !residual_hc || !split || !model_map ||
in_dim == 0 || out_dim == 0 || n_embd == 0 || n_hc == 0 ||
out_dim != (uint64_t)n_embd) {
return 0;
}
const uint64_t blocks = (in_dim + 31) / 32;
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0;
const uint64_t weight_bytes = out_dim * blocks * 34;
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t split_bytes = (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float);
if (weight_bytes > model_size - weight_offset ||
x->bytes < in_dim * sizeof(float) ||
block_out->bytes < out_dim * sizeof(float) ||
residual_hc->bytes < hc_bytes ||
split->bytes < split_bytes ||
out_hc->bytes < hc_bytes ||
(block_add && block_add->bytes < out_dim * sizeof(float)) ||
(block_add2 && block_add2->bytes < out_dim * sizeof(float)) ||
((owned_home_slots || owned_peer_packed || owned_selected) &&
(!owned_home_slots || !owned_peer_packed || !owned_selected ||
owned_expert_split == 0u ||
owned_home_slots->bytes < 6u * out_dim * sizeof(float) ||
owned_peer_packed->bytes < 4u * out_dim * sizeof(float) ||
owned_selected->bytes < 6u * sizeof(int32_t)))) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out_hc);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, label ? label : "q8_0_hc_expand");
if (!wptr) return 0;
const uint64_t xq_bytes = blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "q8_0 hc expand prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32, 0, cuda_decode_stream()>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) return 0;
matmul_q8_0_hc_expand_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)out_hc->ptr,
(float *)block_out->ptr,
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
block_add2 ? (const float *)block_add2->ptr : (const float *)block_out->ptr,
owned_home_slots ? (const float *)owned_home_slots->ptr : NULL,
owned_peer_packed ? (const float *)owned_peer_packed->ptr : NULL,
owned_selected ? (const int32_t *)owned_selected->ptr : NULL,
(const float *)residual_hc->ptr,
(const float *)split->ptr,
reinterpret_cast<const unsigned char *>(wptr),
xq,
xscale,
in_dim,
out_dim,
n_embd,
n_hc,
blocks,
block_add ? 1 : 0,
block_add2 ? 1 : 0,
owned_home_slots ? 1 : 0,
owned_expert_split,
use_dp4a);
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand launch");
}
extern "C" int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
if (!out || !x || !model_map) return 0;
if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0;
uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t);
if (weight_bytes > model_size - weight_offset) return 0;
if (x->bytes < n_tok * in_dim * sizeof(float) ||
out->bytes < n_tok * out_dim * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f16");
if (!wptr) return 0;
const __half *w = (const __half *)wptr;
const int serial_f16 = getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL;
const int router_shape = in_dim == 4096u && out_dim == 256u && n_tok == 1u;
const int serial_router =
!serial_f16 &&
router_shape &&
getenv("DS4_CUDA_SERIAL_ROUTER") != NULL;
const int ordered_router =
!serial_f16 &&
!serial_router &&
n_tok == 1u &&
getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL;
const int small_out_one_token =
!serial_f16 &&
!serial_router &&
!g_quality_mode &&
n_tok == 1u &&
out_dim <= 32u &&
in_dim >= 8192u &&
getenv("DS4_CUDA_F16_SMALL_OUT") != NULL &&
getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") == NULL &&
getenv("DS4_CUDA_NO_F16_SMALL_OUT") == NULL;
if (small_out_one_token) {
matmul_f16_small_out_hx_ordered_chunks_kernel<<<(unsigned)out_dim, 32, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
w,
(const float *)x->ptr,
in_dim,
out_dim);
return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_hx_ordered_chunks launch");
}
const int small_out_batch =
!serial_f16 &&
!serial_router &&
n_tok > 1u &&
out_dim <= 32u &&
in_dim >= 4096u &&
(g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) &&
getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL;
if (small_out_batch) {
matmul_f16_small_out_batch_kernel<<<(unsigned)n_tok, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
w,
(const float *)x->ptr,
in_dim,
out_dim,
n_tok);
return cuda_ok(cudaGetLastError(), "matmul_f16_small_out_batch launch");
}
const int cublas_one_token =
n_tok == 1u &&
getenv("DS4_CUDA_NO_F16_CUBLAS_ONE") == NULL &&
(!g_quality_mode || getenv("DS4_CUDA_F16_CUBLAS_ONE") != NULL);
const int cublas_batch =
n_tok > 1u && getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL;
if (!serial_f16 && g_cublas_ready && (cublas_batch || cublas_one_token)) {
const uint64_t xh_count = n_tok * in_dim;
__half *xh = (__half *)cuda_tmp_alloc_on(logical_tier, xh_count * sizeof(__half), "f16 gemm activations");
if (!xh) return 0;
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256, 0, cuda_decode_stream()>>>(xh, (const float *)x->ptr, xh_count);
if (!cuda_ok(cudaGetLastError(), "f16 activation convert launch")) return 0;
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w,
CUDA_R_16F,
(int)in_dim,
xh,
CUDA_R_16F,
(int)in_dim,
&beta,
out->ptr,
CUDA_R_32F,
(int)out_dim,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
return cublas_ok(st, "f16 matmul");
}
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
if (serial_f16 || serial_router) {
matmul_f16_serial_kernel<<<grid, 1, 0, cuda_decode_stream()>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
return cuda_ok(cudaGetLastError(), serial_router ? "matmul_f16_router_serial launch" : "matmul_f16_serial launch");
}
if (ordered_router) {
matmul_f16_ordered_chunks_kernel<<<grid, 32, 0, cuda_decode_stream()>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
return cuda_ok(cudaGetLastError(), "matmul_f16_ordered_chunks launch");
}
matmul_f16_kernel<<<grid, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
return cuda_ok(cudaGetLastError(), "matmul_f16 launch");
}
extern "C" int ds4_gpu_matmul_f16_rms_fold_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float norm_eps) {
if (!out || !x || !model_map || !g_cublas_ready || n_tok <= 1u ||
(in_dim & 2047u) != 0u || in_dim > UINT32_MAX ||
n_tok > UINT32_MAX || getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL ||
getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") != NULL ||
weight_offset > model_size || out_dim > UINT64_MAX / in_dim) {
return 0;
}
const uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t);
if (weight_bytes > model_size - weight_offset ||
x->bytes < n_tok * in_dim * sizeof(float) ||
out->bytes < n_tok * out_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const char *wptr = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, logical_tier, "f16 rms fold");
if (!wptr) return 0;
const uint64_t xh_count = n_tok * in_dim;
__half *xh = (__half *)cuda_tmp_alloc_on(
logical_tier, xh_count * sizeof(__half), "f16 rms-fold activations");
if (!xh) return 0;
rms_norm_plain_f16_batch8_kernel<<<
(unsigned)n_tok, 256, 0, cuda_decode_stream()>>>(
xh, (const float *)x->ptr, (uint32_t)in_dim,
(uint32_t)n_tok, norm_eps);
if (!cuda_ok(cudaGetLastError(), "f16 rms-fold activation launch")) {
return 0;
}
const float alpha = 1.0f;
const float beta = 0.0f;
const cublasStatus_t st = cublasGemmEx(
cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T, CUBLAS_OP_N,
(int)out_dim, (int)n_tok, (int)in_dim,
&alpha,
(const __half *)wptr, CUDA_R_16F, (int)in_dim,
xh, CUDA_R_16F, (int)in_dim,
&beta,
out->ptr, CUDA_R_32F, (int)out_dim,
CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
return cublas_ok(st, "f16 rms-fold matmul");
}
extern "C" int ds4_gpu_matmul_f16_router_rows_exact_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
const uint64_t in_dim = 4096u;
const uint64_t out_dim = 256u;
if (!out || !x || !model_map || n_rows == 0u ||
weight_offset > model_size) {
return 0;
}
const uint64_t weight_bytes = in_dim * out_dim * sizeof(uint16_t);
if (weight_bytes > model_size - weight_offset ||
x->bytes < (uint64_t)n_rows * in_dim * sizeof(float) ||
out->bytes < (uint64_t)n_rows * out_dim * sizeof(float)) {
return 0;
}
if (n_rows == 1u) {
return ds4_gpu_matmul_f16_tensor(
out, model_map, model_size, weight_offset,
in_dim, out_dim, x, 1);
}
const int logical_tier = ds4_tensor_device_idx(out);
if (ds4_tensor_device_idx(x) != logical_tier || !g_cublas_ready) return 0;
const __half *w = (const __half *)cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, logical_tier,
"f16_router_rows_exact");
if (!w) return 0;
const uint64_t xh_count = (uint64_t)n_rows * in_dim;
__half *xh = (__half *)cuda_tmp_alloc_on(
logical_tier, xh_count * sizeof(__half),
"f16 exact router batch activations");
if (!xh) return 0;
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(
xh, (const float *)x->ptr, xh_count);
if (!cuda_ok(cudaGetLastError(),
"f16 exact router activation convert launch")) {
return 0;
}
const float alpha = 1.0f;
const float beta = 0.0f;
/* Larger batchCount values let cuBLAS select a different reduction and
* change logits. Four-row calls match the one-row decode bit for bit on
* this projection, while still replacing most per-session launches. */
uint32_t row = 0;
for (; row + 4u <= n_rows; row += 4u) {
cublasStatus_t st = cublasGemmStridedBatchedEx(
cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
1,
(int)in_dim,
&alpha,
w,
CUDA_R_16F,
(int)in_dim,
0,
xh + (uint64_t)row * in_dim,
CUDA_R_16F,
(int)in_dim,
(long long int)in_dim,
&beta,
(float *)out->ptr + (uint64_t)row * out_dim,
CUDA_R_32F,
(int)out_dim,
(long long int)out_dim,
4,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
if (!cublas_ok(st, "f16 exact router row batch")) return 0;
}
for (; row < n_rows; row++) {
ds4_gpu_tensor out_row = *out;
ds4_gpu_tensor x_row = *x;
out_row.ptr = (float *)out->ptr + (uint64_t)row * out_dim;
out_row.bytes = out_dim * sizeof(float);
x_row.ptr = (float *)x->ptr + (uint64_t)row * in_dim;
x_row.bytes = in_dim * sizeof(float);
if (!ds4_gpu_matmul_f16_tensor(
&out_row, model_map, model_size, weight_offset,
in_dim, out_dim, &x_row, 1)) {
return 0;
}
}
return 1;
}
extern "C" int ds4_gpu_matmul_f16_pair_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) {
return 0;
}
if (getenv("DS4_CUDA_NO_F16_PAIR_MATMUL") != NULL ||
getenv("DS4_CUDA_SERIAL_F16_MATMUL") != NULL ||
getenv("DS4_CUDA_SERIAL_ROUTER") != NULL ||
getenv("DS4_CUDA_NO_ORDERED_F16_MATMUL") != NULL) {
return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset,
in_dim, out_dim, x, n_tok) &&
ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset,
in_dim, out_dim, x, n_tok);
}
if (weight0_offset > model_size || weight1_offset > model_size ||
out_dim > UINT64_MAX / in_dim ||
n_tok > UINT64_MAX / in_dim ||
n_tok > UINT64_MAX / out_dim) {
return 0;
}
const uint64_t weight_bytes = out_dim * in_dim * sizeof(uint16_t);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (weight_bytes > model_size - weight0_offset ||
weight_bytes > model_size - weight1_offset ||
x->bytes < x_bytes ||
out0->bytes < out_bytes ||
out1->bytes < out_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out0);
if (ds4_tensor_device_idx(out1) != logical_tier) {
return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset,
in_dim, out_dim, x, n_tok) &&
ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset,
in_dim, out_dim, x, n_tok);
}
const __half *w0 = (const __half *)cuda_resolve_weight_ptr(model_map, weight0_offset, weight_bytes, logical_tier, "f16_pair0");
const __half *w1 = (const __half *)cuda_resolve_weight_ptr(model_map, weight1_offset, weight_bytes, logical_tier, "f16_pair1");
if (!w0 || !w1) return 0;
if (n_tok > 1) {
const bool small_out_batch_requested =
out_dim <= 32u &&
in_dim >= 4096u &&
(g_quality_mode || getenv("DS4_CUDA_F16_SMALL_BATCH") != NULL) &&
getenv("DS4_CUDA_NO_F16_SMALL_BATCH") == NULL;
if (!small_out_batch_requested &&
g_cublas_ready &&
getenv("DS4_CUDA_NO_F16_CUBLAS_BATCH") == NULL) {
const uint64_t xh_count = n_tok * in_dim;
__half *xh = (__half *)cuda_tmp_alloc_on(logical_tier,
xh_count * sizeof(__half),
"f16 pair gemm activations");
if (!xh) return 0;
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(
xh, (const float *)x->ptr, xh_count);
if (!cuda_ok(cudaGetLastError(), "f16 pair activation convert launch")) return 0;
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasGemmEx(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w0,
CUDA_R_16F,
(int)in_dim,
xh,
CUDA_R_16F,
(int)in_dim,
&beta,
out0->ptr,
CUDA_R_32F,
(int)out_dim,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
if (!cublas_ok(st, "f16 pair matmul0")) return 0;
st = cublasGemmEx(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w1,
CUDA_R_16F,
(int)in_dim,
xh,
CUDA_R_16F,
(int)in_dim,
&beta,
out1->ptr,
CUDA_R_32F,
(int)out_dim,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
return cublas_ok(st, "f16 pair matmul1");
}
return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset,
in_dim, out_dim, x, n_tok) &&
ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset,
in_dim, out_dim, x, n_tok);
}
matmul_f16_pair_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>(
(float *)out0->ptr,
(float *)out1->ptr,
w0,
w1,
(const float *)x->ptr,
in_dim,
out_dim,
out_dim);
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
}
extern "C" int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
(void)out_kv;
(void)out_score;
(void)state_kv;
(void)state_score;
(void)model_map;
(void)model_size;
(void)weight_kv_offset;
(void)weight_score_offset;
(void)ape_offset;
(void)ape_type;
(void)in_dim;
(void)width;
(void)x;
(void)ratio;
(void)pos;
return 0;
}
extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0) return 0;
if (weight_offset > model_size || out_dim > UINT64_MAX / in_dim) return 0;
uint64_t weight_elems = out_dim * in_dim;
if (weight_elems > UINT64_MAX / sizeof(float)) return 0;
uint64_t weight_bytes = weight_elems * sizeof(float);
if (weight_bytes > model_size - weight_offset) return 0;
if (x->bytes < n_tok * in_dim * sizeof(float) ||
out->bytes < n_tok * out_dim * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "f32");
if (!wptr) return 0;
const float *w = (const float *)wptr;
if (g_cublas_ready && n_tok > 1) {
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasSgemm(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)out_dim,
(int)n_tok,
(int)in_dim,
&alpha,
w,
(int)in_dim,
(const float *)x->ptr,
(int)in_dim,
&beta,
(float *)out->ptr,
(int)out_dim);
return cublas_ok(st, "f32 matmul");
}
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
matmul_f32_kernel<<<grid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
w,
(const float *)x->ptr,
in_dim,
out_dim,
n_tok);
return cuda_ok(cudaGetLastError(), "matmul_f32 launch");
}
extern "C" int ds4_gpu_repeat_hc_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *row, uint32_t n_embd, uint32_t n_hc) {
if (!out || !row || n_embd == 0 || n_hc == 0 ||
row->bytes < (uint64_t)n_embd * sizeof(float) ||
out->bytes < (uint64_t)n_embd * n_hc * sizeof(float)) {
return 0;
}
uint64_t n = (uint64_t)n_embd * n_hc;
repeat_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)row->ptr, n_embd, n_hc);
return cuda_ok(cudaGetLastError(), "repeat_hc launch");
}
/* Non-causal batch attention over a raw KV ring for the DSpark draft block.
* Every query row attends over all n_raw visible rows plus the per-head sink,
* with the same exact one-block max/denominator/value accumulation order as
* the reference decode attention (scores in shared, sequential value pass). */
__global__ static void attention_noncausal_raw_batch_heads_kernel(
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t tok = blockIdx.x;
const uint32_t h = blockIdx.y;
if (tok >= n_tokens || h >= n_head) return;
extern __shared__ float sh_scores[]; /* n_raw floats */
const float *qh = q + ((uint64_t)tok * n_head + h) * head_dim;
const float scale = rsqrtf((float)head_dim);
for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) {
const uint32_t row = (raw_start + r) % raw_cap;
const float *kv = raw_kv + (uint64_t)row * head_dim;
float dot = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) dot += qh[d] * kv[d];
sh_scores[r] = dot * scale;
}
__syncthreads();
__shared__ float partial[256];
__shared__ float max_s;
__shared__ float denom;
float local_max = sinks[h];
for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) {
local_max = fmaxf(local_max, sh_scores[r]);
}
partial[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
partial[threadIdx.x] = fmaxf(partial[threadIdx.x], partial[threadIdx.x + stride]);
}
__syncthreads();
}
if (threadIdx.x == 0) max_s = partial[0];
__syncthreads();
float den_local = 0.0f;
for (uint32_t r = threadIdx.x; r < n_raw; r += blockDim.x) {
sh_scores[r] = expf(sh_scores[r] - max_s);
den_local += sh_scores[r];
}
partial[threadIdx.x] = den_local;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) denom = partial[0] + expf(sinks[h] - max_s);
__syncthreads();
float *oh = heads + ((uint64_t)tok * n_head + h) * head_dim;
for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) {
float acc = 0.0f;
for (uint32_t r = 0; r < n_raw; r++) {
const uint32_t row = (raw_start + r) % raw_cap;
acc += raw_kv[(uint64_t)row * head_dim + d] * sh_scores[r];
}
oh[d] = acc / denom;
}
}
extern "C" int ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_head,
uint32_t head_dim) {
if (!heads || !q || !raw_kv || !model_map ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw ||
raw_start >= raw_cap || n_head == 0 || head_dim == 0 ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier,
"dspark_attn_sinks");
if (!sinks) return 0;
const size_t shmem = (size_t)n_raw * sizeof(float);
if (shmem > 32768) return 0; /* draft blocks are tiny; guard anyway */
dim3 grid(n_tokens, n_head, 1);
attention_noncausal_raw_batch_heads_kernel<<<grid, 256, shmem>>>(
(float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_tokens, n_raw, raw_cap, raw_start, n_head, head_dim);
if (!cuda_ok(cudaGetLastError(), "attention noncausal raw batch heads launch")) return 0;
static int verify_left = -1;
if (verify_left < 0) {
verify_left = getenv("DS4_DSPARK_VERIFY_NONCAUSAL") != NULL ? 3 : 0;
}
if (verify_left > 0) {
verify_left--;
(void)cudaDeviceSynchronize();
const uint64_t qn = (uint64_t)n_tokens * n_head * head_dim;
const uint64_t kn = (uint64_t)raw_cap * head_dim;
std::vector<float> hq(qn), hkv(kn), hout(qn), hsink(n_head);
(void)cudaMemcpy(hq.data(), q->ptr, qn * 4, cudaMemcpyDeviceToHost);
(void)cudaMemcpy(hkv.data(), raw_kv->ptr, kn * 4, cudaMemcpyDeviceToHost);
(void)cudaMemcpy(hout.data(), heads->ptr, qn * 4, cudaMemcpyDeviceToHost);
(void)cudaMemcpy(hsink.data(), sinks, (uint64_t)n_head * 4, cudaMemcpyDeviceToHost);
double max_abs = 0.0, max_rel = 0.0;
const double scale = 1.0 / sqrt((double)head_dim);
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t h = 0; h < n_head; h++) {
std::vector<double> sc(n_raw);
double mx = (double)hsink[h];
for (uint32_t r = 0; r < n_raw; r++) {
const uint32_t row = (raw_start + r) % raw_cap;
double dot = 0.0;
for (uint32_t d = 0; d < head_dim; d++) {
dot += (double)hq[((uint64_t)t * n_head + h) * head_dim + d] *
(double)hkv[(uint64_t)row * head_dim + d];
}
sc[r] = dot * scale;
if (sc[r] > mx) mx = sc[r];
}
double den = exp((double)hsink[h] - mx);
for (uint32_t r = 0; r < n_raw; r++) den += exp(sc[r] - mx);
for (uint32_t d = 0; d < head_dim; d++) {
double acc = 0.0;
for (uint32_t r = 0; r < n_raw; r++) {
const uint32_t row = (raw_start + r) % raw_cap;
acc += exp(sc[r] - mx) * (double)hkv[(uint64_t)row * head_dim + d];
}
const double ref = acc / den;
const double got = (double)hout[((uint64_t)t * n_head + h) * head_dim + d];
const double ad = fabs(ref - got);
if (ad > max_abs) max_abs = ad;
if (fabs(ref) > 1e-3 && ad / fabs(ref) > max_rel) max_rel = ad / fabs(ref);
}
}
}
fprintf(stderr,
"ds4: DSpark noncausal verify n_tok=%u n_raw=%u start=%u cap=%u "
"max_abs=%.3e max_rel=%.3e\n",
n_tokens, n_raw, raw_start, raw_cap, max_abs, max_rel);
}
return 1;
}
extern "C" int ds4_gpu_repeat_hc_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *rows, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) {
uint64_t rows_elems = 0;
uint64_t out_elems = 0;
if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0 ||
(uint64_t)n_tokens > UINT64_MAX / n_embd ||
(rows_elems = (uint64_t)n_tokens * n_embd) > UINT64_MAX / n_hc ||
(out_elems = rows_elems * n_hc) > UINT64_MAX / sizeof(float) ||
rows_elems > UINT64_MAX / sizeof(float) ||
rows->bytes < rows_elems * sizeof(float) ||
out->bytes < out_elems * sizeof(float)) {
return 0;
}
const uint64_t blocks = (out_elems + 255u) / 256u;
if (blocks > UINT32_MAX) return 0;
repeat_hc_rows_kernel<<<(unsigned)blocks, 256>>>((float *)out->ptr, (const float *)rows->ptr, n_tokens, n_embd, n_hc);
return cuda_ok(cudaGetLastError(), "repeat_hc_rows launch");
}
extern "C" int ds4_gpu_rms_norm_plain_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, float eps) {
if (!out || !x || out->bytes < (uint64_t)n * sizeof(float) ||
x->bytes < (uint64_t)n * sizeof(float)) return 0;
if (n == 4096u) {
rms_norm_plain_fast4096_kernel<<<1, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps);
} else if ((n & 2047u) == 0u) {
rms_norm_plain_batch8_kernel<<<1, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps);
} else {
rms_norm_plain_kernel<<<1, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps);
}
return cuda_ok(cudaGetLastError(), "rms_norm_plain launch");
}
extern "C" int ds4_gpu_rms_norm_plain_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, uint32_t rows, float eps) {
if (!out || !x || out->bytes < (uint64_t)n * rows * sizeof(float) ||
x->bytes < (uint64_t)n * rows * sizeof(float)) return 0;
if (n == 4096u) {
rms_norm_plain_fast4096_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps);
} else if ((n & 2047u) == 0u) {
rms_norm_plain_batch8_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps);
} else {
rms_norm_plain_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps);
}
return cuda_ok(cudaGetLastError(), "rms_norm_plain launch");
}
extern "C" int ds4_gpu_rms_norm_weight_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, float eps) {
if (!out || !x || !model_map || weight_offset > model_size ||
model_size - weight_offset < (uint64_t)n * sizeof(float) ||
out->bytes < (uint64_t)n * sizeof(float) ||
x->bytes < (uint64_t)n * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight");
if (!wptr) return 0;
const float *w = (const float *)wptr;
rms_norm_weight_kernel<<<1, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)x->ptr, w, n, 1, eps);
return cuda_ok(cudaGetLastError(), "rms_norm_weight launch");
}
extern "C" int ds4_gpu_rms_norm_weight_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, uint32_t rows, float eps) {
if (!out || !x || !model_map || weight_offset > model_size ||
model_size - weight_offset < (uint64_t)n * sizeof(float) ||
out->bytes < (uint64_t)n * rows * sizeof(float) ||
x->bytes < (uint64_t)n * rows * sizeof(float)) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, (uint64_t)n * sizeof(float), logical_tier, "rms_weight");
if (!wptr) return 0;
const float *w = (const float *)wptr;
rms_norm_weight_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, w, n, rows, eps);
return cuda_ok(cudaGetLastError(), "rms_norm_weight launch");
}
extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
uint32_t rows,
float eps) {
if (!g_cuda_disable_qkv_rms_fused) {
if (!q_out || !q || !kv_out || !kv || !model_map ||
q_weight_offset > model_size ||
kv_weight_offset > model_size ||
model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) ||
model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) ||
q_out->bytes < (uint64_t)q_n * rows * sizeof(float) ||
q->bytes < (uint64_t)q_n * rows * sizeof(float) ||
kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) ||
kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(q_out);
const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map,
q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight");
const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map,
kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight");
if (!q_w || !kv_w) return 0;
dim3 grid(rows, 2u, 1u);
dsv4_qkv_rms_norm_rows_kernel<<<grid, 256>>>(
(float *)q_out->ptr,
(const float *)q->ptr,
q_w,
q_n,
(float *)kv_out->ptr,
(const float *)kv->ptr,
kv_w,
kv_n,
rows,
eps);
return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm rows launch");
}
return ds4_gpu_rms_norm_weight_rows_tensor(q_out, q, model_map, model_size,
q_weight_offset, q_n, rows, eps) &&
ds4_gpu_rms_norm_weight_rows_tensor(kv_out, kv, model_map, model_size,
kv_weight_offset, kv_n, rows, eps);
}
extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_kv_rope_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
uint32_t rows,
uint32_t kv_n_head,
uint32_t kv_head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
if (g_cuda_disable_qkv_rms_fused) return 0;
if (!q_out || !q || !kv_out || !kv || !model_map ||
q_weight_offset > model_size ||
kv_weight_offset > model_size ||
kv_n_head == 0 || kv_head_dim == 0 ||
n_rot > kv_head_dim || (n_rot & 1u) ||
kv_n != kv_n_head * kv_head_dim ||
model_size - q_weight_offset < (uint64_t)q_n * sizeof(float) ||
model_size - kv_weight_offset < (uint64_t)kv_n * sizeof(float) ||
q_out->bytes < (uint64_t)q_n * rows * sizeof(float) ||
q->bytes < (uint64_t)q_n * rows * sizeof(float) ||
kv_out->bytes < (uint64_t)kv_n * rows * sizeof(float) ||
kv->bytes < (uint64_t)kv_n * rows * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(q_out);
const float *q_w = (const float *)cuda_resolve_weight_ptr(model_map,
q_weight_offset, (uint64_t)q_n * sizeof(float), logical_tier, "q_rms_weight");
const float *kv_w = (const float *)cuda_resolve_weight_ptr(model_map,
kv_weight_offset, (uint64_t)kv_n * sizeof(float), logical_tier, "kv_rms_weight");
if (!q_w || !kv_w) return 0;
dim3 grid(rows, 2u, 1u);
dsv4_qkv_rms_norm_rows_kv_rope_kernel<<<grid, 256>>>(
(float *)q_out->ptr,
(const float *)q->ptr,
q_w,
q_n,
(float *)kv_out->ptr,
(const float *)kv->ptr,
kv_w,
kv_n,
rows,
kv_n_head,
kv_head_dim,
n_rot,
pos0,
n_ctx_orig,
inverse ? 1 : 0,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow,
eps);
return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm kv rope launch");
}
extern "C" int ds4_gpu_head_rms_norm_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
if (!x || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0;
head_rms_norm_kernel<<<n_tok * n_head, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, eps);
return cuda_ok(cudaGetLastError(), "head_rms_norm launch");
}
extern "C" int ds4_gpu_head_rms_norm_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, float eps) {
if (!x || n_rot > head_dim || (n_rot & 1u) ||
x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0;
head_rms_norm_rope_tail_kernel<<<n_tok * n_head, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, eps);
return cuda_ok(cudaGetLastError(), "head_rms_norm_rope_tail launch");
}
extern "C" int ds4_gpu_dsv4_fp8_kv_quantize_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) {
if (!x || n_rot > head_dim || x->bytes < (uint64_t)n_tok * head_dim * sizeof(float)) return 0;
fp8_kv_quantize_kernel<<<n_tok, 64>>>((float *)x->ptr, n_tok, head_dim, n_rot);
return cuda_ok(cudaGetLastError(), "fp8_kv_quantize launch");
}
extern "C" int ds4_gpu_dsv4_indexer_qat_tensor(ds4_gpu_tensor *x, uint32_t n_rows, uint32_t head_dim) {
if (!x || n_rows == 0 || head_dim != 128u ||
x->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) {
return 0;
}
indexer_hadamard_fp4_kernel<<<n_rows, 128>>>((float *)x->ptr, n_rows, head_dim);
return cuda_ok(cudaGetLastError(), "indexer_hadamard_fp4 launch");
}
extern "C" int ds4_gpu_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) {
if (!x || n_rot > head_dim || (n_rot & 1) || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0;
uint32_t pairs = n_tok * n_head * (n_rot / 2);
rope_tail_kernel<<<(pairs + 255) / 256, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, 1, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
return cuda_ok(cudaGetLastError(), "rope_tail launch");
}
extern "C" int ds4_gpu_rope_tail_decode_rows_tensor(
ds4_gpu_tensor *x,
const ds4_gpu_attention_decode_row *rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!x || !rows || n_rows == 0u ||
n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u ||
n_rot == 0u || n_rot > head_dim || (n_rot & 1u) != 0u ||
x->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) {
return 0;
}
cuda_attention_decode_row_table table;
memset(&table, 0, sizeof(table));
for (uint32_t i = 0; i < n_rows; i++) table.row[i].pos = rows[i].pos;
const uint32_t pairs = n_rows * n_head * (n_rot / 2u);
rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>(
(float *)x->ptr, table, n_rows, n_head, head_dim, n_rot,
n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
return cuda_ok(cudaGetLastError(), "rope tail decode rows launch");
}
extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim);
extern "C" int ds4_gpu_kv_fp8_store_raw_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *raw_cache,
uint32_t raw_cap,
uint32_t raw_row,
uint32_t head_dim,
uint32_t n_rot) {
if (!kv || !raw_cache || raw_cap == 0u || n_rot > head_dim ||
kv->device_id != raw_cache->device_id ||
kv->bytes < (uint64_t)head_dim * sizeof(float) ||
raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float)) {
return 0;
}
cuda_attention_decode_row_table table;
memset(&table, 0, sizeof(table));
table.row[0].raw_kv = (uint64_t)(uintptr_t)raw_cache->ptr;
table.row[0].raw_cap = raw_cap;
table.row[0].raw_start = raw_row % raw_cap;
fp8_kv_quantize_store_rows_kernel<<<1, 64>>>(
(float *)kv->ptr, table, 1u, head_dim, n_rot);
return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store launch");
}
extern "C" int ds4_gpu_kv_fp8_store_raw_decode_rows_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *const *raw_caches,
const uint32_t *raw_caps,
const uint32_t *raw_rows,
uint32_t n_rows,
uint32_t head_dim,
uint32_t n_rot) {
if (!kv || !raw_caches || !raw_caps || !raw_rows || n_rows == 0u ||
n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_rot > head_dim ||
kv->bytes < (uint64_t)n_rows * head_dim * sizeof(float)) {
return 0;
}
cuda_attention_decode_row_table table;
memset(&table, 0, sizeof(table));
for (uint32_t i = 0; i < n_rows; i++) {
const ds4_gpu_tensor *raw = raw_caches[i];
if (!raw || raw_caps[i] == 0u || raw_rows[i] >= raw_caps[i] ||
raw->device_id != kv->device_id ||
raw->bytes < (uint64_t)raw_caps[i] * head_dim * sizeof(float)) {
return 0;
}
table.row[i].raw_kv = (uint64_t)(uintptr_t)raw->ptr;
table.row[i].raw_cap = raw_caps[i];
table.row[i].raw_start = raw_rows[i];
}
fp8_kv_quantize_store_rows_kernel<<<n_rows, 64>>>(
(float *)kv->ptr, table, n_rows, head_dim, n_rot);
return cuda_ok(cudaGetLastError(), "fp8 KV quantize/store rows launch");
}
extern "C" int ds4_gpu_store_raw_kv_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t row, uint32_t head_dim) {
if (!raw_cache || !kv || raw_cap == 0 ||
raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
kv->bytes < (uint64_t)head_dim * sizeof(float)) return 0;
store_raw_kv_batch_kernel<<<(head_dim + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, row, 1, head_dim);
return cuda_ok(cudaGetLastError(), "store_raw_kv launch");
}
extern "C" int ds4_gpu_store_raw_kv_batch_tensor(ds4_gpu_tensor *raw_cache, const ds4_gpu_tensor *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) {
if (!raw_cache || !kv || raw_cap == 0 ||
raw_cache->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float)) return 0;
uint64_t n = (uint64_t)n_tokens * head_dim;
store_raw_kv_batch_kernel<<<(n + 255) / 256, 256>>>((float *)raw_cache->ptr, (const float *)kv->ptr, raw_cap, pos0, n_tokens, head_dim);
return cuda_ok(cudaGetLastError(), "store_raw_kv_batch launch");
}
extern "C" int ds4_gpu_compressor_store_batch_tensor(
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens) {
if (!kv || !sc || !state_kv || !state_score || !model_map ||
head_dim == 0 || ratio == 0 || n_tokens == 0 ||
(ape_type != 0u && ape_type != 1u)) {
return 0;
}
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
kv->bytes < kv_bytes || sc->bytes < kv_bytes ||
state_kv->bytes < state_bytes || state_score->bytes < state_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(state_kv);
const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape");
if (!ape) return 0;
uint64_t n = (uint64_t)n_tokens * width;
compressor_store_kernel<<<(n + 255) / 256, 256>>>(
(const float *)kv->ptr,
(const float *)sc->ptr,
(float *)state_kv->ptr,
(float *)state_score->ptr,
ape,
0,
ape_type,
head_dim,
ratio,
pos0,
n_tokens);
return cuda_ok(cudaGetLastError(), "compressor store launch");
}
extern "C" int ds4_gpu_compressor_update_tensor(
const ds4_gpu_tensor *kv_cur,
const ds4_gpu_tensor *sc_cur,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
ds4_gpu_tensor *comp_cache,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos,
uint32_t comp_row,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps,
bool state_already_stored,
bool decode_one_token,
bool defer_finalize) {
(void)decode_one_token;
(void)defer_finalize;
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
!model_map || head_dim == 0 || ratio == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) || norm_type != 0u) {
return 0;
}
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset ||
kv_cur->bytes < kv_bytes || sc_cur->bytes < kv_bytes ||
state_kv->bytes < state_bytes || state_score->bytes < state_bytes ||
(emit && comp_cache->bytes < comp_bytes)) {
return 0;
}
if (!state_already_stored) {
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
model_map, model_size, ape_offset, ape_type,
head_dim, ratio, pos, 1)) {
return 0;
}
}
if (!emit) return 1;
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
comp_cache,
(uint64_t)comp_row * head_dim * sizeof(float),
(uint64_t)head_dim * sizeof(float));
if (!comp_row_view) return 0;
compressor_update_pool_kernel<<<(head_dim + 255) / 256, 256>>>(
(float *)comp_row_view->ptr,
(const float *)state_kv->ptr,
(const float *)state_score->ptr,
head_dim,
ratio);
int ok = cuda_ok(cudaGetLastError(), "compressor update pool launch");
if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view,
model_map, model_size, norm_offset,
head_dim, 1, rms_eps);
if (ok) ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, 1, head_dim, n_rot,
pos + 1u - ratio, n_ctx_orig, false,
freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
ds4_gpu_tensor_free(comp_row_view);
if (ok && ratio == 4u) {
uint64_t half = 4ull * width;
compressor_shift_ratio4_kernel<<<(half + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr, width);
ok = cuda_ok(cudaGetLastError(), "compressor ratio4 shift launch");
}
return ok;
}
extern "C" int ds4_gpu_compressor_prefill_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
head_dim == 0 || ratio == 0 || n_tokens == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) || norm_type != 0u) {
return 0;
}
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint32_t n_comp = n_tokens / ratio;
const uint32_t cutoff = n_comp * ratio;
const uint32_t rem = n_tokens - cutoff;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset ||
kv->bytes < kv_bytes || sc->bytes < kv_bytes ||
state_kv->bytes < state_bytes || state_score->bytes < state_bytes ||
(n_comp && comp_cache->bytes < comp_bytes)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(state_kv);
const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape");
if (!ape) return 0;
uint64_t state_n = (uint64_t)state_rows * width;
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
"compressor state kv zero")) return 0;
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
if (ratio == 4u) {
if (cutoff >= ratio) {
uint32_t prev_start = cutoff - ratio;
uint64_t n = (uint64_t)ratio * width;
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr,
(const float *)kv->ptr, (const float *)sc->ptr,
ape, 0, ape_type, width, ratio, pos0,
prev_start, 0, ratio);
if (!cuda_ok(cudaGetLastError(), "compressor prefill prev state launch")) return 0;
}
if (rem != 0) {
uint64_t n = (uint64_t)rem * width;
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr,
(const float *)kv->ptr, (const float *)sc->ptr,
ape, 0, ape_type, width, ratio, pos0,
cutoff, ratio, rem);
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
}
} else if (rem != 0) {
uint64_t n = (uint64_t)rem * width;
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr,
(const float *)kv->ptr, (const float *)sc->ptr,
ape, 0, ape_type, width, ratio, pos0,
cutoff, 0, rem);
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
}
if (n_comp != 0) {
dim3 grid((head_dim + 255) / 256, n_comp, 1);
compressor_prefill_pool_kernel<<<grid, 256>>>(
(float *)comp_cache->ptr,
(const float *)kv->ptr,
(const float *)sc->ptr,
(const float *)state_kv->ptr,
(const float *)state_score->ptr,
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 0);
if (!cuda_ok(cudaGetLastError(), "compressor prefill pool launch")) return 0;
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
model_map, model_size, norm_offset,
head_dim, n_comp, rms_eps)) return 0;
if (n_rot != 0) {
const uint32_t pairs = n_comp * (n_rot / 2u);
rope_tail_kernel<<<(pairs + 255) / 256, 256>>>(
(float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot,
pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
if (!cuda_ok(cudaGetLastError(), "compressor prefill rope launch")) return 0;
}
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
}
return 1;
}
extern "C" int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) || norm_type != 0u) {
return 0;
}
const uint32_t ratio = 4u;
const uint32_t width = 2u * head_dim;
const uint32_t state_rows = 8u;
const uint32_t n_comp = n_tokens / ratio;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset ||
kv->bytes < kv_bytes || sc->bytes < kv_bytes ||
state_kv->bytes < state_bytes || state_score->bytes < state_bytes ||
comp_cache->bytes < comp_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(comp_cache);
const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape");
if (!ape) return 0;
dim3 grid((head_dim + 255) / 256, n_comp, 1);
compressor_prefill_pool_kernel<<<grid, 256>>>(
(float *)comp_cache->ptr,
(const float *)kv->ptr,
(const float *)sc->ptr,
(const float *)state_kv->ptr,
(const float *)state_score->ptr,
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 1);
if (!cuda_ok(cudaGetLastError(), "compressor replay pool launch")) return 0;
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
model_map, model_size, norm_offset,
head_dim, n_comp, rms_eps)) return 0;
if (n_rot != 0) {
const uint32_t pairs = n_comp * (n_rot / 2u);
rope_tail_kernel<<<(pairs + 255) / 256, 256>>>(
(float *)comp_cache->ptr, n_comp, 1, head_dim, n_rot,
pos0, ratio, n_ctx_orig, 0, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
if (!cuda_ok(cudaGetLastError(), "compressor replay rope launch")) return 0;
}
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
uint64_t state_n = (uint64_t)state_rows * width;
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
"compressor replay state kv zero")) return 0;
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
if (!cuda_ok(cudaGetLastError(), "compressor replay state score fill launch")) return 0;
uint32_t prev_start = n_tokens - ratio;
uint64_t n = (uint64_t)ratio * width;
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr,
(const float *)kv->ptr, (const float *)sc->ptr,
ape, 0, ape_type, width, ratio, pos0,
prev_start, 0, ratio);
return cuda_ok(cudaGetLastError(), "compressor replay state launch");
}
extern "C" int ds4_gpu_compressor_prefill_state_ratio4_tensor(
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv_tail,
const ds4_gpu_tensor *sc_tail,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t pos0) {
if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map ||
head_dim == 0 || (ape_type != 0u && ape_type != 1u)) {
return 0;
}
const uint32_t ratio = 4u;
const uint32_t width = 2u * head_dim;
const uint32_t state_rows = 8u;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape;
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
kv_tail->bytes < tail_bytes || sc_tail->bytes < tail_bytes ||
state_kv->bytes < state_bytes || state_score->bytes < state_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(state_kv);
const char *ape = cuda_resolve_weight_ptr(model_map, ape_offset, ape_bytes, logical_tier, "compressor_ape");
if (!ape) return 0;
uint64_t state_n = (uint64_t)state_rows * width;
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
"compressor state kv zero")) return 0;
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
uint64_t n = (uint64_t)ratio * width;
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
(float *)state_kv->ptr, (float *)state_score->ptr,
(const float *)kv_tail->ptr, (const float *)sc_tail->ptr,
ape, 0, ape_type, width, ratio, pos0,
0, 0, ratio);
return cuda_ok(cudaGetLastError(), "compressor state set launch");
}
/* perf-02 split-KV / flash-decode launch helper (opt-in, default OFF).
* Returns 1 if the split path handled the launch, 0 if the caller should fall
* through to the existing attention_decode_mixed_kernel path.
*
* Engages only for the single-token decode shape (n_tokens==1) and only when
* DS4_CUDA_SPLITKV_DECODE is set. S==1 is NOT handled here: the caller dispatches
* the old kernel as the bit-exact anchor when S would be 1.
*/
static int attention_decode_splitkv_launch(
int logical_tier,
float *heads,
const float *sinks,
const float *q,
const float *raw_kv,
const float *comp_kv,
const float *comp_mask,
uint32_t use_comp_mask,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
/* n_tokens is fixed at 1 for the split path; compute the EXACT logical row
* count the kernel will use (raw_count + visible_comp) so S is sized to the
* real work. raw_count MUST apply the same window logic as the kernel /
* reference, otherwise a true-S==1 case (e.g. ratio=1, window=1, n_raw>=2,
* n_comp=0) could be over-estimated to S>1 and engage split-KV instead of
* the bit-exact old-kernel anchor. The count is head-independent. */
const bool single_all = (ratio == 0u);
uint32_t qpos = pos0; /* t==0, n_tokens==1 */
uint32_t first_raw_pos = pos0 + 1u - n_raw;
uint32_t visible_comp = single_all ? n_comp : (n_comp ? (qpos + 1u) / ratio : 0u);
if (visible_comp > n_comp) visible_comp = n_comp;
uint32_t raw_count = 0;
uint32_t raw_first_idx = 0;
if (n_raw != 0) {
const uint32_t raw_last_pos = first_raw_pos + n_raw - 1u;
if (single_all) {
raw_count = n_raw > 256u ? 256u : n_raw;
} else if (qpos >= first_raw_pos) {
uint32_t lo = first_raw_pos;
if (window != 0 && qpos + 1u > window) {
const uint32_t wlo = qpos + 1u - window;
if (wlo > lo) lo = wlo;
}
const uint32_t hi = qpos < raw_last_pos ? qpos : raw_last_pos;
if (hi >= lo) {
raw_first_idx = lo - first_raw_pos;
raw_count = hi - lo + 1u;
if (raw_count > 256u) raw_count = 256u;
}
}
}
uint32_t n_score = raw_count + visible_comp;
if (n_score == 0u) return 0; /* nothing to do; let old path handle it */
const int manual_splitkv = cuda_env_flag_enabled("DS4_CUDA_SPLITKV_DECODE", 0);
const uint32_t scoped_min_score =
(g_decode_fast_attention && !manual_splitkv) ? 512u : 0u;
uint32_t min_score = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_MIN_SCORE",
scoped_min_score, 0u,
DS4_CUDA_ATTENTION_SCORE_CAP,
NULL);
if (n_score < min_score) return 0;
/* S = clamp(ceil(n_score / CHUNK), 1, S_MAX); raise to S_FLOOR for short
* context to fill more SMs, but never exceed n_score (no empty chunks).
* Optional tuning knobs are guarded by min_needed so every block's chunk
* still fits the fixed shared score buffer. */
const uint32_t split_cap = DS4_CUDA_SPLITKV_SCORE_CAP;
const uint32_t min_needed = (n_score + split_cap - 1u) / split_cap;
uint32_t chunk = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_CHUNK",
DS4_CUDA_SPLITKV_CHUNK,
1u, split_cap, NULL);
uint32_t s_floor = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_FLOOR",
DS4_CUDA_SPLITKV_S_FLOOR,
1u, DS4_CUDA_SPLITKV_S_MAX, NULL);
uint32_t s_max = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S_MAX",
DS4_CUDA_SPLITKV_S_MAX,
1u, DS4_CUDA_SPLITKV_S_MAX, NULL);
int exact_present = 0;
uint32_t S = cuda_parse_u32_env_clamped("DS4_CUDA_SPLITKV_S",
0u, 1u, DS4_CUDA_SPLITKV_S_MAX,
&exact_present);
if (!exact_present) {
S = (n_score + chunk - 1u) / chunk;
if (S < s_floor) S = s_floor < n_score ? s_floor : n_score;
if (S > s_max) S = s_max;
}
if (S < min_needed) S = min_needed;
if (S > n_score) S = n_score;
if (S <= 1u) return 0; /* S==1: caller uses the old kernel anchor */
if (cuda_env_flag_enabled("DS4_CUDA_SPLITKV_GLOBAL_SOFTMAX", 0)) {
const uint64_t score_count = (uint64_t)n_head * n_score;
const uint64_t score_bytes = score_count * sizeof(float);
const uint64_t denom_offset = (score_bytes + 255u) & ~255ull;
const uint64_t denom_bytes = (uint64_t)n_head * sizeof(float);
const uint64_t partial_offset = (denom_offset + denom_bytes + 255u) & ~255ull;
const uint64_t partial_bytes = (uint64_t)n_head * S * head_dim * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier,
partial_offset + partial_bytes,
"attention splitkv global softmax");
if (!tmp) return 0;
float *scores = (float *)tmp;
float *denom = (float *)((char *)tmp + denom_offset);
float *partials = (float *)((char *)tmp + partial_offset);
dim3 score_grid(1, n_head, S);
attention_decode_score_split_scores_kernel<<<score_grid, 256>>>(
scores, q, raw_kv, comp_kv, comp_mask, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim, S);
if (!cuda_ok(cudaGetLastError(), "attention splitkv global score launch")) return -1;
attention_decode_global_softmax_kernel<<<n_head, 256>>>(
scores, denom, sinks, n_score, n_head);
if (!cuda_ok(cudaGetLastError(), "attention splitkv global softmax launch")) return -1;
dim3 value_grid(1, n_head, S);
attention_decode_split_value_kernel<<<value_grid, 256>>>(
partials, scores, raw_kv, comp_kv, raw_count, raw_first_idx,
raw_cap, raw_start, n_score, n_head, head_dim, S);
if (!cuda_ok(cudaGetLastError(), "attention splitkv global value launch")) return -1;
dim3 combine_grid(1, n_head, 1);
attention_decode_split_value_combine_kernel<<<combine_grid, 256>>>(
heads, partials, denom, n_head, head_dim, S);
if (!cuda_ok(cudaGetLastError(), "attention splitkv global combine launch")) return -1;
return 1;
}
/* Partials scratch: n_head * S * (head_dim + 2) floats (n_tokens==1). */
uint64_t stride = (uint64_t)head_dim + 2u;
uint64_t count = (uint64_t)n_head * S * stride;
float *partials = (float *)cuda_tmp_alloc_on(logical_tier, count * sizeof(float),
"attention splitkv partials");
if (!partials) return 0;
dim3 split_grid(1, n_head, S);
attention_decode_splitkv_kernel<<<split_grid, 256>>>(partials,
q,
raw_kv,
comp_kv,
comp_mask,
use_comp_mask,
1, pos0, n_raw, raw_cap, raw_start,
n_comp, window, ratio, n_head, head_dim, S);
if (!cuda_ok(cudaGetLastError(), "attention splitkv partial launch")) return -1;
dim3 combine_grid(1, n_head, 1);
attention_decode_splitkv_combine_kernel<<<combine_grid, 256>>>(heads,
sinks,
partials,
1, n_head, head_dim, S);
if (!cuda_ok(cudaGetLastError(), "attention splitkv combine launch")) return -1;
return 1;
}
extern "C" int ds4_gpu_attention_decode_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16 ||
!heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw ||
raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
(n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) ||
(use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float))) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks");
if (!sinks) return 0;
if (!cuda_attention_score_buffer_fits(n_comp)) {
if (!use_mask && head_dim == 512u &&
!g_cuda_no_window_attention) {
const uint32_t synthetic_pos0 = n_raw - 1u;
dim3 online_grid(1, (n_head + 7u) / 8u, 1);
attention_decode_mixed_heads8_online_kernel<<<online_grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
1,
synthetic_pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
0,
0,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention decode online launch");
}
fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp);
return 0;
}
if (!use_mask && head_dim == 512u &&
g_cuda_decode_heads8_online &&
!g_cuda_no_window_attention) {
const uint32_t synthetic_pos0 = n_raw - 1u;
dim3 online_grid(1, (n_head + 7u) / 8u, 1);
attention_decode_mixed_heads8_online_kernel<<<online_grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
1,
synthetic_pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
0,
0,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention decode heads8 online launch");
}
const uint32_t score_lanes =
g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u);
const uint32_t threads =
head_dim == 512u && score_lanes == 0u &&
!g_cuda_no_decode_value512 ? 512u : 256u;
int score_split_rc = attention_decode_score_split_launch(
logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_mask ? (const float *)comp_mask->ptr : NULL, use_mask,
0, n_raw, raw_cap, raw_start, n_comp, 0, 0,
n_head, head_dim, threads, NULL);
if (score_split_rc == 1) {
return cuda_ok(cudaGetLastError(), "attention exact score split launch");
}
if (score_split_rc < 0) return 0;
/* perf-02 split-KV opt-in (default OFF). n_tokens==1 here by construction.
* S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */
if (cuda_splitkv_decode_requested()) {
int rc = attention_decode_splitkv_launch(
logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_mask ? (const float *)comp_mask->ptr : NULL, use_mask,
0, n_raw, raw_cap, raw_start, n_comp, 0, 0, n_head, head_dim);
if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv launch");
if (rc < 0) return 0;
}
dim3 grid(1, n_head, 1);
attention_decode_mixed_kernel<<<grid, threads>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_mask ? (const float *)comp_mask->ptr : NULL,
use_mask,
1, 0, n_raw, raw_cap, raw_start, n_comp,
0, 0, n_head, head_dim,
score_lanes);
return cuda_ok(cudaGetLastError(), "attention decode launch");
}
extern "C" int ds4_gpu_attention_decode_heads_rope_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
int *fused_inv_rope) {
if (fused_inv_rope) *fused_inv_rope = 0;
if (!g_cuda_exact_score_split_fuse_inv_rope ||
n_rot == 0u || n_rot > head_dim || (n_rot & 1u) ||
head_dim != 512u) {
return ds4_gpu_attention_decode_heads_tensor(
heads, model_map, model_size, sinks_offset, q, raw_kv,
n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp,
comp_mask, use_mask, n_head, head_dim);
}
if (!use_mask && g_cuda_decode_heads8_online && !g_cuda_no_window_attention) {
return ds4_gpu_attention_decode_heads_tensor(
heads, model_map, model_size, sinks_offset, q, raw_kv,
n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp,
comp_mask, use_mask, n_head, head_dim);
}
if (comp_kv_f16 ||
!heads || !q || !raw_kv || !model_map || n_raw == 0 || raw_cap < n_raw ||
raw_start >= raw_cap || (n_comp != 0 && !comp_kv) || (use_mask && !comp_mask) ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
(n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) ||
(use_mask && comp_mask->bytes < (uint64_t)n_comp * sizeof(float)) ||
!cuda_attention_score_buffer_fits(n_comp)) {
return ds4_gpu_attention_decode_heads_tensor(
heads, model_map, model_size, sinks_offset, q, raw_kv,
n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp,
comp_mask, use_mask, n_head, head_dim);
}
const uint32_t score_lanes =
g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u);
const uint32_t threads =
score_lanes == 0u && !g_cuda_no_decode_value512 ? 512u : 256u;
if (threads < 512u) {
return ds4_gpu_attention_decode_heads_tensor(
heads, model_map, model_size, sinks_offset, q, raw_kv,
n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp,
comp_mask, use_mask, n_head, head_dim);
}
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float),
logical_tier, "attn_sinks");
if (!sinks) return 0;
cuda_attention_inv_rope_params rope;
rope.n_rot = n_rot;
rope.pos0 = pos0;
rope.n_ctx_orig = n_ctx_orig;
rope.freq_base = freq_base;
rope.freq_scale = freq_scale;
rope.ext_factor = ext_factor;
rope.attn_factor = attn_factor;
rope.beta_fast = beta_fast;
rope.beta_slow = beta_slow;
int score_split_rc = attention_decode_score_split_launch(
logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_mask ? (const float *)comp_mask->ptr : NULL, use_mask,
0, n_raw, raw_cap, raw_start, n_comp, 0, 0,
n_head, head_dim, threads, &rope);
if (score_split_rc == 1) {
if (fused_inv_rope) *fused_inv_rope = 1;
return cuda_ok(cudaGetLastError(),
"attention exact score split fused inv rope launch");
}
if (score_split_rc < 0) return 0;
return ds4_gpu_attention_decode_heads_tensor(
heads, model_map, model_size, sinks_offset, q, raw_kv,
n_raw, raw_cap, raw_start, comp_kv, comp_kv_f16, n_comp,
comp_mask, use_mask, n_head, head_dim);
}
extern "C" int ds4_gpu_attention_decode_rows_rope_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_attention_decode_row *rows,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!heads || !q || !rows || !model_map || n_rows < 2u ||
n_rows > DS4_GPU_ATTENTION_DECODE_BATCH_MAX || n_head == 0u ||
head_dim != 512u || n_rot == 0u || n_rot > head_dim ||
(n_rot & 1u) != 0u || sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_rows * n_head * head_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(heads);
if (logical_tier < 0 || logical_tier >= g_n_gpus ||
ds4_tensor_device_idx(q) != logical_tier) {
return 0;
}
/* This first grouped path mirrors the promoted default decode exactly.
* Alternative score kernels and graph/split-KV experiments retain the
* one-session dispatcher until they gain equivalent row-table variants. */
if (cuda_env_flag_enabled("DS4_CUDA_NO_EXACT_SCORE_SPLIT_DECODE", 0) ||
!cuda_env_flag_enabled("DS4_CUDA_EXACT_SCORE_SPLIT_DECODE", 1) ||
cuda_splitkv_decode_requested() ||
g_cuda_decode_heads8_online || g_cuda_decode_score4 ||
g_cuda_decode_score8 || g_cuda_no_decode_value512 ||
g_cuda_exact_score_split_graph || g_cuda_exact_score_split_ldg ||
g_cuda_exact_score_split_vec4 ||
g_cuda_exact_score_split_vec4_plain ||
g_cuda_exact_score_split_dim2 ||
g_cuda_exact_score_split_fuse_inv_rope ||
getenv("DS4_CUDA_NO_SCORE_TILE") != NULL ||
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_MIN_SCORE") != NULL ||
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_CHUNK") != NULL ||
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_FLOOR") != NULL ||
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S_MAX") != NULL ||
getenv("DS4_CUDA_EXACT_SCORE_SPLIT_S") != NULL) {
return 0;
}
cuda_attention_decode_row_table table;
memset(&table, 0, sizeof(table));
uint32_t max_dense_score = 0u;
bool have_dense = false;
bool have_indexed = false;
for (uint32_t i = 0; i < n_rows; i++) {
const ds4_gpu_attention_decode_row r = rows[i];
if (r.raw_kv == 0u || r.n_raw == 0u || r.raw_cap < r.n_raw ||
r.raw_start >= r.raw_cap || (r.n_comp != 0u && r.comp_kv == 0u)) {
return 0;
}
if (r.indexed) {
if (r.comp_kv == 0u || r.topk == 0u || r.n_comp == 0u ||
r.top_k == 0u || r.top_k > 512u || r.ratio == 0u) {
return 0;
}
have_indexed = true;
} else {
const uint32_t raw_count = r.n_raw > 256u ? 256u : r.n_raw;
const uint32_t n_score = raw_count + r.n_comp;
/* n_score==1 takes the legacy one-block kernel and is not a
* score-split shape. Decode after any nonempty prompt is >1. */
if (n_score <= 1u || n_score > DS4_CUDA_ATTENTION_SCORE_CAP) {
return 0;
}
if (n_score > max_dense_score) max_dense_score = n_score;
have_dense = true;
}
table.row[i] = r;
}
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float),
logical_tier, "attn_sinks_rows");
if (!sinks) return 0;
if (have_dense) {
if ((uint64_t)n_rows > UINT64_MAX / n_head ||
(uint64_t)n_rows * n_head > UINT64_MAX / max_dense_score) {
return 0;
}
const uint64_t score_count =
(uint64_t)n_rows * n_head * max_dense_score;
float *scores = (float *)cuda_tmp_alloc_on(
logical_tier, score_count * sizeof(float),
"attention exact decode rows");
if (!scores) return 0;
const size_t tile_shmem =
(size_t)(DS4_SCORE_TILE_HEADS + DS4_SCORE_TILE_ROWS) *
DS4_SCORE_TILE_STRIDE * sizeof(float);
static int tile_shmem_ready[DS4_MAX_GPUS] = {0};
int physical_device = 0;
if (cudaGetDevice(&physical_device) != cudaSuccess ||
physical_device < 0 || physical_device >= DS4_MAX_GPUS) {
return 0;
}
if (!tile_shmem_ready[physical_device]) {
if (!cuda_ok(cudaFuncSetAttribute(
attention_decode_score_split_scores_tile512_rows_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)tile_shmem),
"attention score rows shared-memory opt-in")) {
return 0;
}
tile_shmem_ready[physical_device] = 1;
}
dim3 score_grid(
(max_dense_score + DS4_SCORE_TILE_ROWS - 1u) /
DS4_SCORE_TILE_ROWS,
(n_head + DS4_SCORE_TILE_HEADS - 1u) /
DS4_SCORE_TILE_HEADS,
n_rows);
attention_decode_score_split_scores_tile512_rows_kernel
<<<score_grid, 256, tile_shmem>>>(
scores, (const float *)q->ptr, table, n_rows,
max_dense_score, n_head, head_dim);
if (!cuda_ok(cudaGetLastError(),
"attention exact score rows launch")) {
return 0;
}
dim3 final_grid(n_rows, n_head, 1u);
attention_decode_score_split_finalize_rows_kernel
<<<final_grid, 512>>>(
(float *)heads->ptr, sinks, scores, table, n_rows,
max_dense_score, n_head, head_dim);
if (!cuda_ok(cudaGetLastError(),
"attention exact finalize rows launch")) {
return 0;
}
}
if (have_indexed) {
dim3 indexed_grid(n_rows, n_head, 1u);
attention_indexed_mixed_decode_rows_kernel<<<indexed_grid, 256>>>(
(float *)heads->ptr, sinks, (const float *)q->ptr, table,
n_rows, n_head, head_dim);
if (!cuda_ok(cudaGetLastError(),
"attention indexed decode rows launch")) {
return 0;
}
}
const uint32_t pairs = n_rows * n_head * (n_rot / 2u);
rope_tail_decode_rows_kernel<<<(pairs + 255u) / 256u, 256>>>(
(float *)heads->ptr, table, n_rows, n_head, head_dim, n_rot,
n_ctx_orig, 1, freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
return cuda_ok(cudaGetLastError(),
"attention decode rows inverse rope launch");
}
extern "C" int ds4_gpu_attention_prefill_raw_heads_tensor(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size, uint64_t sinks_offset, const ds4_gpu_tensor *q, const ds4_gpu_tensor *raw_kv, uint32_t n_tokens, uint32_t window, uint32_t n_head, uint32_t head_dim) {
if (!heads || !q || !raw_kv || !model_map || sinks_offset > model_size ||
model_size - sinks_offset < (uint64_t)n_head * sizeof(float) ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) ||
window > 256) return 0;
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks");
if (!sinks) return 0;
if (n_tokens > 1 && head_dim == 512 &&
getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL &&
(getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) {
dim3 grid(n_tokens, (n_head + 7u) / 8u, 1);
attention_static_mixed_heads8_online_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
(const float *)raw_kv->ptr,
n_tokens,
0,
window,
1,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention raw window launch");
}
if (g_cublas_ready && n_tokens > 1 && head_dim == 512 &&
getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) {
const uint32_t n_keys = n_tokens;
const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys;
const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim;
const uint64_t score_bytes = score_count * sizeof(float);
const uint64_t out_offset = (score_bytes + 255u) & ~255ull;
const uint64_t tmp_bytes = out_offset + out_count * sizeof(float);
float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention raw cublas");
if (!tmp) return 0;
float *scores = tmp;
float *out_tmp = (float *)((char *)tmp + out_offset);
const float alpha = rsqrtf((float)head_dim);
const float beta = 0.0f;
cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)n_keys,
(int)n_tokens,
(int)head_dim,
&alpha,
(const float *)raw_kv->ptr,
(int)head_dim,
0,
(const float *)q->ptr,
(int)(n_head * head_dim),
(long long)head_dim,
&beta,
scores,
(int)n_keys,
(long long)n_keys * n_tokens,
(int)n_head);
if (!cublas_ok(st, "attention raw score gemm")) return 0;
dim3 sgrid(n_tokens, n_head, 1);
attention_prefill_raw_softmax_kernel<<<sgrid, 256>>>(scores, sinks, n_tokens, window, n_keys);
if (!cuda_ok(cudaGetLastError(), "attention raw softmax launch")) return 0;
const float one = 1.0f;
st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_N,
CUBLAS_OP_N,
(int)head_dim,
(int)n_tokens,
(int)n_keys,
&one,
(const float *)raw_kv->ptr,
(int)head_dim,
0,
scores,
(int)n_keys,
(long long)n_keys * n_tokens,
&beta,
out_tmp,
(int)head_dim,
(long long)head_dim * n_tokens,
(int)n_head);
if (!cublas_ok(st, "attention raw value gemm")) return 0;
uint64_t n = (uint64_t)n_tokens * n_head * head_dim;
attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr,
out_tmp,
n_tokens,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention raw unpack launch");
}
dim3 grid(n_tokens, n_head, 1);
attention_prefill_raw_kernel<<<grid, 128>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_tokens, window, n_head, head_dim);
return cuda_ok(cudaGetLastError(), "attention_prefill_raw launch");
}
static int attention_decode_batch_launch(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16 ||
!heads || !q || !raw_kv || !model_map || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
(n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
(n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) ||
(use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) {
return 0;
}
if (n_comp != 0 && ratio == 0) return 0;
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks");
if (!sinks) return 0;
if (g_n_gpus == 1 && !use_comp_mask &&
n_tokens >= 128u && head_dim == kTTHeadDim && n_head == 64u &&
window == kTTRawWindow && ratio != 0u && n_comp <= 32768u &&
n_raw >= n_tokens &&
(uint64_t)n_raw <= (uint64_t)pos0 + n_tokens &&
!g_cuda_no_window_attention && ds4_cuda_attn_tokentile_arch_ok()) {
cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone;
cudaStream_t stream = cuda_decode_stream();
if (cudaStreamIsCapturing(stream, &capture) == cudaSuccess &&
capture == cudaStreamCaptureStatusNone) {
const uint32_t n_tiles =
(n_tokens + kTTTileTokens - 1u) / kTTTileTokens;
const uint32_t rec_stride = n_comp;
const uint32_t n_mirror_rows = n_tokens + kTTRawWindow - 1u;
const uint32_t raw_before = n_raw - n_tokens;
const uint32_t available_before = raw_before < pos0 ? raw_before : pos0;
const uint32_t available_clamped =
available_before < (kTTRawWindow - 1u)
? available_before : (kTTRawWindow - 1u);
const uint32_t raw_row_min =
(kTTRawWindow - 1u) - available_clamped;
const uint32_t first_raw_pos = (uint32_t)(
(uint64_t)pos0 + n_tokens - n_raw);
uint64_t off = 0;
const uint64_t records_off = off;
off = tt_align256_u64(
off + (uint64_t)n_tiles * (rec_stride ? rec_stride : 1u) *
sizeof(int2));
const uint64_t counts_off = off;
off = tt_align256_u64(
off + (uint64_t)n_tiles * sizeof(uint32_t));
const uint64_t raw_mirror_off = off;
off = tt_align256_u64(
off + (uint64_t)n_mirror_rows * kTTHeadDim * sizeof(half));
const uint64_t comp_mirror_off = off;
off = tt_align256_u64(
off + (uint64_t)n_comp * kTTHeadDim * sizeof(half));
unsigned char *scratch = (unsigned char *)tt_scratch_ensure(
off, "allocate dense token-tile attention scratch");
if (scratch) {
int2 *records = (int2 *)(scratch + records_off);
uint32_t *counts = (uint32_t *)(scratch + counts_off);
half *raw_mirror = (half *)(scratch + raw_mirror_off);
half *comp_mirror = (half *)(scratch + comp_mirror_off);
if (!cuda_ok(cudaFuncSetAttribute(
attention_tokentile_hmma_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)tt_TokentileSmemBudget<kTTStageRows, kTTG>::total),
"set dense token-tile HMMA shared-memory limit")) {
return 0;
}
attention_tokentile_dense_build_kernel
<<<n_tiles, kTTThreads, 0, stream>>>(
records, counts, pos0, n_tokens, ratio, n_comp,
rec_stride);
if (!cuda_ok(cudaGetLastError(),
"launch dense token-tile record builder")) return 0;
attention_tokentile_raw_mirror_kernel
<<<n_mirror_rows, 256, 0, stream>>>(
raw_mirror, (const float *)raw_kv->ptr, NULL, pos0,
n_tokens, raw_cap, raw_start, first_raw_pos,
raw_row_min, head_dim);
if (!cuda_ok(cudaGetLastError(),
"launch dense token-tile raw mirror")) return 0;
if (n_comp != 0u) {
attention_tokentile_comp_mirror_kernel
<<<n_comp, 256, 0, stream>>>(
comp_mirror, (const float *)comp_kv->ptr,
n_comp, head_dim);
if (!cuda_ok(cudaGetLastError(),
"launch dense token-tile compressed mirror")) {
return 0;
}
}
const dim3 grid(n_tiles, n_head / kTTG, 1);
attention_tokentile_hmma_kernel
<<<grid, kTTThreads,
tt_TokentileSmemBudget<kTTStageRows, kTTG>::total,
stream>>>(
(float *)heads->ptr, sinks, (const float *)q->ptr,
raw_mirror, comp_mirror, records, counts, rec_stride,
n_tokens, n_head, raw_row_min);
return cuda_ok(cudaGetLastError(),
"launch dense token-tile HMMA attention");
}
} else {
(void)cudaGetLastError();
}
}
if (!cuda_attention_score_buffer_fits(n_comp)) {
if (!use_comp_mask && head_dim == 512u &&
!g_cuda_no_window_attention) {
dim3 online_grid(n_tokens, (n_head + 7u) / 8u, 1);
attention_decode_mixed_heads8_online_kernel<<<online_grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention decode online launch");
}
fprintf(stderr, "ds4: CUDA attention score buffer too small for %u compressed rows\n", n_comp);
return 0;
}
if (!use_comp_mask && n_tokens > 1 && head_dim == 512 &&
!g_cuda_no_window_attention &&
(getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) {
dim3 grid(n_tokens, (n_head + 7u) / 8u, 1);
attention_decode_mixed_heads8_online_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention decode window launch");
}
if (!use_comp_mask && n_tokens == 1u && head_dim == 512 &&
g_cuda_decode_heads8_online &&
!g_cuda_no_window_attention) {
dim3 grid(1, (n_head + 7u) / 8u, 1);
attention_decode_mixed_heads8_online_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention decode heads8 online batch launch");
}
const uint32_t score_lanes =
g_cuda_decode_score4 ? 4u : (g_cuda_decode_score8 ? 8u : 0u);
const uint32_t threads =
n_tokens == 1u && head_dim == 512u && score_lanes == 0u &&
!g_cuda_no_decode_value512 ? 512u : 256u;
if (n_tokens == 1u) {
int score_split_rc = attention_decode_score_split_launch(
logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio,
n_head, head_dim, threads, NULL);
if (score_split_rc == 1) {
return cuda_ok(cudaGetLastError(), "attention exact score split batch launch");
}
if (score_split_rc < 0) return 0;
}
/* perf-02 split-KV opt-in (default OFF). Single-token decode only; multi-
* token batch shapes already fill the grid and fall through unchanged.
* S==1 / disabled / unhandled -> rc 0, fall through to the old kernel. */
if (n_tokens == 1u && cuda_splitkv_decode_requested()) {
int rc = attention_decode_splitkv_launch(
logical_tier, (float *)heads->ptr, sinks, (const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_comp_mask ? (const float *)comp_mask->ptr : NULL, use_comp_mask,
pos0, n_raw, raw_cap, raw_start, n_comp, window, ratio, n_head, head_dim);
if (rc == 1) return cuda_ok(cudaGetLastError(), "attention decode splitkv batch launch");
if (rc < 0) return 0;
}
dim3 grid(n_tokens, n_head, 1);
attention_decode_mixed_kernel<<<grid, threads>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_comp_mask ? (const float *)comp_mask->ptr : NULL,
use_comp_mask, n_tokens, pos0, n_raw, raw_cap,
raw_start, n_comp, window, ratio, n_head, head_dim,
score_lanes);
return cuda_ok(cudaGetLastError(), "attention decode batch launch");
}
extern "C" int ds4_gpu_attention_decode_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset,
q, raw_kv, NULL, 0, NULL, 0, n_tokens, pos0,
n_raw, raw_cap, raw_start, 0, window, 1,
n_head, head_dim);
}
extern "C" int ds4_gpu_attention_decode_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16) return 0;
return attention_decode_batch_launch(heads, model_map, model_size, sinks_offset,
q, raw_kv, comp_kv, comp_kv_f16, comp_mask, use_comp_mask,
n_tokens, pos0, n_raw, raw_cap, raw_start,
n_comp, window, ratio, n_head, head_dim);
}
extern "C" int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16 ||
!heads || !q || !raw_kv || !comp_kv || !topk || !model_map ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
n_comp == 0 || top_k == 0 ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float) ||
topk->bytes < (uint64_t)n_tokens * top_k * sizeof(int32_t)) {
return 0;
}
if (top_k > 512u) return 0;
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks");
if (!sinks) return 0;
const int32_t *topk_ptr = (const int32_t *)topk->ptr;
if (g_n_gpus == 1 && n_tokens >= 128u && head_dim == kTTHeadDim &&
n_head == 64u && top_k == 512u && window == kTTRawWindow &&
ratio != 0u && n_comp <= 32768u &&
n_raw >= n_tokens &&
(uint64_t)n_raw <= (uint64_t)pos0 + n_tokens &&
ds4_cuda_attn_tokentile_arch_ok()) {
cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone;
cudaStream_t stream = cuda_decode_stream();
if (cudaStreamIsCapturing(stream, &capture) == cudaSuccess &&
capture == cudaStreamCaptureStatusNone) {
const uint32_t n_tiles =
(n_tokens + kTTTileTokens - 1u) / kTTTileTokens;
const uint32_t rec_stride =
(kTTTileTokens * top_k) < n_comp
? (kTTTileTokens * top_k) : n_comp;
const uint32_t n_mirror_rows = n_tokens + kTTRawWindow - 1u;
const uint32_t raw_before = n_raw - n_tokens;
const uint32_t available_before = raw_before < pos0 ? raw_before : pos0;
const uint32_t available_clamped =
available_before < (kTTRawWindow - 1u)
? available_before : (kTTRawWindow - 1u);
const uint32_t raw_row_min =
(kTTRawWindow - 1u) - available_clamped;
const uint32_t first_raw_pos = (uint32_t)(
(uint64_t)pos0 + n_tokens - n_raw);
uint64_t off = 0;
const uint64_t records_off = off;
off = tt_align256_u64(
off + (uint64_t)n_tiles * rec_stride * sizeof(int2));
const uint64_t counts_off = off;
off = tt_align256_u64(
off + (uint64_t)n_tiles * sizeof(uint32_t));
const uint64_t raw_mirror_off = off;
off = tt_align256_u64(
off + (uint64_t)n_mirror_rows * kTTHeadDim * sizeof(half));
const uint64_t comp_mirror_off = off;
off = tt_align256_u64(
off + (uint64_t)n_comp * kTTHeadDim * sizeof(half));
unsigned char *scratch = (unsigned char *)tt_scratch_ensure(
off, "allocate token-tile attention scratch");
if (scratch) {
int2 *records = (int2 *)(scratch + records_off);
uint32_t *counts = (uint32_t *)(scratch + counts_off);
half *raw_mirror = (half *)(scratch + raw_mirror_off);
half *comp_mirror = (half *)(scratch + comp_mirror_off);
const uint32_t bitmap_words = (n_comp + 1u) >> 1u;
const size_t bitmap_smem =
(size_t)bitmap_words * sizeof(uint32_t);
/* The launch limit includes the kernel's static scan array,
* so opt in to the largest bitmap admitted by this path even
* when the dynamic bitmap alone is just under 48 KiB. */
static int union_smem_attr_set = 0;
if (!union_smem_attr_set) {
const int max_bitmap_smem =
(int)(((32768u + 1u) >> 1u) * sizeof(uint32_t));
if (!cuda_ok(cudaFuncSetAttribute(
attention_tokentile_union_build_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
max_bitmap_smem),
"set token-tile union shared-memory limit")) {
return 0;
}
union_smem_attr_set = 1;
}
if (!cuda_ok(cudaFuncSetAttribute(
attention_tokentile_hmma_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)tt_TokentileSmemBudget<kTTStageRows, kTTG>::total),
"set token-tile HMMA shared-memory limit")) {
return 0;
}
attention_tokentile_union_build_kernel
<<<n_tiles, kTTThreads, bitmap_smem, stream>>>(
records, counts, topk_ptr, NULL, pos0, n_tokens,
top_k, ratio, n_comp, rec_stride);
if (!cuda_ok(cudaGetLastError(),
"launch token-tile union builder")) return 0;
attention_tokentile_raw_mirror_kernel
<<<n_mirror_rows, 256, 0, stream>>>(
raw_mirror, (const float *)raw_kv->ptr, NULL, pos0,
n_tokens, raw_cap, raw_start, first_raw_pos,
raw_row_min, head_dim);
if (!cuda_ok(cudaGetLastError(),
"launch token-tile raw mirror")) return 0;
attention_tokentile_comp_mirror_kernel
<<<n_comp, 256, 0, stream>>>(
comp_mirror, (const float *)comp_kv->ptr,
n_comp, head_dim);
if (!cuda_ok(cudaGetLastError(),
"launch token-tile compressed mirror")) return 0;
const dim3 grid(n_tiles, n_head / kTTG, 1);
attention_tokentile_hmma_kernel
<<<grid, kTTThreads,
tt_TokentileSmemBudget<kTTStageRows, kTTG>::total,
stream>>>(
(float *)heads->ptr, sinks, (const float *)q->ptr,
raw_mirror, comp_mirror, records, counts, rec_stride,
n_tokens, n_head, raw_row_min);
return cuda_ok(cudaGetLastError(),
"launch token-tile HMMA attention");
}
} else {
(void)cudaGetLastError();
}
}
if (n_tokens > 1u && top_k == 512u &&
getenv("DS4_CUDA_NO_INDEXED_TOPK_SORT") == NULL) {
const uint64_t sort_bytes = (uint64_t)n_tokens * top_k * sizeof(int32_t);
int32_t *sorted = (int32_t *)cuda_tmp_alloc_on(logical_tier, sort_bytes, "indexed attention topk sort");
if (!sorted) return 0;
indexed_topk_sort_512_asc_kernel<<<n_tokens, 512>>>(sorted, topk_ptr, n_tokens);
if (!cuda_ok(cudaGetLastError(), "indexed attention topk sort launch")) return 0;
topk_ptr = sorted;
}
if (n_tokens > 1 && head_dim == 512 && top_k <= 512u &&
getenv("DS4_CUDA_NO_INDEXED_HEADS8") == NULL) {
if (getenv("DS4_CUDA_INDEXED_TWOPASS") == NULL) {
dim3 grid(n_tokens, (n_head + 15u) / 16u, 1);
attention_indexed_mixed_heads8_online_kernel<8, 16><<<grid, 512>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
(const float *)comp_kv->ptr,
topk_ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
top_k,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention indexed online launch");
}
dim3 grid(n_tokens, (n_head + 7u) / 8u, 1);
attention_indexed_mixed_heads8_rb4_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
(const float *)comp_kv->ptr,
topk_ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
top_k,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention indexed heads8 launch");
}
dim3 grid(n_tokens, n_head, 1);
attention_indexed_mixed_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
(const float *)comp_kv->ptr,
topk_ptr,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
top_k,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention indexed mixed launch");
}
static int attention_prefill_mixed_launch(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 || ratio == 0 ||
(n_comp != 0 && !comp_kv) || (use_comp_mask && !comp_mask) ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) ||
(n_comp && comp_kv->bytes < (uint64_t)n_comp * head_dim * sizeof(float)) ||
(use_comp_mask && comp_mask->bytes < (uint64_t)n_tokens * n_comp * sizeof(float))) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float), logical_tier, "attn_sinks");
if (!sinks) return 0;
if (!use_comp_mask && g_n_gpus == 1 && n_tokens >= 128u &&
head_dim == kTTHeadDim && n_head == 64u &&
window == kTTRawWindow && ratio != 0u && n_comp <= 32768u &&
!g_cuda_no_window_attention && ds4_cuda_attn_tokentile_arch_ok()) {
cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone;
if (cudaStreamIsCapturing(cuda_decode_stream(), &capture) == cudaSuccess &&
capture == cudaStreamCaptureStatusNone) {
return attention_decode_batch_launch(
heads, model_map, model_size, sinks_offset, q, raw_kv,
comp_kv, 0, NULL, 0, n_tokens, 0, n_tokens, n_tokens, 0,
n_comp, window, ratio, n_head, head_dim);
}
(void)cudaGetLastError();
}
if (!use_comp_mask && n_tokens > 1 && head_dim == 512 &&
getenv("DS4_CUDA_NO_WINDOW_ATTENTION") == NULL &&
(getenv("DS4_CUDA_WINDOW_ATTENTION") != NULL || (!g_quality_mode && n_tokens >= 128u))) {
dim3 grid(n_tokens, (n_head + 7u) / 8u, 1);
attention_static_mixed_heads8_online_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention mixed window launch");
}
if (g_cublas_ready && n_tokens > 1 && head_dim == 512 &&
getenv("DS4_CUDA_NO_CUBLAS_ATTENTION") == NULL) {
const uint32_t n_keys = n_tokens + n_comp;
const uint64_t kv_count = (uint64_t)n_keys * head_dim;
const uint64_t score_count = (uint64_t)n_head * n_tokens * n_keys;
const uint64_t out_count = (uint64_t)n_head * n_tokens * head_dim;
const uint64_t kv_bytes = kv_count * sizeof(float);
const uint64_t score_offset = (kv_bytes + 255u) & ~255ull;
const uint64_t score_bytes = score_count * sizeof(float);
const uint64_t out_offset = score_offset + ((score_bytes + 255u) & ~255ull);
const uint64_t tmp_bytes = out_offset + out_count * sizeof(float);
float *tmp = (float *)cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention mixed cublas");
if (!tmp) return 0;
float *kv = tmp;
float *scores = (float *)((char *)tmp + score_offset);
float *out_tmp = (float *)((char *)tmp + out_offset);
attention_prefill_pack_mixed_kv_kernel<<<(kv_count + 255) / 256, 256>>>(
kv,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_tokens,
n_comp,
head_dim);
if (!cuda_ok(cudaGetLastError(), "attention mixed kv pack launch")) return 0;
const float alpha = rsqrtf((float)head_dim);
const float beta = 0.0f;
cublasStatus_t st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)n_keys,
(int)n_tokens,
(int)head_dim,
&alpha,
kv,
(int)head_dim,
0,
(const float *)q->ptr,
(int)(n_head * head_dim),
(long long)head_dim,
&beta,
scores,
(int)n_keys,
(long long)n_keys * n_tokens,
(int)n_head);
if (!cublas_ok(st, "attention mixed score gemm")) return 0;
dim3 sgrid(n_tokens, n_head, 1);
attention_prefill_mixed_softmax_kernel<<<sgrid, 256>>>(
scores,
sinks,
use_comp_mask ? (const float *)comp_mask->ptr : NULL,
use_comp_mask,
n_tokens,
n_comp,
window,
ratio,
n_keys);
if (!cuda_ok(cudaGetLastError(), "attention mixed softmax launch")) return 0;
const float one = 1.0f;
st = cublasSgemmStridedBatched(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_N,
CUBLAS_OP_N,
(int)head_dim,
(int)n_tokens,
(int)n_keys,
&one,
kv,
(int)head_dim,
0,
scores,
(int)n_keys,
(long long)n_keys * n_tokens,
&beta,
out_tmp,
(int)head_dim,
(long long)head_dim * n_tokens,
(int)n_head);
if (!cublas_ok(st, "attention mixed value gemm")) return 0;
uint64_t n = (uint64_t)n_tokens * n_head * head_dim;
attention_prefill_unpack_heads_kernel<<<(n + 255) / 256, 256>>>((float *)heads->ptr,
out_tmp,
n_tokens,
n_head,
head_dim);
return cuda_ok(cudaGetLastError(), "attention mixed unpack launch");
}
dim3 grid(n_tokens, n_head, 1);
attention_prefill_mixed_kernel<<<grid, 256>>>((float *)heads->ptr,
sinks,
(const float *)q->ptr,
(const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
use_comp_mask ? (const float *)comp_mask->ptr : NULL,
use_comp_mask, n_tokens, n_comp, window, ratio,
n_head, head_dim);
return cuda_ok(cudaGetLastError(), "attention prefill mixed launch");
}
extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16) return 0;
return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset,
q, raw_kv, comp_kv, NULL, 0, n_tokens,
n_comp, window, ratio, n_head, head_dim);
}
extern "C" int ds4_gpu_attention_prefill_masked_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (comp_kv_f16) return 0;
return attention_prefill_mixed_launch(heads, model_map, model_size, sinks_offset,
q, raw_kv, comp_kv, comp_mask, 1, n_tokens,
n_comp, window, ratio, n_head, head_dim);
}
extern "C" int ds4_gpu_attention_visual_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
const int32_t *tokens,
uint32_t vocab_size,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!heads || !model_map || !q || !raw_kv || !tokens || vocab_size == 0 ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
(n_comp != 0 && (!comp_kv || ratio == 0)) ||
(use_comp_mask && !comp_mask) || comp_kv_f16 || !g_cublas_ready ||
sinks_offset > model_size ||
(uint64_t)n_head * sizeof(float) > model_size - sinks_offset ||
heads->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
raw_kv->bytes < (uint64_t)raw_cap * head_dim * sizeof(float) ||
(n_comp && comp_kv->bytes <
(uint64_t)n_comp * head_dim * sizeof(float)) ||
(use_comp_mask && comp_mask->bytes <
(uint64_t)n_tokens * n_comp * sizeof(float)) ||
pos0 > UINT32_MAX - n_tokens || n_raw > pos0 + n_tokens) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(heads);
const float *sinks = (const float *)cuda_resolve_weight_ptr(
model_map, sinks_offset, (uint64_t)n_head * sizeof(float),
logical_tier, "visual_attn_sinks");
if (!sinks) return 0;
const uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
std::vector<uint32_t> bounds((size_t)n_tokens * 2u);
if (!ds4_deepseek4_attention_bounds(
(const int *)tokens, n_tokens, vocab_size,
pos0, n_raw, window, bounds.data())) return 0;
if (n_raw > UINT32_MAX - n_comp) return 0;
const uint32_t n_keys = n_raw + n_comp;
const uint64_t kv_count = (uint64_t)n_keys * head_dim;
const uint64_t score_count =
(uint64_t)n_head * n_tokens * n_keys;
const uint64_t out_count =
(uint64_t)n_head * n_tokens * head_dim;
if (kv_count > UINT64_MAX / sizeof(float) ||
score_count > UINT64_MAX / sizeof(float) ||
out_count > UINT64_MAX / sizeof(float)) return 0;
const uint64_t kv_bytes = kv_count * sizeof(float);
const uint64_t score_offset = (kv_bytes + 255u) & ~UINT64_C(255);
const uint64_t score_bytes = score_count * sizeof(float);
if (score_offset > UINT64_MAX - score_bytes - 255u) return 0;
const uint64_t out_offset =
(score_offset + score_bytes + 255u) & ~UINT64_C(255);
const uint64_t out_bytes = out_count * sizeof(float);
if (out_offset > UINT64_MAX - out_bytes - 255u) return 0;
const uint64_t bounds_offset =
(out_offset + out_bytes + 255u) & ~UINT64_C(255);
const uint64_t bounds_bytes =
(uint64_t)n_tokens * 2u * sizeof(uint32_t);
if (bounds_offset > UINT64_MAX - bounds_bytes) return 0;
float *tmp = (float *)cuda_tmp_alloc_on(
logical_tier, bounds_offset + bounds_bytes,
"visual mixed attention");
if (!tmp) return 0;
float *kv = tmp;
float *scores = (float *)((char *)tmp + score_offset);
float *out_tmp = (float *)((char *)tmp + out_offset);
uint32_t *device_bounds =
(uint32_t *)((char *)tmp + bounds_offset);
if (!cuda_ok(cudaMemcpy(device_bounds, bounds.data(),
(size_t)bounds_bytes, cudaMemcpyHostToDevice),
"visual attention bounds upload")) return 0;
attention_visual_pack_mixed_kv_kernel<<<
(kv_count + 255u) / 256u, 256u>>>(
kv, (const float *)raw_kv->ptr,
n_comp ? (const float *)comp_kv->ptr : (const float *)raw_kv->ptr,
n_raw, raw_cap, raw_start, n_comp, head_dim);
if (!cuda_ok(cudaGetLastError(), "visual attention KV pack launch"))
return 0;
const float alpha = rsqrtf((float)head_dim);
const float beta = 0.0f;
cublasStatus_t status = cublasSgemmStridedBatched(
cuda_cublas_for_tier(logical_tier), CUBLAS_OP_T, CUBLAS_OP_N,
(int)n_keys, (int)n_tokens, (int)head_dim,
&alpha, kv, (int)head_dim, 0,
(const float *)q->ptr, (int)(n_head * head_dim),
(long long)head_dim, &beta, scores, (int)n_keys,
(long long)n_keys * n_tokens, (int)n_head);
if (!cublas_ok(status, "visual attention score gemm")) return 0;
dim3 score_grid(n_tokens, n_head, 1u);
attention_visual_mixed_softmax_kernel<<<score_grid, 256u>>>(
scores, sinks,
use_comp_mask ? (const float *)comp_mask->ptr : NULL,
device_bounds, use_comp_mask, n_tokens, pos0, first_raw_pos,
n_raw, n_comp, ratio, n_keys);
if (!cuda_ok(cudaGetLastError(), "visual attention softmax launch"))
return 0;
const float one = 1.0f;
status = cublasSgemmStridedBatched(
cuda_cublas_for_tier(logical_tier), CUBLAS_OP_N, CUBLAS_OP_N,
(int)head_dim, (int)n_tokens, (int)n_keys,
&one, kv, (int)head_dim, 0, scores, (int)n_keys,
(long long)n_keys * n_tokens, &beta, out_tmp, (int)head_dim,
(long long)head_dim * n_tokens, (int)n_head);
if (!cublas_ok(status, "visual attention value gemm")) return 0;
const uint64_t values = (uint64_t)n_tokens * n_head * head_dim;
attention_prefill_unpack_heads_kernel<<<
(values + 255u) / 256u, 256u>>>(
(float *)heads->ptr, out_tmp, n_tokens, n_head, head_dim);
return cuda_ok(cudaGetLastError(), "visual attention unpack launch");
}
extern "C" int ds4_gpu_attention_output_q8_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp,
ds4_gpu_tensor *low_tmp,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens) {
(void)group_tmp;
(void)low_tmp;
if (!out || !low || !heads || !model_map ||
group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0) {
return 0;
}
const uint64_t low_dim = (uint64_t)n_groups * rank;
const uint64_t blocks_a = (group_dim + 31) / 32;
const uint64_t blocks_b = (low_dim + 31) / 32;
const uint64_t out_a_bytes = (uint64_t)n_groups * rank * blocks_a * 34;
const uint64_t out_b_bytes = out_dim * blocks_b * 34;
if (out_a_offset > model_size || out_b_offset > model_size ||
out_a_bytes > model_size - out_a_offset ||
out_b_bytes > model_size - out_b_offset ||
heads->bytes < (uint64_t)n_tokens * n_groups * group_dim * sizeof(float) ||
low->bytes < (uint64_t)n_tokens * low_dim * sizeof(float) ||
out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const int physical_device =
(g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus)
? g_gpu[logical_tier].device_id : 0;
const unsigned char *out_a = reinterpret_cast<const unsigned char *>(
cuda_resolve_weight_ptr(model_map, out_a_offset, out_a_bytes, logical_tier, "attn_out_a"));
const unsigned char *out_b = reinterpret_cast<const unsigned char *>(
cuda_resolve_weight_ptr(model_map, out_b_offset, out_b_bytes, logical_tier, "attn_out_b"));
if (!out_a || !out_b) return 0;
const uint32_t profile = getenv("DS4_CUDA_ATTN_OUTPUT_PROFILE") != NULL;
cudaEvent_t prof_ev[3] = {NULL, NULL, NULL};
if (profile) {
for (uint32_t i = 0; i < 3u; i++) {
if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) {
for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]);
memset(prof_ev, 0, sizeof(prof_ev));
break;
}
}
if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0);
}
const __half *out_a_f16 = NULL;
uint32_t out_a_cublas_min_tokens = 2u;
const char *out_a_min_env = getenv("DS4_CUDA_ATTENTION_OUTPUT_A_CUBLAS_MIN");
if (out_a_min_env && out_a_min_env[0]) {
char *endp = NULL;
long v = strtol(out_a_min_env, &endp, 10);
if (endp != out_a_min_env && v > 1 && v < 4096) out_a_cublas_min_tokens = (uint32_t)v;
}
if (!g_quality_mode &&
g_cublas_ready &&
n_tokens >= out_a_cublas_min_tokens &&
getenv("DS4_CUDA_NO_CUBLAS_ATTENTION_OUTPUT_A") == NULL) {
out_a_f16 = cuda_q8_f16_ptr(model_map, out_a_offset, out_a_bytes, group_dim, low_dim, physical_device, "attn_output_a");
}
if (out_a_f16) {
const uint64_t heads_h_count = (uint64_t)n_groups * n_tokens * group_dim;
const uint64_t low_tmp_count = (uint64_t)n_groups * n_tokens * rank;
const uint64_t heads_h_bytes = heads_h_count * sizeof(__half);
const uint64_t low_tmp_offset = (heads_h_bytes + 255u) & ~255ull;
const uint64_t tmp_bytes = low_tmp_offset + low_tmp_count * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a cublas");
if (!tmp) return 0;
__half *heads_h = (__half *)tmp;
float *low_packed = (float *)((char *)tmp + low_tmp_offset);
attention_pack_group_heads_f16_kernel<<<(heads_h_count + 255) / 256, 256, 0, cuda_decode_stream()>>>(
heads_h,
(const float *)heads->ptr,
n_tokens,
n_groups,
group_dim);
if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a pack launch")) return 0;
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasGemmStridedBatchedEx(cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T,
CUBLAS_OP_N,
(int)rank,
(int)n_tokens,
(int)group_dim,
&alpha,
out_a_f16,
CUDA_R_16F,
(int)group_dim,
(long long)rank * group_dim,
heads_h,
CUDA_R_16F,
(int)group_dim,
(long long)n_tokens * group_dim,
&beta,
low_packed,
CUDA_R_32F,
(int)rank,
(long long)rank * n_tokens,
(int)n_groups,
CUDA_R_32F,
CUBLAS_GEMM_DEFAULT);
if (!cublas_ok(st, "attention output a gemm")) return 0;
attention_unpack_group_low_kernel<<<(low_tmp_count + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)low->ptr,
low_packed,
n_tokens,
n_groups,
rank);
if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a unpack launch")) return 0;
} else {
const uint64_t x_rows = (uint64_t)n_tokens * n_groups;
const uint64_t xq_bytes = x_rows * blocks_a * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output a q8 prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1);
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(xq,
xscale,
(const float *)heads->ptr,
group_dim,
blocks_a);
if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a prequant launch")) return 0;
int grouped_mma_done = 0;
if (n_tokens >= 8u) {
/* One mma launch per group: T=32 matches the warp tree of the
* grouped reference kernels (multi-term slots for blocks > 32). */
grouped_mma_done = 1;
for (uint32_t g = 0; g < n_groups && grouped_mma_done == 1; g++) {
const int rc = cuda_q8_mma_try_launch(
(float *)low->ptr + (uint64_t)g * rank,
reinterpret_cast<const unsigned char *>(out_a) +
(uint64_t)g * rank * blocks_a * 34u,
xq + (uint64_t)g * blocks_a * 32u,
xscale + (uint64_t)g * blocks_a,
group_dim, rank, n_tokens, blocks_a,
(uint64_t)n_groups * blocks_a, low_dim, 32u);
if (rc < 0) return 0;
if (rc == 0) grouped_mma_done = 0;
}
}
if (grouped_mma_done) {
/* handled */
} else if (getenv("DS4_CUDA_NO_ATTN_A_TOK2") == NULL && n_tokens >= 2u) {
dim3 grid_a(((unsigned)low_dim + 7u) / 8u, ((unsigned)n_tokens + 1u) / 2u, 1);
grouped_q8_0_a_preq_warp8_tok2_kernel<<<grid_a, 512, 0, cuda_decode_stream()>>>((float *)low->ptr,
out_a,
xq,
xscale,
group_dim,
rank,
n_groups,
n_tokens,
blocks_a,
use_dp4a);
} else {
dim3 grid_a(((unsigned)low_dim + 7u) / 8u, (unsigned)n_tokens, 1);
grouped_q8_0_a_preq_warp8_kernel<<<grid_a, 256, 0, cuda_decode_stream()>>>((float *)low->ptr,
out_a,
xq,
xscale,
group_dim,
rank,
n_groups,
n_tokens,
blocks_a,
use_dp4a);
}
if (!cuda_ok(cudaGetLastError(), "attention_output_q8_a preq launch")) return 0;
}
if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0);
(void)out_b;
int ok = cuda_matmul_q8_0_tensor_labeled(out,
model_map,
model_size,
out_b_offset,
low_dim,
out_dim,
low,
n_tokens,
"attn_output_b");
if (prof_ev[2]) {
(void)cudaEventRecord(prof_ev[2], 0);
if (cudaEventSynchronize(prof_ev[2]) == cudaSuccess) {
float ms_a = 0.0f, ms_b = 0.0f, ms_total = 0.0f;
(void)cudaEventElapsedTime(&ms_a, prof_ev[0], prof_ev[1]);
(void)cudaEventElapsedTime(&ms_b, prof_ev[1], prof_ev[2]);
(void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[2]);
fprintf(stderr,
"ds4: CUDA attention output profile tokens=%u groups=%u group_dim=%llu rank=%llu low=%llu out=%llu A=%.3f B=%.3f total=%.3f ms\n",
n_tokens,
n_groups,
(unsigned long long)group_dim,
(unsigned long long)rank,
(unsigned long long)low_dim,
(unsigned long long)out_dim,
ms_a,
ms_b,
ms_total);
}
for (uint32_t i = 0; i < 3u; i++) (void)cudaEventDestroy(prof_ev[i]);
}
return ok;
}
extern "C" int ds4_gpu_attention_output_low_q8_rows_exact_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
const ds4_gpu_tensor *heads,
uint32_t n_rows) {
if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
n_groups_total == 0 || group_cnt == 0 ||
group0 > n_groups_total || group_cnt > n_groups_total - group0 ||
n_rows == 0 || (uint64_t)n_rows * group_cnt > 65535u) {
return 0;
}
const uint64_t low_dim = (uint64_t)group_cnt * rank;
const uint64_t blocks_a = (group_dim + 31) / 32;
const uint64_t row_a_bytes = blocks_a * 34u;
const uint64_t a_offset =
out_a_offset + (uint64_t)group0 * rank * row_a_bytes;
const uint64_t out_a_bytes = low_dim * row_a_bytes;
if (a_offset < out_a_offset || a_offset > model_size ||
out_a_bytes > model_size - a_offset ||
heads->bytes < (uint64_t)n_rows * n_groups_total * group_dim *
sizeof(float) ||
low->bytes < (uint64_t)n_rows * low_dim * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(low);
const unsigned char *out_a = reinterpret_cast<const unsigned char *>(
cuda_resolve_weight_ptr(model_map, a_offset, out_a_bytes,
logical_tier, "attn_out_a_rows"));
if (!out_a) return 0;
const uint64_t x_rows = (uint64_t)n_rows * group_cnt;
const uint64_t xq_bytes = x_rows * blocks_a * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + x_rows * blocks_a * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes, "attention output low q8 prequant");
if (!tmp) return 0;
int8_t *xq = (int8_t *)tmp;
float *xscale = (float *)((char *)tmp + scale_offset);
const int use_dp4a = cuda_q8_use_dp4a();
dim3 qgrid((unsigned)blocks_a, (unsigned)x_rows, 1);
quantize_q8_0_group_slice_rows_kernel<<<qgrid, 32, 0, cuda_decode_stream()>>>(
xq,
xscale,
(const float *)heads->ptr,
group_dim,
blocks_a,
n_groups_total,
group0,
group_cnt);
if (!cuda_ok(cudaGetLastError(),
"attention_output_low_q8 rows prequant launch")) return 0;
if (n_rows >= 8u && cuda_decode_stream() == (cudaStream_t)0) {
bool complete = true;
for (uint32_t group = 0; group < group_cnt; group++) {
const int rc = cuda_q8_mma_try_launch(
(float *)low->ptr + (uint64_t)group * rank,
out_a + (uint64_t)group * rank * row_a_bytes,
xq + (uint64_t)group * blocks_a * 32u,
xscale + (uint64_t)group * blocks_a,
group_dim, rank, n_rows, blocks_a,
(uint64_t)group_cnt * blocks_a, low_dim, 32u);
if (rc < 0) return 0;
if (!rc) { complete = false; break; }
}
if (complete) return 1;
}
dim3 grid_a(((unsigned)low_dim + 7u) / 8u, n_rows, 1u);
grouped_q8_0_a_preq_warp8_kernel<<<grid_a, 256, 0, cuda_decode_stream()>>>((float *)low->ptr,
out_a,
xq,
xscale,
group_dim,
rank,
group_cnt,
n_rows,
blocks_a,
use_dp4a);
return cuda_ok(cudaGetLastError(),
"attention_output_low_q8 rows launch");
}
extern "C" int ds4_gpu_attention_output_low_q8_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
const ds4_gpu_tensor *heads) {
return ds4_gpu_attention_output_low_q8_rows_exact_tensor(
low, model_map, model_size, out_a_offset, group_dim, rank,
n_groups, 0u, n_groups, heads, 1u);
}
extern "C" int ds4_gpu_attention_output_q8_tp_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *heads) {
if (!out || !low || !heads || !model_map ||
group_dim == 0 || rank == 0 || n_groups_total == 0 ||
group_cnt == 0 || group0 > n_groups_total ||
group_cnt > n_groups_total - group0 || out_dim == 0) {
return 0;
}
const uint64_t blocks_a = (group_dim + 31u) / 32u;
const uint64_t row_a_bytes = blocks_a * 34u;
const uint64_t low_dim_total = (uint64_t)n_groups_total * rank;
const uint64_t k_off = (uint64_t)group0 * rank;
const uint64_t k_cnt = (uint64_t)group_cnt * rank;
if ((k_off % 32u) != 0 || (k_cnt % 32u) != 0) return 0;
if (heads->bytes < (uint64_t)(group0 + group_cnt) * group_dim * sizeof(float) ||
low->bytes < k_cnt * sizeof(float) ||
out->bytes < out_dim * sizeof(float)) {
return 0;
}
ds4_gpu_tensor heads_slice = *heads;
heads_slice.ptr = (char *)heads->ptr + (uint64_t)group0 * group_dim * sizeof(float);
heads_slice.bytes = (uint64_t)group_cnt * group_dim * sizeof(float);
heads_slice.owner = 0;
const uint64_t a_off = out_a_offset + (uint64_t)group0 * rank * row_a_bytes;
return ds4_gpu_attention_output_low_q8_tensor(low,
model_map,
model_size,
a_off,
group_dim,
rank,
group_cnt,
&heads_slice) &&
ds4_gpu_matmul_q8_0_kslice_rows_tensor(out,
model_map,
model_size,
out_b_offset,
low_dim_total,
out_dim,
k_off,
k_cnt,
low,
1);
}
extern "C" int ds4_gpu_swiglu_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *gate, const ds4_gpu_tensor *up, uint32_t n, float clamp, float weight) {
if (!out || !gate || !up ||
out->bytes < (uint64_t)n * sizeof(float) ||
gate->bytes < (uint64_t)n * sizeof(float) ||
up->bytes < (uint64_t)n * sizeof(float)) return 0;
swiglu_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)gate->ptr, (const float *)up->ptr, n, clamp, weight);
return cuda_ok(cudaGetLastError(), "swiglu launch");
}
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
if (getenv("DS4_CUDA_DISABLE_SHARED_GATE_UP_PAIR") == NULL) {
return ds4_gpu_matmul_q8_0_pair_tensor(gate, up,
model_map, model_size,
gate_offset, up_offset,
in_dim, out_dim, out_dim,
x, 1) &&
ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f);
}
return ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size,
gate_offset, in_dim, out_dim, x, 1) &&
ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size,
up_offset, in_dim, out_dim, x, 1) &&
ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f);
}
extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor(
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *prequant,
uint32_t expert_split,
bool home_rank) {
if (!mid || !x || !model_map || in_dim == 0u || out_dim == 0u ||
x->bytes < in_dim * sizeof(float) ||
mid->bytes < out_dim * sizeof(float) ||
(selected && (selected->bytes < 6u * sizeof(int32_t) ||
expert_split == 0u))) {
return 0;
}
const uint64_t blocks = (in_dim + 31u) / 32u;
if (gate_offset > model_size || up_offset > model_size ||
out_dim > UINT64_MAX / (blocks * 34u)) {
return 0;
}
const uint64_t weight_bytes = out_dim * blocks * 34u;
if (weight_bytes > model_size - gate_offset ||
weight_bytes > model_size - up_offset) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(x);
if (logical_tier < 0 || logical_tier >= g_n_gpus) return 0;
if (selected && ds4_tensor_device_idx(selected) != logical_tier) return 0;
if (prequant && ds4_tensor_device_idx(prequant) != logical_tier) return 0;
const int mid_tier = ds4_tensor_device_idx(mid);
if (mid_tier != logical_tier && !g_gpu_peer_ok[logical_tier][mid_tier]) {
return 0;
}
const char *gate_w = cuda_resolve_weight_ptr(
model_map, gate_offset, weight_bytes, logical_tier,
"shared_mid_gate_exact");
const char *up_w = cuda_resolve_weight_ptr(
model_map, up_offset, weight_bytes, logical_tier,
"shared_mid_up_exact");
if (!gate_w || !up_w) return 0;
const uint64_t xq_bytes = blocks * 32u;
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
int8_t *xq;
float *xscale;
if (prequant) {
if (prequant->bytes < tmp_bytes) return 0;
xq = (int8_t *)prequant->ptr;
xscale = (float *)((char *)prequant->ptr + scale_offset);
} else {
void *tmp = cuda_tmp_alloc_on(logical_tier, tmp_bytes,
"shared mid q8 exact prequant");
if (!tmp) return 0;
xq = (int8_t *)tmp;
xscale = (float *)((char *)tmp + scale_offset);
quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32, 0, cuda_decode_stream()>>>(
xq, xscale, (const float *)x->ptr, in_dim, blocks);
if (!cuda_ok(cudaGetLastError(),
"shared mid q8 exact quantize launch")) {
return 0;
}
}
shared_mid_q8_0_preq_warp8_exact_kernel<<<
((unsigned)out_dim + 7u) / 8u, 256, 0, cuda_decode_stream()>>>(
(float *)mid->ptr,
(const unsigned char *)gate_w,
(const unsigned char *)up_w,
xq,
xscale,
in_dim,
out_dim,
blocks,
clamp,
selected ? (const int32_t *)selected->ptr : NULL,
expert_split,
home_rank,
cuda_q8_use_dp4a());
return cuda_ok(cudaGetLastError(), "shared mid q8 exact launch");
}
extern "C" int ds4_gpu_add_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, uint32_t n) {
if (!out || !a || !b ||
out->bytes < (uint64_t)n * sizeof(float) ||
a->bytes < (uint64_t)n * sizeof(float) ||
b->bytes < (uint64_t)n * sizeof(float)) return 0;
add_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)a->ptr, (const float *)b->ptr, n);
return cuda_ok(cudaGetLastError(), "add launch");
}
extern "C" int ds4_gpu_add_xdev_tensor(ds4_gpu_tensor *out,
const ds4_gpu_tensor *local,
const ds4_gpu_tensor *remote,
ds4_gpu_tensor *remote_tmp,
uint32_t n) {
if (!out || !local || !remote ||
out->bytes < (uint64_t)n * sizeof(float) ||
local->bytes < (uint64_t)n * sizeof(float) ||
remote->bytes < (uint64_t)n * sizeof(float)) return 0;
if (n == 0) return 1;
const int od = ds4_tensor_device_idx(out);
const int ld = ds4_tensor_device_idx(local);
const int rd = ds4_tensor_device_idx(remote);
if (od != ld) return 0;
const ds4_gpu_tensor *rhs = remote;
if (rd != od) {
if (!remote_tmp ||
remote_tmp->bytes < (uint64_t)n * sizeof(float) ||
ds4_tensor_device_idx(remote_tmp) != od) return 0;
if (!ds4_gpu_tensor_copy_xdev(remote_tmp, remote,
(uint64_t)n * sizeof(float))) return 0;
rhs = remote_tmp;
}
int ok = 0;
WITH_DEVICE(g_gpu[od].device_id) {
cudaStream_t s = (cudaStream_t)g_gpu[od].stream;
add_kernel<<<(n + 255u) / 256u, 256, 0, s>>>(
(float *)out->ptr,
(const float *)local->ptr,
(const float *)rhs->ptr,
n);
ok = cuda_ok(cudaGetLastError(), "xdev add launch");
cudaEvent_t e = (cudaEvent_t)g_gpu[od].boundary_event;
if (ok) ok = cuda_ok(cudaEventRecord(e, s), "xdev add event record");
if (ok) ok = cuda_ok(cudaStreamWaitEvent(0, e, 0), "xdev add default wait");
if (ok && g_xdev_sync_debug) {
ok = cuda_ok(cudaStreamSynchronize(s), "xdev add sync");
}
}
return ok;
}
extern "C" int ds4_gpu_directional_steering_project_tensor(
ds4_gpu_tensor *x,
const ds4_gpu_tensor *directions,
uint32_t layer,
uint32_t width,
uint32_t rows,
float scale) {
if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0;
const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float);
const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float);
if (x->bytes < x_bytes || directions->bytes < dir_bytes) return 0;
uint32_t nth = 256u;
while (nth > width && nth > 1u) nth >>= 1;
directional_steering_project_kernel<<<rows, nth>>>(
(float *)x->ptr,
(const float *)directions->ptr,
layer,
width,
rows,
scale);
return cuda_ok(cudaGetLastError(), "directional steering launch");
}
extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t token, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits) {
if (!selected || !weights || !probs || !logits || !model_map || n_expert_groups > 1u || n_group_used > 0u) return 0;
if ((n_expert != 256u && n_expert != 384u) || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f ||
logits->bytes < n_expert * sizeof(float) || probs->bytes < n_expert * sizeof(float) ||
selected->bytes < 6u * sizeof(int32_t) || weights->bytes < 6u * sizeof(float)) return 0;
int32_t tok = (int32_t)token;
int ok = 1;
const float *bias = NULL;
const int32_t *hash = NULL;
const int logical_tier = ds4_tensor_device_idx(selected);
if (ok && has_bias && !hash_mode) {
if (bias_offset > model_size || model_size - bias_offset < n_expert * sizeof(float)) ok = 0;
else bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, n_expert * sizeof(float), logical_tier, "router_bias");
if (!bias) ok = 0;
}
if (ok && hash_mode) {
const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t);
if (hash_offset > model_size || hash_bytes > model_size - hash_offset) ok = 0;
else hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash");
if (!hash) ok = 0;
}
if (ok) {
if (n_expert == 384u) {
router_select_warp_topk_kernel<384><<<1, dim3(32, 4), 0, cuda_decode_stream()>>>(
(int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
has_bias && !hash_mode, hash_mode);
} else if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL &&
getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) {
dim3 block(32, 4, 1);
router_select_warp_topk_kernel<<<1, block, 0, cuda_decode_stream()>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
has_bias && !hash_mode, hash_mode);
} else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) {
router_select_parallel_kernel<<<1, 256, 0, cuda_decode_stream()>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
has_bias && !hash_mode, hash_mode);
} else {
router_select_kernel<<<1, 1, 0, cuda_decode_stream()>>>((int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
has_bias && !hash_mode, hash_mode);
}
ok = cuda_ok(cudaGetLastError(), "router_select launch");
}
return ok;
}
extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits, const ds4_gpu_tensor *tokens, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_tokens) {
if ((n_expert != 256u && n_expert != 384u) || n_expert_used != 6u || fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0;
if (!selected || !weights || !probs || !logits || !tokens || !model_map || n_tokens == 0 ||
n_expert_groups > 1u || n_group_used > 0u ||
logits->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) ||
probs->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) ||
tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t) ||
selected->bytes < (uint64_t)n_tokens * 6u * sizeof(int32_t) ||
weights->bytes < (uint64_t)n_tokens * 6u * sizeof(float)) {
return 0;
}
const float *bias = NULL;
const int32_t *hash = NULL;
const int logical_tier = ds4_tensor_device_idx(selected);
if (has_bias && !hash_mode) {
if (bias_offset > model_size || model_size - bias_offset < n_expert * sizeof(float)) return 0;
bias = (const float *)cuda_resolve_weight_ptr(model_map, bias_offset, n_expert * sizeof(float), logical_tier, "router_bias");
if (!bias) return 0;
}
if (hash_mode) {
const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t);
if (hash_offset > model_size || hash_bytes > model_size - hash_offset) return 0;
hash = (const int32_t *)cuda_resolve_weight_ptr(model_map, hash_offset, hash_bytes, logical_tier, "router_hash");
if (!hash) return 0;
}
if (n_expert == 384u) {
router_select_warp_topk_kernel<384><<<(n_tokens + 3u) / 4u, dim3(32, 4), 0, cuda_decode_stream()>>>(
(int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
bias, hash, (const float *)logits->ptr, (const int32_t *)tokens->ptr,
0, hash_rows, n_tokens, has_bias && !hash_mode, hash_mode);
} else if (getenv("DS4_CUDA_NO_WARP_ROUTER_SELECT") == NULL &&
getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) {
dim3 block(32, 4, 1);
router_select_warp_topk_kernel<<<(n_tokens + 3u) / 4u, block>>>((int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
hash,
(const float *)logits->ptr,
(const int32_t *)tokens->ptr,
0,
hash_rows,
n_tokens,
has_bias && !hash_mode,
hash_mode);
} else if (getenv("DS4_CUDA_NO_PARALLEL_ROUTER_SELECT") == NULL) {
router_select_parallel_kernel<<<n_tokens, 256>>>((int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
hash,
(const float *)logits->ptr,
(const int32_t *)tokens->ptr,
0,
hash_rows,
n_tokens,
has_bias && !hash_mode,
hash_mode);
} else {
router_select_kernel<<<n_tokens, 1>>>((int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
hash,
(const float *)logits->ptr,
(const int32_t *)tokens->ptr,
0,
hash_rows,
n_tokens,
has_bias && !hash_mode,
hash_mode);
}
return cuda_ok(cudaGetLastError(), "router_select launch");
}
extern "C" int ds4_gpu_router_select_batch_visual_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
bool has_bias,
bool hash_mode,
const void *vision_map,
uint64_t vision_size,
uint64_t visual_bias_offset,
const ds4_gpu_tensor *logits,
const ds4_gpu_tensor *tokens,
uint32_t vocab_size,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
if (!selected || !weights || !probs || !logits || !tokens ||
!model_map || !vision_map || n_tokens == 0 || vocab_size == 0 ||
n_expert != 256u || n_expert_used != 6u ||
fabsf(expert_weight_scale - 1.5f) > 1.0e-6f ||
logits->bytes < (uint64_t)n_tokens * 256u * sizeof(float) ||
probs->bytes < (uint64_t)n_tokens * 256u * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * 6u * sizeof(int32_t) ||
weights->bytes < (uint64_t)n_tokens * 6u * sizeof(float) ||
tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t)) return 0;
const int logical_tier = ds4_tensor_device_idx(selected);
const float *bias = NULL;
const int32_t *hash = NULL;
if (has_bias && !hash_mode) {
if (bias_offset > model_size || model_size - bias_offset < 256u * sizeof(float)) return 0;
bias = (const float *)cuda_resolve_weight_ptr(
model_map, bias_offset, 256u * sizeof(float),
logical_tier, "router_bias");
if (!bias) return 0;
}
if (hash_mode) {
const uint64_t hash_bytes = (uint64_t)hash_rows * 6u * sizeof(int32_t);
if (hash_offset > model_size || model_size - hash_offset < hash_bytes) return 0;
hash = (const int32_t *)cuda_resolve_weight_ptr(
model_map, hash_offset, hash_bytes,
logical_tier, "router_hash");
if (!hash) return 0;
}
if (visual_bias_offset > vision_size ||
vision_size - visual_bias_offset < 256u * sizeof(float)) return 0;
const float *visual_bias = (const float *)cuda_resolve_weight_ptr(
vision_map, visual_bias_offset, 256u * sizeof(float),
logical_tier, "visual_router_bias");
if (!visual_bias) return 0;
router_select_visual_parallel_kernel<<<n_tokens, 256>>>(
(int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
visual_bias,
hash,
(const float *)logits->ptr,
(const int32_t *)tokens->ptr,
hash_rows,
vocab_size,
n_tokens,
has_bias && !hash_mode,
hash_mode);
return cuda_ok(cudaGetLastError(), "visual router select launch");
}
__device__ static float dev_f16_to_f32(uint16_t v) {
return __half2float(*reinterpret_cast<const __half *>(&v));
}
__device__ __forceinline__ static uint32_t dev_unpack_iq2_signs(uint32_t v) {
const uint32_t p = __popc(v) & 1u;
const uint32_t s = v ^ (p << 7u);
return s * 0x01010101u;
}
__device__ __forceinline__ static int32_t dev_iq2_dp4a_8(uint64_t grid, uint32_t sign, const int8_t *q8, int32_t acc) {
const uint32_t signs = dev_unpack_iq2_signs(sign);
const int32_t sm0 = __vcmpne4(signs & 0x08040201u, 0);
const int32_t sm1 = __vcmpne4(signs & 0x80402010u, 0);
const int32_t g0 = __vsub4((int32_t)(uint32_t)grid ^ sm0, sm0);
const int32_t g1 = __vsub4((int32_t)(uint32_t)(grid >> 32) ^ sm1, sm1);
acc = __dp4a(g0, *(const int32_t *)(q8 + 0), acc);
acc = __dp4a(g1, *(const int32_t *)(q8 + 4), acc);
return acc;
}
__device__ static int32_t dev_dot_q2_16(const uint8_t *q2, const int8_t *q8, int shift) {
int32_t sum = 0;
#pragma unroll
for (uint32_t i = 0; i < 16; i += 4) {
const int32_t v = (*(const int32_t *)(q2 + i) >> shift) & 0x03030303;
sum = __dp4a(v, *(const int32_t *)(q8 + i), sum);
}
return sum;
}
__device__ static int32_t dev_dot_iq2_pair_16(uint8_t grid0, uint32_t sign0, uint8_t grid1, uint32_t sign1, const int8_t *q8) {
int32_t sum = 0;
sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid0], cuda_ksigns_iq2xs[sign0], q8, sum);
sum = dev_iq2_dp4a_8(cuda_iq2xxs_grid[grid1], cuda_ksigns_iq2xs[sign1], q8 + 8, sum);
return sum;
}
__device__ __forceinline__ static void dev_iq2_i8x8_lut(
const uint64_t *grid,
const uint8_t *signs,
uint8_t grid_idx,
uint32_t sign_idx,
int32_t *w0,
int32_t *w1) {
const uint32_t s = dev_unpack_iq2_signs(signs[sign_idx]);
const int32_t sm0 = __vcmpne4(s & 0x08040201u, 0);
const int32_t sm1 = __vcmpne4(s & 0x80402010u, 0);
const uint64_t g = grid[grid_idx];
*w0 = __vsub4((int32_t)(uint32_t)g ^ sm0, sm0);
*w1 = __vsub4((int32_t)(uint32_t)(g >> 32) ^ sm1, sm1);
}
__device__ static float dev_dot_iq2_xxs_q8_K_block_lut(
const cuda_block_iq2_xxs *x,
const cuda_block_q8_K *y,
const uint64_t *grid,
const uint8_t *signs) {
const float xd = dev_f16_to_f32(x->d);
const uint16_t *q2 = x->qs;
const int8_t *q8 = y->qs;
int32_t bsum = 0;
for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u);
int32_t w[8];
dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]);
int32_t sumi = 0;
sumi = __dp4a(w[0], *(const int32_t *)(q8 + ib32 * 32u + 0), sumi);
sumi = __dp4a(w[1], *(const int32_t *)(q8 + ib32 * 32u + 4), sumi);
sumi = __dp4a(w[2], *(const int32_t *)(q8 + ib32 * 32u + 8), sumi);
sumi = __dp4a(w[3], *(const int32_t *)(q8 + ib32 * 32u + 12), sumi);
sumi = __dp4a(w[4], *(const int32_t *)(q8 + ib32 * 32u + 16), sumi);
sumi = __dp4a(w[5], *(const int32_t *)(q8 + ib32 * 32u + 20), sumi);
sumi = __dp4a(w[6], *(const int32_t *)(q8 + ib32 * 32u + 24), sumi);
sumi = __dp4a(w[7], *(const int32_t *)(q8 + ib32 * 32u + 28), sumi);
bsum += sumi * ls;
}
return 0.125f * xd * y->d * (float)bsum;
}
__device__ static float dev_dot_iq2_xxs_q8_K_block(const cuda_block_iq2_xxs *x, const cuda_block_q8_K *y) {
const float d = dev_f16_to_f32(x->d) * y->d;
const uint16_t *q2 = x->qs;
const int8_t *q8 = y->qs;
int32_t bsum = 0;
for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const uint32_t ls = 2u * (aux1 >> 28) + 1u;
const uint8_t a0 = (uint8_t)(aux0 & 0xffu);
const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu);
const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu);
const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu);
int32_t sumi = 0;
sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8);
q8 += 16;
sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8);
q8 += 16;
bsum += sumi * (int32_t)ls;
}
return 0.125f * d * (float)bsum;
}
__device__ static void dev_dot_iq2_xxs_q8_K_block8_deq_lut(
const cuda_block_iq2_xxs *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
const cuda_block_q8_K *y4,
const cuda_block_q8_K *y5,
const cuda_block_q8_K *y6,
const cuda_block_q8_K *y7,
uint32_t n,
float acc[8],
const uint64_t *grid,
const uint8_t *signs) {
const float xd = dev_f16_to_f32(x->d);
const uint16_t *q2 = x->qs;
int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const int8_t *q8[8] = {
y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL,
y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL,
};
for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u);
int32_t w[8];
dev_iq2_i8x8_lut(grid, signs, (uint8_t)(aux0 & 0xffu), (aux1 >> 0) & 127u, &w[0], &w[1]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 8) & 0xffu), (aux1 >> 7) & 127u, &w[2], &w[3]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 16) & 0xffu), (aux1 >> 14) & 127u, &w[4], &w[5]);
dev_iq2_i8x8_lut(grid, signs, (uint8_t)((aux0 >> 24) & 0xffu), (aux1 >> 21) & 127u, &w[6], &w[7]);
for (uint32_t p = 0; p < n; p++) {
const int8_t *q = q8[p] + ib32 * 32;
int32_t sumi = 0;
sumi = __dp4a(w[0], *(const int32_t *)(q + 0), sumi);
sumi = __dp4a(w[1], *(const int32_t *)(q + 4), sumi);
sumi = __dp4a(w[2], *(const int32_t *)(q + 8), sumi);
sumi = __dp4a(w[3], *(const int32_t *)(q + 12), sumi);
sumi = __dp4a(w[4], *(const int32_t *)(q + 16), sumi);
sumi = __dp4a(w[5], *(const int32_t *)(q + 20), sumi);
sumi = __dp4a(w[6], *(const int32_t *)(q + 24), sumi);
sumi = __dp4a(w[7], *(const int32_t *)(q + 28), sumi);
bsum[p] += sumi * ls;
}
}
const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 };
for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p];
}
__device__ static void dev_dot_iq2_xxs_q8_K_block4(
const cuda_block_iq2_xxs *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
uint32_t n,
float acc[4]) {
const float xd = dev_f16_to_f32(x->d);
const uint16_t *q2 = x->qs;
int32_t bsum[4] = {0, 0, 0, 0};
const int8_t *q8[4] = {
y0 ? y0->qs : NULL,
y1 ? y1->qs : NULL,
y2 ? y2->qs : NULL,
y3 ? y3->qs : NULL,
};
for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const uint32_t ls = 2u * (aux1 >> 28) + 1u;
const uint8_t a0 = (uint8_t)(aux0 & 0xffu);
const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu);
const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu);
const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu);
for (uint32_t p = 0; p < n; p++) {
int32_t sumi = 0;
sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32);
sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16);
bsum[p] += sumi * (int32_t)ls;
}
}
const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 };
for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p];
}
__device__ static DS4_CUDA_UNUSED void dev_dot_iq2_xxs_q8_K_block8(
const cuda_block_iq2_xxs *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
const cuda_block_q8_K *y4,
const cuda_block_q8_K *y5,
const cuda_block_q8_K *y6,
const cuda_block_q8_K *y7,
uint32_t n,
float acc[8]) {
const float xd = dev_f16_to_f32(x->d);
const uint16_t *q2 = x->qs;
int32_t bsum[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const int8_t *q8[8] = {
y0 ? y0->qs : NULL, y1 ? y1->qs : NULL, y2 ? y2->qs : NULL, y3 ? y3->qs : NULL,
y4 ? y4->qs : NULL, y5 ? y5->qs : NULL, y6 ? y6->qs : NULL, y7 ? y7->qs : NULL,
};
for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux0 = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux1 = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const uint32_t ls = 2u * (aux1 >> 28) + 1u;
const uint8_t a0 = (uint8_t)(aux0 & 0xffu);
const uint8_t a1 = (uint8_t)((aux0 >> 8) & 0xffu);
const uint8_t a2 = (uint8_t)((aux0 >> 16) & 0xffu);
const uint8_t a3 = (uint8_t)((aux0 >> 24) & 0xffu);
for (uint32_t p = 0; p < n; p++) {
int32_t sumi = 0;
sumi += dev_dot_iq2_pair_16(a0, (aux1 >> 0) & 127u, a1, (aux1 >> 7) & 127u, q8[p] + ib32 * 32);
sumi += dev_dot_iq2_pair_16(a2, (aux1 >> 14) & 127u, a3, (aux1 >> 21) & 127u, q8[p] + ib32 * 32 + 16);
bsum[p] += sumi * (int32_t)ls;
}
}
const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 };
for (uint32_t p = 0; p < n; p++) acc[p] += 0.125f * xd * ys[p]->d * (float)bsum[p];
}
__device__ static void dev_q4_K_get_scale_min(
uint32_t j,
const uint8_t *scales,
uint8_t *d_out,
uint8_t *m_out) {
if (j < 4u) {
*d_out = scales[j] & 63u;
*m_out = scales[j + 4u] & 63u;
} else {
*d_out = (scales[j + 4u] & 0x0fu) | ((scales[j - 4u] >> 6u) << 4u);
*m_out = (scales[j + 4u] >> 4u) | ((scales[j] >> 6u) << 4u);
}
}
__device__ __forceinline__ static int32_t dev_dot_q4_32(const uint8_t *qs, const int8_t *q8, int shift) {
int32_t sum = 0;
#pragma unroll
for (uint32_t i = 0; i < 32u; i += 4u) {
const int32_t v = (*(const int32_t *)(qs + i) >> shift) & 0x0f0f0f0f;
sum = __dp4a(v, *(const int32_t *)(q8 + i), sum);
}
return sum;
}
__device__ static float dev_dot_q4_K_q8_K_block(const cuda_block_q4_K *x, const cuda_block_q8_K *y) {
const float xd = dev_f16_to_f32(x->d);
const float xmin = dev_f16_to_f32(x->dmin);
int isum = 0;
int summs = 0;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
uint8_t sc, m;
dev_q4_K_get_scale_min(j, x->scales, &sc, &m);
summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]);
const uint32_t byte_off = (j >> 1u) * 32u;
const int shift = (j & 1u) ? 4 : 0;
isum += (int)sc * dev_dot_q4_32(x->qs + byte_off, y->qs + j * 32u, shift);
}
return y->d * xd * (float)isum - y->d * xmin * (float)summs;
}
/* Vector-load variant of dev_dot_q4_K_q8_K_block: loads the whole 144-byte
* Q4_K block with nine 16B loads (requires a 16B-aligned tensor base; block
* stride 144 and row strides are 16B multiples), then computes the exact same
* integer sums and float finish. Same values in the same order, so results
* are bit-identical; the wide loads just improve DRAM/memory-level
* parallelism for the bandwidth-bound decode matvecs. */
__device__ __forceinline__ static void dev_dot_q4_K_q8_K_block_vec(
const cuda_block_q4_K *x,
const cuda_block_q8_K *y,
float *out_acc) {
const uint4 hdr = *(const uint4 *)x; /* d, dmin, scales[12] */
uint4 qv[8];
#pragma unroll
for (uint32_t i = 0; i < 8u; i++) qv[i] = ((const uint4 *)(x->qs))[i];
const uint16_t xd_u = (uint16_t)(hdr.x & 0xffffu);
const uint16_t xmin_u = (uint16_t)(hdr.x >> 16u);
uint8_t scales[12];
scales[0] = (uint8_t)(hdr.y);
scales[1] = (uint8_t)(hdr.y >> 8);
scales[2] = (uint8_t)(hdr.y >> 16);
scales[3] = (uint8_t)(hdr.y >> 24);
scales[4] = (uint8_t)(hdr.z);
scales[5] = (uint8_t)(hdr.z >> 8);
scales[6] = (uint8_t)(hdr.z >> 16);
scales[7] = (uint8_t)(hdr.z >> 24);
scales[8] = (uint8_t)(hdr.w);
scales[9] = (uint8_t)(hdr.w >> 8);
scales[10] = (uint8_t)(hdr.w >> 16);
scales[11] = (uint8_t)(hdr.w >> 24);
const float xd = dev_f16_to_f32(xd_u);
const float xmin = dev_f16_to_f32(xmin_u);
int isum = 0;
int summs = 0;
const int32_t *qw = (const int32_t *)qv;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
uint8_t sc, m;
dev_q4_K_get_scale_min(j, scales, &sc, &m);
summs += (int)m * (int)(y->bsums[2u * j] + y->bsums[2u * j + 1u]);
const uint32_t word_off = (j >> 1u) * 8u;
const int shift = (j & 1u) ? 4 : 0;
int32_t sum = 0;
#pragma unroll
for (uint32_t i = 0; i < 8u; i++) {
const int32_t v = (qw[word_off + i] >> shift) & 0x0f0f0f0f;
sum = __dp4a(v, *(const int32_t *)(y->qs + j * 32u + i * 4u), sum);
}
isum += (int)sc * sum;
}
*out_acc += y->d * xd * (float)isum - y->d * xmin * (float)summs;
}
__device__ static void dev_dot_q4_K_q8_K_block8(
const cuda_block_q4_K *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
const cuda_block_q8_K *y4,
const cuda_block_q8_K *y5,
const cuda_block_q8_K *y6,
const cuda_block_q8_K *y7,
uint32_t n,
float acc[8]) {
const float xd = dev_f16_to_f32(x->d);
const float xmin = dev_f16_to_f32(x->dmin);
const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 };
int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0};
int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0};
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
uint8_t sc, m;
dev_q4_K_get_scale_min(j, x->scales, &sc, &m);
const uint32_t byte_off = (j >> 1u) * 32u;
const int shift = (j & 1u) ? 4 : 0;
for (uint32_t p = 0; p < n; p++) {
summs[p] += (int)m * (int)(ys[p]->bsums[2u * j] + ys[p]->bsums[2u * j + 1u]);
isum[p] += (int)sc * dev_dot_q4_32(x->qs + byte_off, ys[p]->qs + j * 32u, shift);
}
}
for (uint32_t p = 0; p < n; p++) {
acc[p] += ys[p]->d * xd * (float)isum[p] - ys[p]->d * xmin * (float)summs[p];
}
}
__device__ static float dev_dot_q2_K_q8_K_block(const cuda_block_q2_K *x, const cuda_block_q8_K *y) {
const uint8_t *q2 = x->qs;
const int8_t *q8 = y->qs;
const uint8_t *sc = x->scales;
int summs = 0;
for (int j = 0; j < 16; j++) summs += y->bsums[j] * (sc[j] >> 4);
const float dall = y->d * dev_f16_to_f32(x->d);
const float dmin = y->d * dev_f16_to_f32(x->dmin);
int isum = 0;
int is = 0;
for (int k = 0; k < CUDA_QK_K / 128; k++) {
int shift = 0;
for (int j = 0; j < 4; j++) {
int d = sc[is++] & 0x0f;
isum += d * dev_dot_q2_16(q2, q8, shift);
d = sc[is++] & 0x0f;
isum += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift);
shift += 2;
q8 += 32;
}
q2 += 32;
}
return dall * (float)isum - dmin * (float)summs;
}
__device__ static void dev_dot_q2_K_q8_K_block4(
const cuda_block_q2_K *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
uint32_t n,
float acc[4]) {
const uint8_t *sc = x->scales;
const float xd = dev_f16_to_f32(x->d);
const float xmin = dev_f16_to_f32(x->dmin);
const cuda_block_q8_K *ys[4] = { y0, y1, y2, y3 };
int isum[4] = {0, 0, 0, 0};
int summs[4] = {0, 0, 0, 0};
for (uint32_t p = 0; p < n; p++) {
for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4);
}
for (uint32_t p = 0; p < n; p++) {
const uint8_t *q2 = x->qs;
const int8_t *q8 = ys[p]->qs;
int is = 0;
for (int k = 0; k < CUDA_QK_K / 128; k++) {
int shift = 0;
for (int j = 0; j < 4; j++) {
int d = sc[is++] & 0x0f;
isum[p] += d * dev_dot_q2_16(q2, q8, shift);
d = sc[is++] & 0x0f;
isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift);
shift += 2;
q8 += 32;
}
q2 += 32;
}
}
for (uint32_t p = 0; p < n; p++) {
const float yd = ys[p]->d;
acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p];
}
}
__device__ static void dev_dot_q2_K_q8_K_block8(
const cuda_block_q2_K *x,
const cuda_block_q8_K *y0,
const cuda_block_q8_K *y1,
const cuda_block_q8_K *y2,
const cuda_block_q8_K *y3,
const cuda_block_q8_K *y4,
const cuda_block_q8_K *y5,
const cuda_block_q8_K *y6,
const cuda_block_q8_K *y7,
uint32_t n,
float acc[8]) {
const uint8_t *sc = x->scales;
const float xd = dev_f16_to_f32(x->d);
const float xmin = dev_f16_to_f32(x->dmin);
const cuda_block_q8_K *ys[8] = { y0, y1, y2, y3, y4, y5, y6, y7 };
int isum[8] = {0, 0, 0, 0, 0, 0, 0, 0};
int summs[8] = {0, 0, 0, 0, 0, 0, 0, 0};
for (uint32_t p = 0; p < n; p++) {
for (int j = 0; j < 16; j++) summs[p] += ys[p]->bsums[j] * (sc[j] >> 4);
}
for (uint32_t p = 0; p < n; p++) {
const uint8_t *q2 = x->qs;
const int8_t *q8 = ys[p]->qs;
int is = 0;
for (int k = 0; k < CUDA_QK_K / 128; k++) {
int shift = 0;
for (int j = 0; j < 4; j++) {
int d = sc[is++] & 0x0f;
isum[p] += d * dev_dot_q2_16(q2, q8, shift);
d = sc[is++] & 0x0f;
isum[p] += d * dev_dot_q2_16(q2 + 16, q8 + 16, shift);
shift += 2;
q8 += 32;
}
q2 += 32;
}
}
for (uint32_t p = 0; p < n; p++) {
const float yd = ys[p]->d;
acc[p] += yd * xd * (float)isum[p] - yd * xmin * (float)summs[p];
}
}
__device__ static float half_warp_sum_f32(float v, uint32_t lane16) {
uint32_t mask = 0xffffu << (threadIdx.x & 16u);
for (int offset = 8; offset > 0; offset >>= 1) {
v += __shfl_down_sync(mask, v, offset, 16);
}
(void)lane16;
return v;
}
__device__ static float quarter_warp_sum_f32(float v, uint32_t lane8) {
uint32_t mask = 0xffu << (threadIdx.x & 24u);
for (int offset = 4; offset > 0; offset >>= 1) {
v += __shfl_down_sync(mask, v, offset, 8);
}
(void)lane8;
return v;
}
__global__ static void q8_K_quantize_kernel(cuda_block_q8_K *out, const float *x, uint32_t in_dim, uint32_t n_rows) {
uint32_t b = blockIdx.x;
uint32_t row = blockIdx.y;
if (row >= n_rows || b >= in_dim / CUDA_QK_K) return;
const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K;
cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b;
__shared__ float abs_part[256];
__shared__ float val_part[256];
__shared__ float maxv_s;
__shared__ float iscale_s;
uint32_t tid = threadIdx.x;
float v = tid < CUDA_QK_K ? xr[tid] : 0.0f;
abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f;
val_part[tid] = v;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (tid < stride && abs_part[tid + stride] > abs_part[tid]) {
abs_part[tid] = abs_part[tid + stride];
val_part[tid] = val_part[tid + stride];
}
__syncthreads();
}
float amax = abs_part[0];
if (amax == 0.0f) {
if (tid == 0) yb->d = 0.0f;
if (tid < CUDA_QK_K) yb->qs[tid] = 0;
if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0;
return;
}
if (tid == 0) {
maxv_s = val_part[0];
iscale_s = -127.0f / maxv_s;
}
__syncthreads();
if (tid < CUDA_QK_K) {
int qv = (int)lrintf(iscale_s * xr[tid]);
if (qv > 127) qv = 127;
if (qv < -128) qv = -128;
yb->qs[tid] = (int8_t)qv;
}
__syncthreads();
if (tid < CUDA_QK_K / 16) {
int sum = 0;
for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i];
yb->bsums[tid] = (int16_t)sum;
}
if (tid == 0) yb->d = 1.0f / iscale_s;
}
/* Decode-only dual quantizer. The Q8_0 half mirrors
* quantize_q8_0_f32_kernel's 32-thread reduction and expression order, while
* the Q8_K half remains byte-for-byte the ordinary routed-MoE quantizer. */
__global__ static void q8_K_q8_0_quantize_kernel(
cuda_block_q8_K *out,
int8_t *q8_0,
float *q8_0_scale,
const float *x,
uint32_t in_dim,
uint32_t n_rows) {
const uint32_t b = blockIdx.x;
const uint32_t row = blockIdx.y;
if (row >= n_rows || b >= in_dim / CUDA_QK_K) return;
const float *xr = x + (uint64_t)row * in_dim +
(uint64_t)b * CUDA_QK_K;
cuda_block_q8_K *yb = out +
(uint64_t)row * (in_dim / CUDA_QK_K) + b;
__shared__ float abs_part[256];
__shared__ float val_part[256];
__shared__ float maxv_s;
__shared__ float iscale_s;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f;
abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f;
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (lane < stride) {
abs_part[tid] = fmaxf(abs_part[tid], abs_part[tid + stride]);
}
__syncthreads();
}
const uint32_t q8_blocks = in_dim / 32u;
const uint32_t q8_block = b * 8u + warp;
const float d = abs_part[warp * 32u] / 127.0f;
const float id = d != 0.0f ? 1.0f / d : 0.0f;
if (lane == 0u) {
q8_0_scale[(uint64_t)row * q8_blocks + q8_block] = d;
}
int qv = (int)lrintf(v * id);
qv = qv > 127 ? 127 : (qv < -128 ? -128 : qv);
q8_0[((uint64_t)row * q8_blocks + q8_block) * 32u + lane] =
(int8_t)qv;
__syncthreads();
abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f;
val_part[tid] = v;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (tid < stride && abs_part[tid + stride] > abs_part[tid]) {
abs_part[tid] = abs_part[tid + stride];
val_part[tid] = val_part[tid + stride];
}
__syncthreads();
}
const float amax = abs_part[0];
if (amax == 0.0f) {
if (tid == 0) yb->d = 0.0f;
if (tid < CUDA_QK_K) yb->qs[tid] = 0;
if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0;
return;
}
if (tid == 0) {
maxv_s = val_part[0];
iscale_s = -127.0f / maxv_s;
}
__syncthreads();
if (tid < CUDA_QK_K) {
int kv = (int)lrintf(iscale_s * xr[tid]);
if (kv > 127) kv = 127;
if (kv < -128) kv = -128;
yb->qs[tid] = (int8_t)kv;
}
__syncthreads();
if (tid < CUDA_QK_K / 16) {
int sum = 0;
for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i];
yb->bsums[tid] = (int16_t)sum;
}
if (tid == 0) yb->d = 1.0f / iscale_s;
}
__device__ __forceinline__ static bool moe_owned_local_expert(
int32_t expert,
uint32_t expert_base,
uint32_t expert_count,
uint32_t *local_expert) {
if (expert < 0) return false;
const uint32_t e = (uint32_t)expert;
if (e < expert_base || e - expert_base >= expert_count) return false;
if (local_expert) *local_expert = e - expert_base;
return true;
}
__global__ static void mxfp4_prepare_owned_assignments_kernel(
int32_t *local_ids,
float *local_weights,
const int32_t *selected,
const float *weights,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t slot = threadIdx.x;
if (slot >= 6u) return;
uint32_t local_expert = 0;
const bool owned = moe_owned_local_expert(
selected[slot], expert_base, expert_count, &local_expert);
local_ids[slot] = owned ? (int32_t)local_expert : -1;
local_weights[slot] = owned ? weights[slot] : 0.0f;
}
/* Quantize only selected slots owned by this expert-parallel rank. Rows keep
* their original slot index so the final rank-local reduction can visit slots
* in canonical order without compaction or a host synchronization. */
__global__ static void q8_K_quantize_owned_kernel(
cuda_block_q8_K *out,
const float *x,
const int32_t *selected,
uint32_t in_dim,
uint32_t n_rows,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t b = blockIdx.x;
const uint32_t row = blockIdx.y;
if (row >= n_rows || b >= in_dim / CUDA_QK_K) return;
if (!moe_owned_local_expert(selected[row], expert_base, expert_count, NULL)) return;
const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K;
cuda_block_q8_K *yb = out + (uint64_t)row * (in_dim / CUDA_QK_K) + b;
__shared__ float abs_part[256];
__shared__ float val_part[256];
__shared__ float maxv_s;
__shared__ float iscale_s;
const uint32_t tid = threadIdx.x;
const float v = tid < CUDA_QK_K ? xr[tid] : 0.0f;
abs_part[tid] = tid < CUDA_QK_K ? fabsf(v) : 0.0f;
val_part[tid] = v;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (tid < stride && abs_part[tid + stride] > abs_part[tid]) {
abs_part[tid] = abs_part[tid + stride];
val_part[tid] = val_part[tid + stride];
}
__syncthreads();
}
const float amax = abs_part[0];
if (amax == 0.0f) {
if (tid == 0) yb->d = 0.0f;
if (tid < CUDA_QK_K) yb->qs[tid] = 0;
if (tid < CUDA_QK_K / 16) yb->bsums[tid] = 0;
return;
}
if (tid == 0) {
maxv_s = val_part[0];
iscale_s = -127.0f / maxv_s;
}
__syncthreads();
if (tid < CUDA_QK_K) {
int qv = (int)lrintf(iscale_s * xr[tid]);
if (qv > 127) qv = 127;
if (qv < -128) qv = -128;
yb->qs[tid] = (int8_t)qv;
}
__syncthreads();
if (tid < CUDA_QK_K / 16) {
int sum = 0;
for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16 + i];
yb->bsums[tid] = (int16_t)sum;
}
if (tid == 0) yb->d = 1.0f / iscale_s;
}
__global__ static void moe_filter_owned_pairs_kernel(
int32_t *selected,
float *weights,
uint64_t pair_count,
uint32_t n_total_expert,
uint32_t expert_base,
uint32_t expert_count) {
const uint64_t pair = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (pair >= pair_count) return;
const int32_t expert_i = selected[pair];
if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert &&
(uint32_t)expert_i >= expert_base &&
(uint32_t)expert_i - expert_base < expert_count) {
selected[pair] = expert_i - (int32_t)expert_base;
} else {
selected[pair] = -1;
weights[pair] = 0.0f;
}
}
__global__ static void q8_K_quantize_sidecar_kernel(
cuda_block_q8_K *out,
const float *x,
const float *amax_sidecar,
uint32_t in_dim,
uint32_t n_rows) {
uint32_t b = blockIdx.x;
uint32_t row = blockIdx.y;
const uint32_t blocks = in_dim / CUDA_QK_K;
if (row >= n_rows || b >= blocks) return;
const float *xr = x + (uint64_t)row * in_dim + (uint64_t)b * CUDA_QK_K;
const float *sc = amax_sidecar + ((uint64_t)row * blocks + b) * 32u;
cuda_block_q8_K *yb = out + (uint64_t)row * blocks + b;
__shared__ float abs_part[32];
__shared__ float val_part[32];
__shared__ float iscale_s;
const uint32_t tid = threadIdx.x;
if (tid < 32u) {
const float v = sc[tid];
abs_part[tid] = fabsf(v);
val_part[tid] = v;
}
__syncthreads();
for (uint32_t stride = 16u; stride > 0u; stride >>= 1u) {
if (tid < stride && abs_part[tid + stride] > abs_part[tid]) {
abs_part[tid] = abs_part[tid + stride];
val_part[tid] = val_part[tid + stride];
}
__syncthreads();
}
const float amax = abs_part[0];
if (amax == 0.0f) {
if (tid == 0u) yb->d = 0.0f;
if (tid < CUDA_QK_K) yb->qs[tid] = 0;
if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0;
return;
}
if (tid == 0u) iscale_s = -127.0f / val_part[0];
__syncthreads();
if (tid < CUDA_QK_K) {
int qv = (int)lrintf(iscale_s * xr[tid]);
if (qv > 127) qv = 127;
if (qv < -128) qv = -128;
yb->qs[tid] = (int8_t)qv;
}
__syncthreads();
if (tid < CUDA_QK_K / 16u) {
int sum = 0;
for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i];
yb->bsums[tid] = (int16_t)sum;
}
if (tid == 0u) yb->d = 1.0f / iscale_s;
}
__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t row = blockIdx.x;
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = threadIdx.x; b < xq_blocks; b += blockDim.x) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
__shared__ float partial_gate[256];
__shared__ float partial_up[256];
partial_gate[threadIdx.x] = gate;
partial_up[threadIdx.x] = up;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride];
partial_up[threadIdx.x] += partial_up[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
gate = partial_gate[0];
up = partial_up[0];
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_warp8_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t lane = threadIdx.x & 31u;
uint32_t warp = threadIdx.x >> 5u;
uint32_t row = blockIdx.x * 8u + warp;
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_hwarp16_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t lane = threadIdx.x & 15u;
uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u);
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 16u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = half_warp_sum_f32(gate, lane);
up = half_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
// perf-04: launch-geometry tuning for the routed-MoE gate/up decode kernels
// (moe_gate_up_mid_qwarp32 / _decode_lut_qwarp32 / _decode_q4K_qwarp32). Each
// block processes MOE_DECODE_ROW_TILES tiles of 32 rows (row_lane in [0,32)).
// The historical value was 4 (128 rows/block -> ~96 blocks, occupancy ~16%,
// "grid too small to fill the device"). Lowering it issues correspondingly more
// blocks (e.g. 1 tile -> 32 rows/block -> ~4x more blocks -> ~384) to fill the
// SMs. The per-row arithmetic is identical regardless of this value, so output
// is bit-identical; only the qgrid.x divisor must match MOE_DECODE_ROWS_PER_BLOCK.
#ifndef MOE_DECODE_ROW_TILES
#define MOE_DECODE_ROW_TILES 1u
#endif
#define MOE_DECODE_ROWS_PER_BLOCK (32u * MOE_DECODE_ROW_TILES)
__global__ static void moe_gate_up_mid_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t pair = blockIdx.y;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) {
for (uint32_t i = threadIdx.x; i < MOE_DECODE_ROWS_PER_BLOCK; i += blockDim.x) {
const uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + i;
if (row >= expert_mid_dim) continue;
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = up_out[off] = mid_out[off] = 0.0f;
}
return;
}
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) {
uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
}
__global__ static void moe_gate_up_mid_decode_lut_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t pair = blockIdx.y;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) {
for (uint32_t i = threadIdx.x; i < MOE_DECODE_ROWS_PER_BLOCK; i += blockDim.x) {
const uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + i;
if (row >= expert_mid_dim) continue;
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) gate_out[off] = up_out[off] = 0.0f;
mid_out[off] = 0.0f;
}
return;
}
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[32];
__shared__ uint64_t s_iq2_grid[256];
__shared__ uint8_t s_iq2_signs[128];
if (xq_blocks <= 32u) {
/* Cooperatively copy complete Q8_K blocks, including their scales. */
const uint32_t words = xq_blocks * sizeof(cuda_block_q8_K) / sizeof(uint32_t);
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x)
((uint32_t *)sxq)[i] = ((const uint32_t *)xqb)[i];
xqb = sxq;
}
/* Wide inputs bypass the shared activation cache, not the lookup tables. */
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) {
uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs);
up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
}
__global__ static void moe_gate_up_mid_decode_lut_owned_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t expert_base,
uint32_t expert_count,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t pair = blockIdx.y;
uint32_t expert = 0u;
if (!moe_owned_local_expert(selected[pair], expert_base, expert_count,
&expert)) return;
const cuda_block_q8_K *xqb = xq;
__shared__ cuda_block_q8_K sxq[32];
__shared__ uint64_t s_iq2_grid[256];
__shared__ uint8_t s_iq2_signs[128];
if (xq_blocks <= 32u) {
const uint32_t words = xq_blocks * sizeof(cuda_block_q8_K) / sizeof(uint32_t);
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x)
((uint32_t *)sxq)[i] = ((const uint32_t *)xqb)[i];
xqb = sxq;
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) {
uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block_lut(gr + b, xqb + b, s_iq2_grid, s_iq2_signs);
up += dev_dot_iq2_xxs_q8_K_block_lut(ur + b, xqb + b, s_iq2_grid, s_iq2_signs);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair];
}
}
}
__device__ __forceinline__ static int32_t iq2_aligned_q8k_partial(
uint2 packed, const int8_t *q8, const uint64_t *grid, const uint8_t *signs) {
int32_t sum = 0;
#pragma unroll
for (uint32_t i = 0; i < 4u; i++) {
int32_t w0, w1;
dev_iq2_i8x8_lut(grid, signs, (packed.x >> (i * 8u)) & 255u,
(packed.y >> (i * 7u)) & 127u, &w0, &w1);
sum = __dp4a(w0, ((const int32_t *)q8)[i * 2u], sum);
sum = __dp4a(w1, ((const int32_t *)q8)[i * 2u + 1u], sum);
}
return sum * (int32_t)(2u * (packed.y >> 28) + 1u);
}
/* Eight lanes unpack one IQ2 block. Integer sums are exact; paired block
* accumulators reproduce the original eight-lane floating reduction. */
__global__ static void moe_gate_up_mid_aligned_q8k_kernel(
float *mid, const char *gate, const char *up, const cuda_block_q8_K *xq,
const int32_t *selected, const float *weights, uint32_t blocks,
uint32_t rows, uint32_t slots, uint64_t scale_bytes, float clamp) {
const uint32_t pair = blockIdx.y, token = pair / slots;
const int32_t expert = selected[pair];
const uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u);
if (expert < 0) {
if ((threadIdx.x & 31u) == 0 && row < rows) mid[(uint64_t)pair * rows + row] = 0;
return;
}
const cuda_block_q8_K *activation = xq + (uint64_t)token * blocks;
__shared__ cuda_block_q8_K sxq[32];
__shared__ uint64_t grid[256];
__shared__ uint8_t signs[128];
if (blocks <= 32u) {
for (uint32_t i = threadIdx.x; i < blocks * sizeof(cuda_block_q8_K) / 4u; i += blockDim.x)
((uint32_t *)sxq)[i] = ((const uint32_t *)activation)[i];
activation = sxq;
}
grid[threadIdx.x] = cuda_iq2xxs_grid[threadIdx.x];
if (threadIdx.x < 128u) signs[threadIdx.x] = cuda_ksigns_iq2xs[threadIdx.x];
__syncthreads();
if (row >= rows) return;
const uint32_t lane = threadIdx.x & 31u, sub = lane & 7u, group = lane >> 3u;
const uint64_t base = ((uint64_t)(uint32_t)expert * rows + row) * blocks;
float g[2] = {0, 0}, u[2] = {0, 0};
for (uint32_t k = 0; k < blocks; k += 8u) {
#pragma unroll
for (uint32_t half = 0; half < 2u; half++) {
const uint32_t b = k + group + half * 4u;
int32_t gs = 0, us = 0;
if (b < blocks) {
const uint64_t off = (base + b) * 8u + sub;
const int8_t *q8 = activation[b].qs + sub * 32u;
gs = iq2_aligned_q8k_partial(((const uint2 *)(gate + scale_bytes))[off], q8, grid, signs);
us = iq2_aligned_q8k_partial(((const uint2 *)(up + scale_bytes))[off], q8, grid, signs);
}
#pragma unroll
for (uint32_t stride = 4; stride; stride >>= 1) {
gs += __shfl_down_sync(0xffffffffu, gs, stride, 8);
us += __shfl_down_sync(0xffffffffu, us, stride, 8);
}
if (sub == 0 && b < blocks) {
g[half] += 0.125f * dev_f16_to_f32(((const uint16_t *)gate)[base + b]) * activation[b].d * (float)gs;
u[half] += 0.125f * dev_f16_to_f32(((const uint16_t *)up)[base + b]) * activation[b].d * (float)us;
}
}
}
float gv = g[0] + g[1], uv = u[0] + u[1];
gv += __shfl_down_sync(0xffffffffu, gv, 16);
uv += __shfl_down_sync(0xffffffffu, uv, 16);
gv += __shfl_down_sync(0xffffffffu, gv, 8);
uv += __shfl_down_sync(0xffffffffu, uv, 8);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gv > clamp) gv = clamp;
if (uv > clamp) uv = clamp;
if (uv < -clamp) uv = -clamp;
}
mid[(uint64_t)pair * rows + row] = (gv / (1.0f + expf(-gv))) * uv * weights[pair];
}
}
__global__ static void moe_count_sorted_pairs_kernel(
uint32_t *counts,
const int32_t *selected,
uint32_t pair_count,
uint32_t n_total_expert) {
uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
if (pair >= pair_count) return;
int32_t expert_i = selected[pair];
if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return;
atomicAdd(counts + (uint32_t)expert_i, 1u);
}
__global__ static void moe_prefix_sorted_pairs_kernel(
uint32_t *offsets,
uint32_t *cursors,
const uint32_t *counts,
uint32_t n_total_expert) {
if (threadIdx.x == 0) {
uint32_t sum = 0;
for (uint32_t e = 0; e < n_total_expert; e++) {
offsets[e] = sum;
cursors[e] = sum;
sum += counts[e];
}
offsets[n_total_expert] = sum;
}
}
__global__ static void moe_scatter_sorted_pairs_kernel(
uint32_t *sorted_pairs,
uint32_t *cursors,
const int32_t *selected,
uint32_t pair_count,
uint32_t n_total_expert) {
uint32_t pair = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
if (pair >= pair_count) return;
int32_t expert_i = selected[pair];
if (expert_i < 0 || (uint32_t)expert_i >= n_total_expert) return;
uint32_t pos = atomicAdd(cursors + (uint32_t)expert_i, 1u);
sorted_pairs[pos] = pair;
}
__global__ static void moe_build_expert_tile_offsets_kernel(
uint32_t *tile_offsets,
uint32_t *tile_total,
const uint32_t *counts,
uint32_t block_m,
uint32_t n_total_expert) {
if (threadIdx.x == 0) {
uint32_t sum = 0;
for (uint32_t e = 0; e < n_total_expert; e++) {
tile_offsets[e] = sum;
sum += (counts[e] + block_m - 1u) / block_m;
}
tile_offsets[n_total_expert] = sum;
*tile_total = sum;
}
}
__global__ static void moe_build_expert_tiles_kernel(
uint32_t *tile_experts,
uint32_t *tile_starts,
const uint32_t *tile_offsets,
const uint32_t *counts,
uint32_t block_m,
uint32_t n_total_expert) {
uint32_t e = (uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
if (e >= n_total_expert) return;
uint32_t ntiles = (counts[e] + block_m - 1u) / block_m;
uint32_t off = tile_offsets[e];
for (uint32_t t = 0; t < ntiles; t++) {
tile_experts[off + t] = e;
tile_starts[off + t] = t * block_m;
}
}
/* Decode-sized routed batches spend more host time launching metadata kernels
* than doing the <= 96 pair / 128 expert setup. Build both tile lists in one
* deterministic block; the expensive expert kernels remain unchanged. */
__global__ static void moe_prepare_sorted_tiles_small_kernel(
uint32_t *counts,
uint32_t *offsets,
uint32_t *cursors,
uint32_t *sorted_pairs,
uint32_t *tile_offsets,
uint32_t *tile_total,
uint32_t *tile_experts,
uint32_t *tile_starts,
uint32_t *tile16_offsets,
uint32_t *tile16_total,
uint32_t *tile16_experts,
uint32_t *tile16_starts,
const int32_t *selected,
uint32_t pair_count,
uint32_t n_total_expert,
uint32_t block_m,
bool build_tile16) {
if (blockIdx.x != 0) return;
const uint32_t tid = threadIdx.x;
__shared__ uint32_t local_counts[128];
__shared__ int32_t local_selected[96];
for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) {
local_counts[e] = 0u;
}
for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) {
local_selected[pair] = selected[pair];
}
__syncthreads();
for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) {
const int32_t expert_i = local_selected[pair];
if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) {
atomicAdd(local_counts + (uint32_t)expert_i, 1u);
}
}
__syncthreads();
for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) {
counts[e] = local_counts[e];
}
if (tid == 0u) {
uint32_t pair_sum = 0u;
uint32_t tile_sum = 0u;
uint32_t tile16_sum = 0u;
for (uint32_t e = 0; e < n_total_expert; e++) {
const uint32_t count = local_counts[e];
offsets[e] = pair_sum;
pair_sum += count;
cursors[e] = pair_sum;
tile_offsets[e] = tile_sum;
tile_sum += (count + block_m - 1u) / block_m;
if (build_tile16) {
tile16_offsets[e] = tile16_sum;
tile16_sum += (count + 15u) / 16u;
}
}
offsets[n_total_expert] = pair_sum;
tile_offsets[n_total_expert] = tile_sum;
*tile_total = tile_sum;
if (build_tile16) {
tile16_offsets[n_total_expert] = tile16_sum;
*tile16_total = tile16_sum;
}
}
__syncthreads();
for (uint32_t pair = tid; pair < pair_count; pair += blockDim.x) {
const int32_t expert_i = local_selected[pair];
if (expert_i >= 0 && (uint32_t)expert_i < n_total_expert) {
uint32_t rank = 0u;
for (uint32_t prev = 0; prev < pair; prev++) {
rank += local_selected[prev] == expert_i;
}
sorted_pairs[offsets[(uint32_t)expert_i] + rank] = pair;
}
}
for (uint32_t e = tid; e < n_total_expert; e += blockDim.x) {
const uint32_t count = local_counts[e];
const uint32_t ntiles = (count + block_m - 1u) / block_m;
const uint32_t tile_off = tile_offsets[e];
for (uint32_t t = 0; t < ntiles; t++) {
tile_experts[tile_off + t] = e;
tile_starts[tile_off + t] = t * block_m;
}
if (build_tile16) {
const uint32_t ntiles16 = (count + 15u) / 16u;
const uint32_t tile16_off = tile16_offsets[e];
for (uint32_t t = 0; t < ntiles16; t++) {
tile16_experts[tile16_off + t] = e;
tile16_starts[tile16_off + t] = t * 16u;
}
}
}
}
__global__ static void moe_gate_up_mid_sorted_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t pair = sorted_pairs[blockIdx.y];
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static DS4_CUDA_UNUSED void moe_gate_up_mid_expert_tile8_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t group = threadIdx.x >> 3u;
uint32_t lane = threadIdx.x & 7u;
uint32_t pair_slot = group & 7u;
uint32_t row_lane = group >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_pair = tile_starts[tile] + pair_slot;
if (local_pair >= counts[expert]) return;
uint32_t sorted_idx = offsets[expert] + local_pair;
uint32_t pair = sorted_pairs[sorted_idx];
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
for (uint32_t rr = 0; rr < 2u; rr++) {
uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
}
__global__ static void moe_gate_up_mid_expert_tile4_row32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[4][16];
uint32_t pair[4] = {0, 0, 0, 0};
uint32_t tok[4] = {0, 0, 0, 0};
uint32_t slot[4] = {0, 0, 0, 0};
const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 4u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
tok[np] = pair[np] / n_expert;
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
if (row >= expert_mid_dim) return;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float up[4] = {0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
dev_dot_iq2_xxs_q8_K_block4(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, gate);
dev_dot_iq2_xxs_q8_K_block4(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, up);
}
for (uint32_t p = 0; p < np; p++) {
gate[p] = quarter_warp_sum_f32(gate[p], lane);
up[p] = quarter_warp_sum_f32(up[p], lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate[p] > clamp) gate[p] = clamp;
if (up[p] > clamp) up[p] = clamp;
if (up[p] < -clamp) up[p] = -clamp;
}
const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate[p];
up_out[off] = up[p];
}
mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]];
}
}
}
__global__ static void moe_gate_up_mid_expert_tile8_row32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][16];
__shared__ uint64_t s_iq2_grid[256];
__shared__ uint8_t s_iq2_signs[128];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
tok[np] = pair[np] / n_expert;
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
if (row >= expert_mid_dim) return;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate,
s_iq2_grid, s_iq2_signs);
dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up,
s_iq2_grid, s_iq2_signs);
}
for (uint32_t p = 0; p < np; p++) {
gate[p] = quarter_warp_sum_f32(gate[p], lane);
up[p] = quarter_warp_sum_f32(up[p], lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate[p] > clamp) gate[p] = clamp;
if (up[p] > clamp) up[p] = clamp;
if (up[p] < -clamp) up[p] = -clamp;
}
const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate[p];
up_out[off] = up[p];
}
mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]];
}
}
}
__global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][16];
__shared__ uint64_t s_iq2_grid[256];
__shared__ uint8_t s_iq2_signs[128];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
tok[np] = pair[np] / n_expert;
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
for (uint32_t rr = 0; rr < 64u; rr++) {
uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate,
s_iq2_grid, s_iq2_signs);
dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up,
s_iq2_grid, s_iq2_signs);
}
for (uint32_t p = 0; p < np; p++) {
gate[p] = quarter_warp_sum_f32(gate[p], lane);
up[p] = quarter_warp_sum_f32(up[p], lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate[p] > clamp) gate[p] = clamp;
if (up[p] > clamp) up[p] = clamp;
if (up[p] < -clamp) up[p] = -clamp;
}
const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate[p];
up_out[off] = up[p];
}
mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]];
}
}
}
}
template <uint32_t ROW_SPAN>
__global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][16];
__shared__ uint64_t s_iq2_grid[256];
__shared__ uint8_t s_iq2_signs[128];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
tok[np] = pair[np] / n_expert;
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) {
uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
dev_dot_iq2_xxs_q8_K_block8_deq_lut(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate,
s_iq2_grid, s_iq2_signs);
dev_dot_iq2_xxs_q8_K_block8_deq_lut(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up,
s_iq2_grid, s_iq2_signs);
}
for (uint32_t p = 0; p < np; p++) {
gate[p] = quarter_warp_sum_f32(gate[p], lane);
up[p] = quarter_warp_sum_f32(up[p], lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate[p] > clamp) gate[p] = clamp;
if (up[p] > clamp) up[p] = clamp;
if (up[p] < -clamp) up[p] = -clamp;
}
const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate[p];
up_out[off] = up[p];
}
mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]];
}
}
}
}
__global__ static void moe_gate_up_mid_sorted_p2_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t pair_count,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t pair_lane = (threadIdx.x >> 3u) & 1u;
uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u);
uint32_t sorted_idx = blockIdx.y * 2u + pair_lane;
if (row >= expert_mid_dim || sorted_idx >= pair_count) return;
uint32_t pair = sorted_pairs[sorted_idx];
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_iq2_xxs_q8_K_block(gr + b, xqb + b);
up += dev_dot_iq2_xxs_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static DS4_CUDA_UNUSED void moe_down_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t row = blockIdx.x;
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = threadIdx.x; b < midq_blocks; b += blockDim.x) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
__shared__ float partial[256];
partial[threadIdx.x] = acc;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0];
}
__global__ static DS4_CUDA_UNUSED void moe_down_warp8_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t lane = threadIdx.x & 31u;
uint32_t warp = threadIdx.x >> 5u;
uint32_t row = blockIdx.x * 8u + warp;
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 32u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = warp_sum_f32(acc);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
__global__ static DS4_CUDA_UNUSED void moe_down_hwarp16_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t lane = threadIdx.x & 15u;
uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u);
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 16u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = half_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
__global__ static void moe_down_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
__global__ static void moe_gate_up_mid_decode_q4K_qwarp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t pair = blockIdx.y;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) {
uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
}
__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 15u;
uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u);
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 16u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = half_warp_sum_f32(gate, lane);
up = half_warp_sum_f32(up, lane);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 15u;
uint32_t group = threadIdx.x >> 4u;
uint32_t pair = blockIdx.y;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
if (group >= 8u) return;
uint32_t row = blockIdx.x * 8u + group;
if (row >= expert_mid_dim) return;
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 16u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = half_warp_sum_f32(gate, lane);
up = half_warp_sum_f32(up, lane);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_gate_up_mid_decode_q4K_warp32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 31u;
uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u);
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_kernel(
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t lane = threadIdx.x & 31u;
uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u);
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
/* Word-wise cooperative staging copy (same bytes, all lanes busy). */
const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
uint32_t *dst = (uint32_t *)sxq;
const uint32_t *srcw = (const uint32_t *)xqb;
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i];
__syncthreads();
xqb = sxq;
}
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base |
gate_row_bytes | gate_expert_bytes) & 15u) == 0u);
if (vec_ok) {
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate);
dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up);
}
} else {
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel(
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t expert_base,
uint32_t expert_count,
float clamp) {
uint32_t lane = threadIdx.x & 31u;
uint32_t row = blockIdx.x * 8u + (threadIdx.x >> 5u);
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t expert = 0u;
if (!moe_owned_local_expert(selected[pair], expert_base, expert_count,
&expert)) return;
const cuda_block_q8_K *xqb = xq;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
const uint32_t words = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
uint32_t *dst = (uint32_t *)sxq;
const uint32_t *srcw = (const uint32_t *)xqb;
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) dst[i] = srcw[i];
__syncthreads();
xqb = sxq;
}
const bool vec_ok = ((((uintptr_t)gate_base | (uintptr_t)up_base |
gate_row_bytes | gate_expert_bytes) & 15u) == 0u);
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
if (vec_ok) {
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
dev_dot_q4_K_q8_K_block_vec(gr + b, xqb + b, &gate);
dev_dot_q4_K_q8_K_block_vec(ur + b, xqb + b, &up);
}
} else {
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[pair];
}
}
__global__ static void moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel(
float *mid_out,
float *amax_sidecar,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t row = blockIdx.x * 8u + warp;
const uint32_t pair = blockIdx.y;
const uint32_t tok = pair / n_expert;
const uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
__shared__ float tile_vals[8];
__shared__ float tile_abs[8];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
float midv = 0.0f;
const bool valid = row < expert_mid_dim;
if (valid) {
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
midv = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
mid_out[off] = midv;
}
}
if (lane == 0u) {
tile_vals[warp] = midv;
tile_abs[warp] = valid ? fabsf(midv) : 0.0f;
}
__syncthreads();
if (threadIdx.x == 0u) {
float best_abs = tile_abs[0];
float best_val = tile_vals[0];
#pragma unroll
for (uint32_t i = 1u; i < 8u; i++) {
if (tile_abs[i] > best_abs) {
best_abs = tile_abs[i];
best_val = tile_vals[i];
}
}
const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K;
const uint32_t qblock = blockIdx.x / 32u;
const uint32_t tile = blockIdx.x & 31u;
if (qblock < midq_blocks) {
amax_sidecar[((uint64_t)pair * midq_blocks + qblock) * 32u + tile] = best_val;
}
}
}
__global__ static void moe_gate_up_mid_decode_q4K_warp32_row16_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t lane = threadIdx.x & 31u;
uint32_t warp = threadIdx.x >> 5u;
uint32_t row = blockIdx.x * 16u + warp;
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = warp_sum_f32(gate);
up = warp_sum_f32(up);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_gate_up_midq_decode_q4K_qwarp32_kernel(
float *mid_out,
cuda_block_q8_K *midq,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row_lane = threadIdx.x >> 3u;
const uint32_t qblock = blockIdx.x;
const uint32_t pair = blockIdx.y;
const uint32_t tok = pair / n_expert;
const uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const uint32_t expert = (uint32_t)expert_i;
const cuda_block_q8_K *xqb = xq + (uint64_t)tok * xq_blocks;
__shared__ cuda_block_q8_K sxq[16];
__shared__ float vals[CUDA_QK_K];
__shared__ float abs_part[CUDA_QK_K];
__shared__ float val_part[CUDA_QK_K];
__shared__ float iscale_s;
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) sxq[i] = xqb[i];
__syncthreads();
xqb = sxq;
}
const float w = weights[(uint64_t)tok * n_expert + slot];
#pragma unroll
for (uint32_t rr = 0; rr < 8u; rr++) {
const uint32_t row_in_block = row_lane + rr * 32u;
const uint32_t row = qblock * CUDA_QK_K + row_in_block;
float midv = 0.0f;
if (row < expert_mid_dim) {
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
gate += dev_dot_q4_K_q8_K_block(gr + b, xqb + b);
up += dev_dot_q4_K_q8_K_block(ur + b, xqb + b);
}
gate = quarter_warp_sum_f32(gate, lane);
up = quarter_warp_sum_f32(up, lane);
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
midv = (gate / (1.0f + expf(-gate))) * up * w;
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
mid_out[off] = midv;
}
}
if (lane == 0u) vals[row_in_block] = midv;
}
__syncthreads();
cuda_block_q8_K *yb = midq + (uint64_t)pair * (expert_mid_dim / CUDA_QK_K) + qblock;
const uint32_t tid = threadIdx.x;
const float v = vals[tid];
abs_part[tid] = fabsf(v);
val_part[tid] = v;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (tid < stride && abs_part[tid + stride] > abs_part[tid]) {
abs_part[tid] = abs_part[tid + stride];
val_part[tid] = val_part[tid + stride];
}
__syncthreads();
}
const float amax = abs_part[0];
if (amax == 0.0f) {
if (tid == 0u) yb->d = 0.0f;
if (tid < CUDA_QK_K) yb->qs[tid] = 0;
if (tid < CUDA_QK_K / 16u) yb->bsums[tid] = 0;
return;
}
if (tid == 0u) {
iscale_s = -127.0f / val_part[0];
}
__syncthreads();
int qv = (int)lrintf(iscale_s * v);
if (qv > 127) qv = 127;
if (qv < -128) qv = -128;
yb->qs[tid] = (int8_t)qv;
__syncthreads();
if (tid < CUDA_QK_K / 16u) {
int sum = 0;
for (int i = 0; i < 16; i++) sum += yb->qs[tid * 16u + (uint32_t)i];
yb->bsums[tid] = (int16_t)sum;
}
if (tid == 0u) yb->d = 1.0f / iscale_s;
}
template <uint32_t ROW_SPAN>
__global__ static void moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][16];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t tok[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t slot[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
tok[np] = pair[np] / n_expert;
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) {
uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u;
if (row >= expert_mid_dim) continue;
const cuda_block_q4_K *gr = (const cuda_block_q4_K *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_q4_K *ur = (const cuda_block_q4_K *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
float gate[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float up[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 8u) {
dev_dot_q4_K_q8_K_block8(gr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, gate);
dev_dot_q4_K_q8_K_block8(ur + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, up);
}
for (uint32_t p = 0; p < np; p++) {
gate[p] = quarter_warp_sum_f32(gate[p], lane);
up[p] = quarter_warp_sum_f32(up[p], lane);
if (lane == 0) {
if (clamp > 1.0e-6f) {
if (gate[p] > clamp) gate[p] = clamp;
if (up[p] > clamp) up[p] = clamp;
if (up[p] < -clamp) up[p] = -clamp;
}
const uint64_t off = (uint64_t)pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate[p];
up_out[off] = up[p];
}
mid_out[off] = (gate[p] / (1.0f + expf(-gate[p]))) * up[p] * weights[(uint64_t)tok[p] * n_expert + slot[p]];
}
}
}
}
template<uint32_t SLOTS, bool ALIGNED = false>
__global__ static void moe_down_sum_qwarp32_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t total_experts = 0) {
out += (uint64_t)blockIdx.y * out_dim;
selected += (uint64_t)blockIdx.y * SLOTS;
midq += (uint64_t)blockIdx.y * SLOTS * midq_blocks;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
if (row >= out_dim) return;
float total = 0.0f;
#pragma unroll
for (uint32_t slot = 0; slot < SLOTS; slot++) {
int32_t expert_i = selected[slot];
/* Unowned slots may have no initialized quantized intermediate. */
if (expert_i < 0) continue;
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
if constexpr (ALIGNED) {
const uint64_t pairs = (uint64_t)total_experts * (out_dim / 2u) * midq_blocks;
const uint64_t dm_bytes = (pairs * 8u + 63u) & ~63ull;
const uint64_t sc_bytes = (pairs * 32u + 63u) & ~63ull;
const uint64_t pair = ((uint64_t)(uint32_t)expert_i * (out_dim / 2u) + row / 2u) * midq_blocks + b;
const uint32_t parity = row & 1u;
cuda_block_q2_K w;
const uint32_t dm = ((const uint32_t *)down_base)[pair * 2u + parity];
w.d = (uint16_t)dm;
w.dmin = (uint16_t)(dm >> 16);
#pragma unroll
for (uint32_t i = 0; i < 4; i++) {
const uint32_t sc = ((const uint32_t *)(down_base + dm_bytes))[
pair * 8u + (i >> 1) * 4u + parity * 2u + (i & 1u)];
memcpy(w.scales + i * 4u, &sc, sizeof(sc));
}
#pragma unroll
for (uint32_t i = 0; i < 16; i++) {
const uint32_t q = ((const uint32_t *)(down_base + dm_bytes + sc_bytes))[
pair * 32u + i * 2u + parity];
memcpy(w.qs + i * 4u, &q, sizeof(q));
}
acc += dev_dot_q2_K_q8_K_block(&w, xq + b);
} else {
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
}
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) total += acc;
}
if (lane == 0) out[row] = total;
}
__global__ static void moe_down_owned_slots_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
const uint32_t slot = blockIdx.y;
if (row >= out_dim || slot >= 6u) return;
uint32_t expert = 0;
if (!moe_owned_local_expert(selected[slot], expert_base,
expert_count, &expert)) {
return;
}
const cuda_block_q2_K *wr =
(const cuda_block_q2_K *)(down_base +
(uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc;
}
/* Map one of two packed operands for a three-slot reduction group. The only
* multi-slot operand is the peer-owned prefix (slots 0+1 within the group),
* which can be pre-added exactly because the reference reduction starts from
* +0. Every other peer slot remains a distinct operand in original order. */
__device__ __forceinline__ static int moe_owned_packed_component(
const int32_t *selected,
uint32_t group,
uint32_t component,
uint32_t expert_base,
uint32_t expert_count,
bool *prefix_pair) {
const uint32_t slot0 = group * 3u;
uint32_t mask = 0u;
#pragma unroll
for (uint32_t i = 0; i < 3u; i++) {
if (moe_owned_local_expert(selected[slot0 + i], expert_base,
expert_count, NULL)) {
mask |= 1u << i;
}
}
*prefix_pair = false;
if ((mask & 3u) == 3u) {
if (component == 0u) {
*prefix_pair = true;
return (int)slot0;
}
return (mask & 4u) != 0u ? (int)(slot0 + 2u) : -1;
}
uint32_t ordinal = 0u;
#pragma unroll
for (uint32_t i = 0; i < 3u; i++) {
if ((mask & (1u << i)) == 0u) continue;
if (ordinal++ == component) return (int)(slot0 + i);
}
return -1;
}
/* One thread owns an output column and loads all source slots before writing.
* This permits packed_out == slots for the non-peer-copy fallback. */
__global__ static void moe_down_owned_pack_f32_slots_kernel(
float *packed_out,
const float *slots,
const int32_t *selected,
uint32_t out_dim,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t col =
(uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
if (col >= out_dim) return;
float slotv[6];
#pragma unroll
for (uint32_t slot = 0; slot < 6u; slot++) {
slotv[slot] = slots[(uint64_t)slot * out_dim + col];
}
float packed[4];
#pragma unroll
for (uint32_t packed_slot = 0; packed_slot < 4u; packed_slot++) {
bool prefix_pair = false;
const int first_slot = moe_owned_packed_component(
selected, packed_slot / 2u, packed_slot & 1u,
expert_base, expert_count, &prefix_pair);
if (first_slot < 0) {
packed[packed_slot] = 0.0f;
} else if (prefix_pair) {
float value = __fadd_rn(0.0f, slotv[first_slot]);
packed[packed_slot] = __fadd_rn(value, slotv[first_slot + 1]);
} else {
packed[packed_slot] = slotv[first_slot];
}
}
#pragma unroll
for (uint32_t packed_slot = 0; packed_slot < 4u; packed_slot++) {
packed_out[(uint64_t)packed_slot * out_dim + col] = packed[packed_slot];
}
}
__global__ static void moe_down_owned_copy_f32_slots_kernel(
float *dst,
const float *src,
uint64_t count) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < count) dst[i] = src[i];
}
__global__ static void moe_down_owned_packed_qwarp32_kernel(
float *packed_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
const uint32_t packed_slot = blockIdx.y;
if (row >= out_dim || packed_slot >= 4u) return;
bool prefix_pair = false;
const int first_slot = moe_owned_packed_component(
selected, packed_slot / 2u, packed_slot & 1u,
expert_base, expert_count, &prefix_pair);
if (first_slot < 0) {
if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f;
return;
}
float packed = 0.0f;
const uint32_t n_slots = prefix_pair ? 2u : 1u;
#pragma unroll
for (uint32_t i = 0; i < 2u; i++) {
if (i >= n_slots) break;
const uint32_t slot = (uint32_t)first_slot + i;
uint32_t expert = 0;
if (!moe_owned_local_expert(selected[slot], expert_base,
expert_count, &expert)) {
continue;
}
const cuda_block_q2_K *wr =
(const cuda_block_q2_K *)(down_base +
(uint64_t)expert * down_expert_bytes +
(uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0u) {
packed = prefix_pair ? __fadd_rn(packed, acc) : acc;
}
}
if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed;
}
__global__ static void moe_down_q4K_sum6_qwarp32_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
if (row >= out_dim) return;
const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u);
float total = 0.0f;
#pragma unroll
for (uint32_t slot = 0; slot < 6u; slot++) {
int32_t expert_i = selected[slot];
if (expert_i < 0) continue;
const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
if (vec_ok) {
for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc);
} else {
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) total += acc;
}
if (lane == 0) out[row] = total;
}
__global__ static void moe_down_q4K_owned_slots_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
const uint32_t slot = blockIdx.y;
if (row >= out_dim || slot >= 6u) return;
uint32_t expert = 0;
if (!moe_owned_local_expert(selected[slot], expert_base,
expert_count, &expert)) {
return;
}
const cuda_block_q4_K *wr =
(const cuda_block_q4_K *)(down_base +
(uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes |
down_expert_bytes) & 15u) == 0u);
float acc = 0.0f;
if (vec_ok) {
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc);
}
} else {
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
}
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)slot * out_dim + row] = acc;
}
__global__ static void moe_down_q4K_owned_packed_qwarp32_kernel(
float *packed_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t expert_base,
uint32_t expert_count) {
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
const uint32_t packed_slot = blockIdx.y;
if (row >= out_dim || packed_slot >= 4u) return;
bool prefix_pair = false;
const int first_slot = moe_owned_packed_component(
selected, packed_slot / 2u, packed_slot & 1u,
expert_base, expert_count, &prefix_pair);
if (first_slot < 0) {
if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f;
return;
}
const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes |
down_expert_bytes) & 15u) == 0u);
float packed = 0.0f;
const uint32_t n_slots = prefix_pair ? 2u : 1u;
#pragma unroll
for (uint32_t i = 0; i < 2u; i++) {
if (i >= n_slots) break;
const uint32_t slot = (uint32_t)first_slot + i;
uint32_t expert = 0;
if (!moe_owned_local_expert(selected[slot], expert_base,
expert_count, &expert)) {
continue;
}
const cuda_block_q4_K *wr =
(const cuda_block_q4_K *)(down_base +
(uint64_t)expert * down_expert_bytes +
(uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
if (vec_ok) {
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc);
}
} else {
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
}
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0u) {
packed = prefix_pair ? __fadd_rn(packed, acc) : acc;
}
}
if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed;
}
__global__ static void moe_owned_slots_combine_fixed3_kernel(
float *out,
const float *home_slots,
const float *peer_slots,
const int32_t *selected,
uint32_t out_dim,
uint32_t expert_split) {
const uint32_t col =
(uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
const uint32_t row = blockIdx.y;
if (col >= out_dim) return;
out += (uint64_t)row * out_dim;
home_slots += (uint64_t)row * 6u * out_dim;
peer_slots += (uint64_t)row * 6u * out_dim;
selected += (uint64_t)row * 6u;
float slotv[6];
#pragma unroll
for (uint32_t slot = 0; slot < 6u; slot++) {
const int32_t expert = selected[slot];
if (expert < 0 ||
(uint32_t)expert >= 2u * expert_split) {
slotv[slot] = 0.0f;
} else {
const bool on_home = (uint32_t)expert < expert_split;
const float *src = on_home ? home_slots : peer_slots;
slotv[slot] = src[(uint64_t)slot * out_dim + col];
}
}
float home = __fadd_rn(0.0f, slotv[0]);
home = __fadd_rn(home, slotv[1]);
home = __fadd_rn(home, slotv[2]);
float peer = __fadd_rn(0.0f, slotv[3]);
peer = __fadd_rn(peer, slotv[4]);
peer = __fadd_rn(peer, slotv[5]);
out[col] = __fadd_rn(home, peer);
}
__device__ static float moe_owned_packed_combine_row(
const float *home_slots,
const float *peer_packed,
const int32_t *selected,
uint32_t row,
uint32_t out_dim,
uint32_t expert_split) {
float groups[2];
#pragma unroll
for (uint32_t group = 0; group < 2u; group++) {
const uint32_t slot0 = group * 3u;
uint32_t peer_mask = 0u;
uint32_t valid_mask = 0u;
#pragma unroll
for (uint32_t i = 0; i < 3u; i++) {
const int32_t expert = selected[slot0 + i];
if (expert >= 0 && (uint32_t)expert < 2u * expert_split) {
valid_mask |= 1u << i;
}
if (expert >= 0 && (uint32_t)expert >= expert_split &&
(uint32_t)expert < 2u * expert_split) {
peer_mask |= 1u << i;
}
}
const float *packed = peer_packed +
(uint64_t)group * 2u * out_dim + row;
float acc;
if ((peer_mask & 3u) == 3u) {
/* packed[0] is already (+0 + slot0) + slot1. */
acc = packed[0];
float slot2 = 0.0f;
if ((peer_mask & 4u) != 0u) {
slot2 = packed[out_dim];
} else if ((valid_mask & 4u) != 0u) {
slot2 = home_slots[(uint64_t)(slot0 + 2u) * out_dim + row];
}
acc = __fadd_rn(acc, slot2);
} else {
acc = 0.0f;
uint32_t peer_operand = 0u;
#pragma unroll
for (uint32_t i = 0; i < 3u; i++) {
float value;
if ((peer_mask & (1u << i)) != 0u) {
value = packed[(uint64_t)peer_operand * out_dim];
peer_operand++;
} else if ((valid_mask & (1u << i)) != 0u) {
value = home_slots[(uint64_t)(slot0 + i) * out_dim + row];
} else {
value = 0.0f;
}
acc = __fadd_rn(acc, value);
}
}
groups[group] = acc;
}
return __fadd_rn(groups[0], groups[1]);
}
__global__ static void moe_owned_packed_combine_fixed3_kernel(
float *out,
const float *home_slots,
const float *peer_packed,
const int32_t *selected,
uint32_t out_dim,
uint32_t expert_split) {
const uint32_t row =
(uint32_t)((uint64_t)blockIdx.x * blockDim.x + threadIdx.x);
if (row >= out_dim) return;
out[row] = moe_owned_packed_combine_row(
home_slots, peer_packed, selected, row, out_dim, expert_split);
}
__global__ static void moe_down_q4K_sum3_qwarp32_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
if (row >= out_dim) return;
const bool vec_ok = ((((uintptr_t)down_base | down_row_bytes | down_expert_bytes) & 15u) == 0u);
float total = 0.0f;
#pragma unroll
for (uint32_t slot = 0; slot < 3u; slot++) {
int32_t expert_i = selected[slot];
if (expert_i < 0) continue;
const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
float acc = 0.0f;
if (vec_ok) {
for (uint32_t b = lane; b < midq_blocks; b += 8u) dev_dot_q4_K_q8_K_block_vec(wr + b, xq + b, &acc);
} else {
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
}
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) total += acc;
}
if (lane == 0) out[row] = total;
}
__global__ static void moe_down_q4K_sum3_slotwarp_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim) {
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t slot = lane >> 3u;
const uint32_t qlane = lane & 7u;
const uint32_t row = blockIdx.x * 8u + warp;
if (row >= out_dim) return;
float acc = 0.0f;
if (slot < 3u && selected[slot] >= 0) {
int32_t expert_i = selected[slot];
const cuda_block_q4_K *wr =
(const cuda_block_q4_K *)(down_base +
(uint64_t)(uint32_t)expert_i * down_expert_bytes +
(uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks;
for (uint32_t b = qlane; b < midq_blocks; b += 8u) {
acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
}
acc = quarter_warp_sum_f32(acc, qlane);
}
const float s1 = __shfl_sync(0xffffffffu, acc, 8);
const float s2 = __shfl_sync(0xffffffffu, acc, 16);
if (lane == 0u) {
const float s0 = acc;
out[row] = (s0 + s1) + s2;
}
}
static void routed_moe_decode_graph_destroy_one(int logical_tier) {
if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return;
cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier];
if (c->exec) (void)cudaGraphExecDestroy(c->exec);
if (c->graph) (void)cudaGraphDestroy(c->graph);
memset(c, 0, sizeof(*c));
}
static int routed_moe_decode_q4_graph_launch(
int logical_tier,
float *out,
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_w,
const char *up_w,
const char *down_w,
cuda_block_q8_K *xq,
cuda_block_q8_K *midq,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp,
const float *x) {
if (logical_tier < 0 || logical_tier >= DS4_MAX_GPUS) return 0;
if (n_expert != 3u && n_expert != 6u) return 0;
uint32_t xq_blocks = expert_in_dim / CUDA_QK_K;
uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K;
if (xq_blocks == 0u || midq_blocks == 0u) return 0;
cuda_moe_decode_graph_cache *c = &g_moe_decode_graph[logical_tier];
const bool shape_match =
c->valid &&
c->n_expert == n_expert &&
c->expert_in_dim == expert_in_dim &&
c->expert_mid_dim == expert_mid_dim &&
c->out_dim == out_dim;
if (c->valid && !shape_match) {
routed_moe_decode_graph_destroy_one(logical_tier);
c = &g_moe_decode_graph[logical_tier];
}
uint32_t x_rows = 1u;
uint32_t mid_rows = n_expert;
dim3 xq_grid(xq_blocks, 1, 1);
dim3 gate_grid((expert_mid_dim + 7u) / 8u, n_expert, 1);
dim3 midq_grid(midq_blocks, n_expert, 1);
dim3 down_grid((out_dim + 31u) / 32u, 1, 1);
dim3 block(256, 1, 1);
void *xq_args[] = { &xq, &x, &expert_in_dim, &x_rows };
cudaKernelNodeParams xq_params;
memset(&xq_params, 0, sizeof(xq_params));
xq_params.func = (void *)q8_K_quantize_kernel;
xq_params.gridDim = xq_grid;
xq_params.blockDim = block;
xq_params.kernelParams = xq_args;
void *gate_args[] = {
&gate_out, &up_out, &mid_out, &gate_w, &up_w, &xq, &selected,
&weights, &gate_expert_bytes, &gate_row_bytes, &xq_blocks,
&expert_mid_dim, &n_expert, &write_aux, &clamp
};
cudaKernelNodeParams gate_params;
memset(&gate_params, 0, sizeof(gate_params));
gate_params.func = (void *)moe_gate_up_mid_decode_q4K_warp32_kernel;
gate_params.gridDim = gate_grid;
gate_params.blockDim = block;
gate_params.kernelParams = gate_args;
void *midq_args[] = { &midq, &mid_out, &expert_mid_dim, &mid_rows };
cudaKernelNodeParams midq_params;
memset(&midq_params, 0, sizeof(midq_params));
midq_params.func = (void *)q8_K_quantize_kernel;
midq_params.gridDim = midq_grid;
midq_params.blockDim = block;
midq_params.kernelParams = midq_args;
void *down_args[] = {
&out, &down_w, &midq, &selected, &down_expert_bytes,
&down_row_bytes, &midq_blocks, &out_dim
};
cudaKernelNodeParams down_params;
memset(&down_params, 0, sizeof(down_params));
down_params.func = n_expert == 6u
? (void *)moe_down_q4K_sum6_qwarp32_kernel
: (void *)moe_down_q4K_sum3_qwarp32_kernel;
down_params.gridDim = down_grid;
down_params.blockDim = block;
down_params.kernelParams = down_args;
if (!c->valid) {
cudaError_t err = cudaGraphCreate(&c->graph, 0);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: routed MoE decode graph create failed: %s\n",
cudaGetErrorString(err));
routed_moe_decode_graph_destroy_one(logical_tier);
return -1;
}
err = cudaGraphAddKernelNode(&c->xq_node, c->graph, NULL, 0,
&xq_params);
if (err == cudaSuccess) {
err = cudaGraphAddKernelNode(&c->gate_node, c->graph,
&c->xq_node, 1, &gate_params);
}
if (err == cudaSuccess) {
err = cudaGraphAddKernelNode(&c->midq_node, c->graph,
&c->gate_node, 1, &midq_params);
}
if (err == cudaSuccess) {
err = cudaGraphAddKernelNode(&c->down_node, c->graph,
&c->midq_node, 1, &down_params);
}
if (err == cudaSuccess) {
err = cudaGraphInstantiate(&c->exec, c->graph, NULL, NULL, 0);
}
if (err != cudaSuccess) {
fprintf(stderr, "ds4: routed MoE decode graph instantiate failed: %s\n",
cudaGetErrorString(err));
routed_moe_decode_graph_destroy_one(logical_tier);
return -1;
}
c->n_expert = n_expert;
c->expert_in_dim = expert_in_dim;
c->expert_mid_dim = expert_mid_dim;
c->out_dim = out_dim;
c->valid = 1;
} else {
cudaError_t err =
cudaGraphExecKernelNodeSetParams(c->exec, c->xq_node,
&xq_params);
if (err == cudaSuccess) {
err = cudaGraphExecKernelNodeSetParams(c->exec, c->gate_node,
&gate_params);
}
if (err == cudaSuccess) {
err = cudaGraphExecKernelNodeSetParams(c->exec, c->midq_node,
&midq_params);
}
if (err == cudaSuccess) {
err = cudaGraphExecKernelNodeSetParams(c->exec, c->down_node,
&down_params);
}
if (err != cudaSuccess) {
fprintf(stderr, "ds4: routed MoE decode graph update failed: %s\n",
cudaGetErrorString(err));
routed_moe_decode_graph_destroy_one(logical_tier);
return -1;
}
}
cudaError_t err = cudaGraphLaunch(c->exec, 0);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: routed MoE decode graph launch failed: %s\n",
cudaGetErrorString(err));
routed_moe_decode_graph_destroy_one(logical_tier);
return -1;
}
return 1;
}
/* Q4_K prefill (n_tokens > 1) down kernel. Mirrors moe_down_qwarp32_kernel
* geometry exactly; only the weight block type and dot helper differ. The
* pair = blockIdx.y indexing means the same grid shape (out_dim/32, n_tokens*n_expert)
* used by the IQ2 path applies here. The downstream moe_sum_kernel is
* weight-type-agnostic and sums these per-pair outputs into the final output. */
__global__ static void moe_down_q4K_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q4_K_q8_K_block(wr + b, xq + b);
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
template <uint32_t ROW_SPAN>
__global__ static void moe_down_q4K_expert_tile8_rowspan_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][8];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) {
uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u;
if (row >= out_dim) continue;
const cuda_block_q4_K *wr = (const cuda_block_q4_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q4_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc);
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
/* INT8 tensor-core (m8n8k16) exact MoE prefill tile kernels.
*
* Each warp computes an 8-token x 8-row tile. The Q4_K x Q8_K superblock dot
* keeps its integer sums (order-invariant, exact) but computes the 32-wide
* group dots on tensor cores; every output element keeps 8 float slot
* accumulators (slot[b & 7] += term_b, b ascending) and reduces them with the
* exact quarter_warp_sum_f32 grouping, so results are bit-identical to the
* scalar expert-tile kernels (fuzz-verified). Requires sm_75+, 16B-aligned
* expert tensors, and the staged activation-block counts (<=16 gate/up,
* <=8 down). Rollback: DS4_CUDA_MOE_NO_Q4_MMA=1. */
__device__ __forceinline__ static void mma_m8n8k16_s8(int32_t &c0, int32_t &c1, uint32_t a, uint32_t b) {
#if __CUDA_ARCH__ >= 750
asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%0,%1};"
: "+r"(c0), "+r"(c1) : "r"(a), "r"(b));
#else
(void)a; (void)b; (void)c0; (void)c1;
#endif
}
template <uint32_t ROW_SPAN>
__global__ static void moe_gate_up_mid_q4K_tile8_mma_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][16];
__shared__ uint32_t s_pair[8];
__shared__ uint32_t s_tok[8];
__shared__ uint32_t s_slot[8];
__shared__ uint32_t s_np;
if (threadIdx.x == 0) {
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
uint32_t pr = sorted_pairs[offsets[expert] + local_pair];
s_pair[np] = pr;
s_tok[np] = pr / n_expert;
s_slot[np] = pr - s_tok[np] * n_expert;
}
s_np = np;
}
__syncthreads();
const uint32_t np = s_np;
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u); i += blockDim.x) {
const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
uint32_t p = i / words_per_tok;
uint32_t w = i - p * words_per_tok;
((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w];
}
__syncthreads();
}
const uint32_t mtok = lane >> 2u; /* token row of this thread's C elems */
const uint32_t n0 = (lane & 3u) * 2u; /* first C column (weight row) */
/* 8 warps x 8 rows = 64 rows per pass */
for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) {
const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u;
if (row0 >= expert_mid_dim) continue;
const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes;
const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes;
/* per-element slot accumulators (2 elements x 8 slots) */
float sg0[8] = {0,0,0,0,0,0,0,0}, sg1[8] = {0,0,0,0,0,0,0,0};
float su0[8] = {0,0,0,0,0,0,0,0}, su1[8] = {0,0,0,0,0,0,0,0};
for (uint32_t b = 0; b < xq_blocks; b++) {
/* headers for this thread's two C columns */
const uint4 ghdr0 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)n0 * gate_row_bytes) + b);
const uint4 ghdr1 = *(const uint4 *)((const cuda_block_q4_K *)(grow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b);
const uint4 uhdr0 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)n0 * gate_row_bytes) + b);
const uint4 uhdr1 = *(const uint4 *)((const cuda_block_q4_K *)(urow + (uint64_t)(n0 + 1u) * gate_row_bytes) + b);
/* B-fragment source rows for loads: n_load = lane>>2.
* Batch all global loads for this superblock upfront so the
* memory system sees independent requests instead of a
* load->mma dependency chain. */
const uint32_t *gqw = (const uint32_t *)(((const cuda_block_q4_K *)(grow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs);
const uint32_t *uqw = (const uint32_t *)(((const cuda_block_q4_K *)(urow + (uint64_t)(lane >> 2u) * gate_row_bytes) + b)->qs);
const int8_t *aqs = sxq[mtok][b].qs;
uint32_t gw8[8], uw8[8];
#pragma unroll
for (uint32_t k = 0; k < 8u; k++) {
gw8[k] = gqw[k * 4u + (lane & 3u)];
uw8[k] = uqw[k * 4u + (lane & 3u)];
}
int gi0 = 0, gi1 = 0, ui0 = 0, ui1 = 0;
int gs0 = 0, gs1 = 0, us0 = 0, us1 = 0;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const int shift = (j & 1u) ? 4 : 0;
/* dot32 via two chained k16 mmas, per matrix */
int32_t gc0 = 0, gc1 = 0, uc0 = 0, uc1 = 0;
#pragma unroll
for (uint32_t h = 0; h < 2u; h++) {
const uint32_t koff = h * 16u + (lane & 3u) * 4u;
const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff);
const uint32_t gw = (gw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu;
const uint32_t uw = (uw8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu;
mma_m8n8k16_s8(gc0, gc1, a, gw);
mma_m8n8k16_s8(uc0, uc1, a, uw);
}
/* integer scale application for this thread's two columns */
uint8_t sc, m;
dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr0.y, &sc, &m);
gi0 += (int)sc * gc0;
const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u];
gs0 += (int)m * bs;
dev_q4_K_get_scale_min(j, (const uint8_t *)&ghdr1.y, &sc, &m);
gi1 += (int)sc * gc1;
gs1 += (int)m * bs;
dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr0.y, &sc, &m);
ui0 += (int)sc * uc0;
us0 += (int)m * bs;
dev_q4_K_get_scale_min(j, (const uint8_t *)&uhdr1.y, &sc, &m);
ui1 += (int)sc * uc1;
us1 += (int)m * bs;
}
/* float finish, exact dev_dot_q4_K_q8_K_block8 expression */
const float yd = sxq[mtok][b].d;
const uint32_t sl = b & 7u;
sg0[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr0.x & 0xffffu)) * (float)gi0 -
yd * dev_f16_to_f32((uint16_t)(ghdr0.x >> 16u)) * (float)gs0;
sg1[sl] += yd * dev_f16_to_f32((uint16_t)(ghdr1.x & 0xffffu)) * (float)gi1 -
yd * dev_f16_to_f32((uint16_t)(ghdr1.x >> 16u)) * (float)gs1;
su0[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr0.x & 0xffffu)) * (float)ui0 -
yd * dev_f16_to_f32((uint16_t)(uhdr0.x >> 16u)) * (float)us0;
su1[sl] += yd * dev_f16_to_f32((uint16_t)(uhdr1.x & 0xffffu)) * (float)ui1 -
yd * dev_f16_to_f32((uint16_t)(uhdr1.x >> 16u)) * (float)us1;
}
/* quarter_warp_sum_f32 order: ((s0+s4)+(s2+s6)) + ((s1+s5)+(s3+s7)) */
const uint32_t p = mtok;
if (p < np) {
const uint32_t rowa = row0 + n0;
const uint32_t rowb = row0 + n0 + 1u;
float gate2[2], up2[2];
{
float a0 = sg0[0] + sg0[4], a1 = sg0[1] + sg0[5], a2 = sg0[2] + sg0[6], a3 = sg0[3] + sg0[7];
gate2[0] = (a0 + a2) + (a1 + a3);
a0 = sg1[0] + sg1[4]; a1 = sg1[1] + sg1[5]; a2 = sg1[2] + sg1[6]; a3 = sg1[3] + sg1[7];
gate2[1] = (a0 + a2) + (a1 + a3);
a0 = su0[0] + su0[4]; a1 = su0[1] + su0[5]; a2 = su0[2] + su0[6]; a3 = su0[3] + su0[7];
up2[0] = (a0 + a2) + (a1 + a3);
a0 = su1[0] + su1[4]; a1 = su1[1] + su1[5]; a2 = su1[2] + su1[6]; a3 = su1[3] + su1[7];
up2[1] = (a0 + a2) + (a1 + a3);
}
#pragma unroll
for (uint32_t e = 0; e < 2u; e++) {
const uint32_t row = e ? rowb : rowa;
if (row >= expert_mid_dim) continue;
float gate = gate2[e];
float up = up2[e];
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]];
}
}
}
}
template <uint32_t ROW_SPAN>
__global__ static void moe_down_q4K_tile8_mma_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][8];
__shared__ uint32_t s_pair[8];
__shared__ uint32_t s_np;
if (threadIdx.x == 0) {
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
s_pair[np] = sorted_pairs[offsets[expert] + local_pair];
}
s_np = np;
}
__syncthreads();
const uint32_t np = s_np;
if (midq_blocks <= 8u) {
const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) {
uint32_t p = i / words_per_tok;
uint32_t w = i - p * words_per_tok;
((uint32_t *)sxq[p])[w] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w];
}
__syncthreads();
}
const uint32_t mtok = lane >> 2u;
const uint32_t n0 = (lane & 3u) * 2u;
for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) {
const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u;
if (row0 >= out_dim) continue;
const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes;
float s0[8] = {0,0,0,0,0,0,0,0}, s1[8] = {0,0,0,0,0,0,0,0};
for (uint32_t b = 0; b < midq_blocks; b++) {
const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b);
const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b);
const uint32_t *wqw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs);
const int8_t *aqs = sxq[mtok][b].qs;
uint32_t w8[8];
#pragma unroll
for (uint32_t k = 0; k < 8u; k++) w8[k] = wqw[k * 4u + (lane & 3u)];
int i0 = 0, i1 = 0, m0 = 0, m1 = 0;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const int shift = (j & 1u) ? 4 : 0;
int32_t c0 = 0, c1 = 0;
#pragma unroll
for (uint32_t h = 0; h < 2u; h++) {
const uint32_t koff = h * 16u + (lane & 3u) * 4u;
const uint32_t a = *(const uint32_t *)(aqs + j * 32u + koff);
const uint32_t w = (w8[(j >> 1u) * 2u + h] >> shift) & 0x0f0f0f0fu;
mma_m8n8k16_s8(c0, c1, a, w);
}
uint8_t sc, m;
const int bs = (int)sxq[mtok][b].bsums[2u * j] + (int)sxq[mtok][b].bsums[2u * j + 1u];
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc, &m);
i0 += (int)sc * c0;
m0 += (int)m * bs;
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc, &m);
i1 += (int)sc * c1;
m1 += (int)m * bs;
}
const float yd = sxq[mtok][b].d;
const uint32_t sl = b & 7u;
s0[sl] += yd * dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu)) * (float)i0 -
yd * dev_f16_to_f32((uint16_t)(hdr0.x >> 16u)) * (float)m0;
s1[sl] += yd * dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu)) * (float)i1 -
yd * dev_f16_to_f32((uint16_t)(hdr1.x >> 16u)) * (float)m1;
}
const uint32_t p = mtok;
if (p < np) {
float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7];
const float r0 = (a0 + a2) + (a1 + a3);
a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7];
const float r1 = (a0 + a2) + (a1 + a3);
if (row0 + n0 < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0] = r0;
if (row0 + n0 + 1u < out_dim) down_out[(uint64_t)s_pair[p] * out_dim + row0 + n0 + 1u] = r1;
}
}
}
/* 16-pair MoE expert tile kernels on sm_80+ m16n8k32 INT8 tensor cores.
*
* Same per-output math and reduction order as the 8-pair expert tile
* kernels (slot[b & 7] += term_b with b ascending, then the exact
* quarter_warp_sum_f32 grouping), so results are bit-identical; grouping
* 16 pairs per tile just halves how often each expert's weights are
* streamed from DRAM. Gate and up run as two passes over the superblocks
* to keep register pressure at the 8-pair kernel's level. */
__device__ __forceinline__ static void mma16_m16n8k32_s8(
int32_t &c0, int32_t &c1, int32_t &c2, int32_t &c3,
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
uint32_t b0, uint32_t b1) {
#if __CUDA_ARCH__ >= 800
asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
: "+r"(c0),"+r"(c1),"+r"(c2),"+r"(c3)
: "r"(a0),"r"(a1),"r"(a2),"r"(a3),"r"(b0),"r"(b1));
#else
(void)a0;(void)a1;(void)a2;(void)a3;(void)b0;(void)b1;(void)c0;(void)c1;(void)c2;(void)c3;
#endif
}
/* One matrix pass over all superblocks for this thread's 4 C elements
* (tokens mtokA/mtokB x rows n0/n0+1). Returns the quarter-tree-reduced
* values in r[4] with the exact reference ordering. */
__device__ __forceinline__ static void moe_tile16_mma_pass(
const char *wrow, /* row0 base of this matrix */
uint64_t row_bytes,
const cuda_block_q8_K (*sxq)[16],
uint32_t xq_blocks,
uint32_t lane,
float r[4]) {
const uint32_t mtokA = lane >> 2u;
const uint32_t mtokB = mtokA + 8u;
const uint32_t n0 = (lane & 3u) * 2u;
float s0[8] = {0,0,0,0,0,0,0,0};
float s1[8] = {0,0,0,0,0,0,0,0};
float s2[8] = {0,0,0,0,0,0,0,0};
float s3[8] = {0,0,0,0,0,0,0,0};
for (uint32_t b = 0; b < xq_blocks; b++) {
const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * row_bytes) + b);
const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * row_bytes) + b);
const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * row_bytes) + b)->qs);
uint32_t w8[8];
#pragma unroll
for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)];
const int8_t *aqsA = sxq[mtokA][b].qs;
const int8_t *aqsB = sxq[mtokB][b].qs;
int i0 = 0, i1 = 0, i2 = 0, i3 = 0;
int m0 = 0, m1 = 0, m2 = 0, m3 = 0;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const int shift = (j & 1u) ? 4 : 0;
const uint32_t koff = (lane & 3u) * 4u;
const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff);
const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff);
const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff);
const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff);
const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu;
const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu;
int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0;
mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1);
uint8_t sc0, sm0, sc1, sm1;
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0);
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1);
const int bsA = (int)sxq[mtokA][b].bsums[2u * j] + (int)sxq[mtokA][b].bsums[2u * j + 1u];
const int bsB = (int)sxq[mtokB][b].bsums[2u * j] + (int)sxq[mtokB][b].bsums[2u * j + 1u];
i0 += (int)sc0 * c0;
i1 += (int)sc1 * c1;
i2 += (int)sc0 * c2;
i3 += (int)sc1 * c3;
m0 += (int)sm0 * bsA;
m1 += (int)sm1 * bsA;
m2 += (int)sm0 * bsB;
m3 += (int)sm1 * bsB;
}
const float ydA = sxq[mtokA][b].d;
const float ydB = sxq[mtokB][b].d;
const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu));
const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u));
const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu));
const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u));
const uint32_t sl = b & 7u;
s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0;
s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1;
s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2;
s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3;
}
{
float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7];
r[0] = (a0 + a2) + (a1 + a3);
a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7];
r[1] = (a0 + a2) + (a1 + a3);
a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7];
r[2] = (a0 + a2) + (a1 + a3);
a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7];
r[3] = (a0 + a2) + (a1 + a3);
}
}
template <uint32_t ROW_SPAN>
__global__ static void moe_gate_up_mid_q4K_tile16_mma_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t write_aux,
float clamp) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
extern __shared__ unsigned char t16_sh[];
cuda_block_q8_K (*sxq)[16] = (cuda_block_q8_K (*)[16])t16_sh; /* [16][16] */
__shared__ uint32_t s_pair[16];
__shared__ uint32_t s_tok[16];
__shared__ uint32_t s_slot[16];
__shared__ uint32_t s_np;
if (threadIdx.x == 0) {
uint32_t np = 0;
for (; np < 16u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
uint32_t pr = sorted_pairs[offsets[expert] + local_pair];
s_pair[np] = pr;
s_tok[np] = pr / n_expert;
s_slot[np] = pr - s_tok[np] * n_expert;
}
s_np = np;
}
__syncthreads();
const uint32_t np = s_np;
if (xq_blocks <= 16u) {
const uint32_t words_per_tok = xq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) {
uint32_t p = i / words_per_tok;
uint32_t w = i - p * words_per_tok;
((uint32_t *)sxq[p])[w] = ((const uint32_t *)(xq + (uint64_t)s_tok[p] * xq_blocks))[w];
}
/* zero-fill missing pairs so the A fragments are defined */
const uint32_t total_words = 16u * words_per_tok;
for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) {
((uint32_t *)t16_sh)[i] = 0u;
}
__syncthreads();
}
const uint32_t mtokA = lane >> 2u;
const uint32_t mtokB = mtokA + 8u;
const uint32_t n0 = (lane & 3u) * 2u;
for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) {
const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u;
if (row0 >= expert_mid_dim) continue;
const char *grow = gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes;
const char *urow = up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row0 * gate_row_bytes;
float gr[4], ur[4];
moe_tile16_mma_pass(grow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, gr);
moe_tile16_mma_pass(urow, gate_row_bytes, (const cuda_block_q8_K (*)[16])sxq, xq_blocks, lane, ur);
#pragma unroll
for (uint32_t e = 0; e < 4u; e++) {
const uint32_t p = (e < 2u) ? mtokA : mtokB;
const uint32_t row = row0 + n0 + (e & 1u);
if (p >= np || row >= expert_mid_dim) continue;
float gate = gr[e];
float up = ur[e];
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)s_pair[p] * expert_mid_dim + row;
if (write_aux) {
gate_out[off] = gate;
up_out[off] = up;
}
mid_out[off] = (gate / (1.0f + expf(-gate))) * up *
weights[(uint64_t)s_tok[p] * n_expert + s_slot[p]];
}
}
}
template <uint32_t ROW_SPAN>
__global__ static void moe_down_q4K_tile16_mma_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t warp = threadIdx.x >> 5u;
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
extern __shared__ unsigned char t16_sh[];
__shared__ uint32_t s_pair[16];
__shared__ uint32_t s_np;
if (threadIdx.x == 0) {
uint32_t np = 0;
for (; np < 16u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
s_pair[np] = sorted_pairs[offsets[expert] + local_pair];
}
s_np = np;
}
__syncthreads();
const uint32_t np = s_np;
if (midq_blocks <= 16u) {
const uint32_t words_per_tok = midq_blocks * (uint32_t)(sizeof(cuda_block_q8_K) / 4u);
for (uint32_t i = threadIdx.x; i < np * words_per_tok; i += blockDim.x) {
uint32_t p = i / words_per_tok;
uint32_t w = i - p * words_per_tok;
((uint32_t *)t16_sh)[i] = ((const uint32_t *)(midq + (uint64_t)s_pair[p] * midq_blocks))[w];
}
const uint32_t total_words = 16u * words_per_tok;
for (uint32_t i = threadIdx.x + np * words_per_tok; i < total_words; i += blockDim.x) {
((uint32_t *)t16_sh)[i] = 0u;
}
__syncthreads();
}
const uint32_t mtokA = lane >> 2u;
const uint32_t mtokB = mtokA + 8u;
const uint32_t n0 = (lane & 3u) * 2u;
for (uint32_t rr = 0; rr < ROW_SPAN / 64u; rr++) {
const uint32_t row0 = blockIdx.x * ROW_SPAN + rr * 64u + warp * 8u;
if (row0 >= out_dim) continue;
const char *wrow = down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row0 * down_row_bytes;
float s0[8] = {0,0,0,0,0,0,0,0};
float s1[8] = {0,0,0,0,0,0,0,0};
float s2[8] = {0,0,0,0,0,0,0,0};
float s3[8] = {0,0,0,0,0,0,0,0};
for (uint32_t b = 0; b < midq_blocks; b++) {
const uint4 hdr0 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)n0 * down_row_bytes) + b);
const uint4 hdr1 = *(const uint4 *)((const cuda_block_q4_K *)(wrow + (uint64_t)(n0 + 1u) * down_row_bytes) + b);
const uint32_t *qw = (const uint32_t *)(((const cuda_block_q4_K *)(wrow + (uint64_t)(lane >> 2u) * down_row_bytes) + b)->qs);
uint32_t w8[8];
#pragma unroll
for (uint32_t k = 0; k < 8u; k++) w8[k] = qw[k * 4u + (lane & 3u)];
/* activation rows: midq_blocks stride within the staged region */
const int8_t *aqsA = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b)->qs;
const int8_t *aqsB = ((const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b)->qs;
const cuda_block_q8_K *blkA = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokA * midq_blocks + b;
const cuda_block_q8_K *blkB = (const cuda_block_q8_K *)t16_sh + (uint64_t)mtokB * midq_blocks + b;
int i0 = 0, i1 = 0, i2 = 0, i3 = 0;
int m0 = 0, m1 = 0, m2 = 0, m3 = 0;
#pragma unroll
for (uint32_t j = 0; j < 8u; j++) {
const int shift = (j & 1u) ? 4 : 0;
const uint32_t koff = (lane & 3u) * 4u;
const uint32_t a0 = *(const uint32_t *)(aqsA + j * 32u + koff);
const uint32_t a1 = *(const uint32_t *)(aqsB + j * 32u + koff);
const uint32_t a2 = *(const uint32_t *)(aqsA + j * 32u + 16u + koff);
const uint32_t a3 = *(const uint32_t *)(aqsB + j * 32u + 16u + koff);
const uint32_t b0 = (w8[(j >> 1u) * 2u + 0u] >> shift) & 0x0f0f0f0fu;
const uint32_t b1 = (w8[(j >> 1u) * 2u + 1u] >> shift) & 0x0f0f0f0fu;
int32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0;
mma16_m16n8k32_s8(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1);
uint8_t sc0, sm0, sc1, sm1;
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr0.y, &sc0, &sm0);
dev_q4_K_get_scale_min(j, (const uint8_t *)&hdr1.y, &sc1, &sm1);
const int bsA = (int)blkA->bsums[2u * j] + (int)blkA->bsums[2u * j + 1u];
const int bsB = (int)blkB->bsums[2u * j] + (int)blkB->bsums[2u * j + 1u];
i0 += (int)sc0 * c0;
i1 += (int)sc1 * c1;
i2 += (int)sc0 * c2;
i3 += (int)sc1 * c3;
m0 += (int)sm0 * bsA;
m1 += (int)sm1 * bsA;
m2 += (int)sm0 * bsB;
m3 += (int)sm1 * bsB;
}
const float ydA = blkA->d;
const float ydB = blkB->d;
const float xd0 = dev_f16_to_f32((uint16_t)(hdr0.x & 0xffffu));
const float xmin0 = dev_f16_to_f32((uint16_t)(hdr0.x >> 16u));
const float xd1 = dev_f16_to_f32((uint16_t)(hdr1.x & 0xffffu));
const float xmin1 = dev_f16_to_f32((uint16_t)(hdr1.x >> 16u));
const uint32_t sl = b & 7u;
s0[sl] += ydA * xd0 * (float)i0 - ydA * xmin0 * (float)m0;
s1[sl] += ydA * xd1 * (float)i1 - ydA * xmin1 * (float)m1;
s2[sl] += ydB * xd0 * (float)i2 - ydB * xmin0 * (float)m2;
s3[sl] += ydB * xd1 * (float)i3 - ydB * xmin1 * (float)m3;
}
float rr4[4];
{
float a0 = s0[0] + s0[4], a1 = s0[1] + s0[5], a2 = s0[2] + s0[6], a3 = s0[3] + s0[7];
rr4[0] = (a0 + a2) + (a1 + a3);
a0 = s1[0] + s1[4]; a1 = s1[1] + s1[5]; a2 = s1[2] + s1[6]; a3 = s1[3] + s1[7];
rr4[1] = (a0 + a2) + (a1 + a3);
a0 = s2[0] + s2[4]; a1 = s2[1] + s2[5]; a2 = s2[2] + s2[6]; a3 = s2[3] + s2[7];
rr4[2] = (a0 + a2) + (a1 + a3);
a0 = s3[0] + s3[4]; a1 = s3[1] + s3[5]; a2 = s3[2] + s3[6]; a3 = s3[3] + s3[7];
rr4[3] = (a0 + a2) + (a1 + a3);
}
#pragma unroll
for (uint32_t e = 0; e < 4u; e++) {
const uint32_t p = (e < 2u) ? mtokA : mtokB;
const uint32_t row = row0 + n0 + (e & 1u);
if (p >= np || row >= out_dim) continue;
down_out[(uint64_t)s_pair[p] * out_dim + row] = rr4[e];
}
}
}
static int cuda_q4_mma_tile16_shmem_ok(int which_down) {
/* Opt the tile16 kernels into >48KB dynamic shared memory, per device. */
static int ready[DS4_MAX_GPUS][2];
static int failed = 0;
if (failed) return 0;
int dev = 0;
cudaGetDevice(&dev);
if (dev < 0 || dev >= DS4_MAX_GPUS) return 0;
if (ready[dev][which_down]) return 1;
cudaFuncAttributes fn_attr;
cudaError_t err = which_down
? cudaFuncGetAttributes(&fn_attr, moe_down_q4K_tile16_mma_kernel<512>)
: cudaFuncGetAttributes(&fn_attr, moe_gate_up_mid_q4K_tile16_mma_kernel<512>);
if (err != cudaSuccess || fn_attr.binaryVersion < 80 ||
fn_attr.ptxVersion < 80) {
failed = 1;
return 0;
}
const int bytes = (int)(16u * 16u * sizeof(cuda_block_q8_K));
if (which_down) {
err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<512>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
if (err == cudaSuccess)
err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<1024>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
if (err == cudaSuccess)
err = cudaFuncSetAttribute(moe_down_q4K_tile16_mma_kernel<2048>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
} else {
err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<512>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
if (err == cudaSuccess)
err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<1024>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
if (err == cudaSuccess)
err = cudaFuncSetAttribute(moe_gate_up_mid_q4K_tile16_mma_kernel<2048>,
cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}
if (err != cudaSuccess) {
failed = 1;
return 0;
}
ready[dev][which_down] = 1;
return 1;
}
__global__ static void moe_down_sorted_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t pair = sorted_pairs[blockIdx.y];
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
__global__ static DS4_CUDA_UNUSED void moe_down_expert_tile8_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t group = threadIdx.x >> 3u;
uint32_t lane = threadIdx.x & 7u;
uint32_t pair_slot = group & 7u;
uint32_t row_lane = group >> 3u;
uint32_t expert = tile_experts[tile];
uint32_t local_pair = tile_starts[tile] + pair_slot;
if (local_pair >= counts[expert]) return;
uint32_t sorted_idx = offsets[expert] + local_pair;
uint32_t pair = sorted_pairs[sorted_idx];
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
for (uint32_t rr = 0; rr < 2u; rr++) {
uint32_t row = blockIdx.x * 8u + row_lane + rr * 4u;
if (row >= out_dim) continue;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
}
__global__ static void moe_down_expert_tile4_row32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t atomic_out) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[4][8];
uint32_t pair[4] = {0, 0, 0, 0};
const cuda_block_q8_K *xqb[4] = {NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 4u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
if (row >= out_dim) return;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q2_K_q8_K_block4(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL, np, acc);
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) {
if (atomic_out) {
uint32_t tok = pair[p] / n_expert;
atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]);
} else {
down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
}
__global__ static void moe_down_expert_tile8_row32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t atomic_out) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t expert = tile_experts[tile];
uint32_t local_start = tile_starts[tile];
__shared__ cuda_block_q8_K sxq[8][8];
uint32_t pair[8] = {0, 0, 0, 0, 0, 0, 0, 0};
const cuda_block_q8_K *xqb[8] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL};
uint32_t np = 0;
for (; np < 8u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
if (row >= out_dim) return;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np, acc);
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) {
if (atomic_out) {
uint32_t tok = pair[p] / n_expert;
atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]);
} else {
down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
}
__global__ static void moe_down_expert_tile16_row32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t atomic_out) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t local_start = tile_starts[tile];
if (local_start & 8u) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
uint32_t expert = tile_experts[tile];
__shared__ cuda_block_q8_K sxq[16][8];
uint32_t pair[16] = {0};
const cuda_block_q8_K *xqb[16] = {NULL};
uint32_t np = 0;
for (; np < 16u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
if (row >= out_dim) return;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[16] = {0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc);
if (np > 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL,
xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL,
xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL,
xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8);
}
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) {
if (atomic_out) {
uint32_t tok = pair[p] / n_expert;
atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]);
} else {
down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
}
__global__ static void moe_down_expert_tile16_row2048_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t atomic_out) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t local_start = tile_starts[tile];
if (local_start & 8u) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
__shared__ cuda_block_q8_K sxq[16][8];
uint32_t pair[16] = {0};
const cuda_block_q8_K *xqb[16] = {NULL};
uint32_t np = 0;
for (; np < 16u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < 64u; rr++) {
uint32_t row = blockIdx.x * 2048u + row_lane + rr * 32u;
if (row >= out_dim) continue;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[16] = {0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc);
if (np > 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL,
xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL,
xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL,
xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8);
}
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) {
if (atomic_out) {
uint32_t tok = pair[p] / n_expert;
atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]);
} else {
down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
}
}
template <uint32_t ROW_SPAN>
__global__ static void moe_down_expert_tile16_rowspan_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const uint32_t *offsets,
const uint32_t *counts,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t atomic_out) {
uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
uint32_t local_start = tile_starts[tile];
if (local_start & 8u) return;
uint32_t lane = threadIdx.x & 7u;
uint32_t row_lane = threadIdx.x >> 3u;
uint32_t expert = tile_experts[tile];
__shared__ cuda_block_q8_K sxq[16][8];
uint32_t pair[16] = {0};
const cuda_block_q8_K *xqb[16] = {NULL};
uint32_t np = 0;
for (; np < 16u; np++) {
uint32_t local_pair = local_start + np;
if (local_pair >= counts[expert]) break;
pair[np] = sorted_pairs[offsets[expert] + local_pair];
xqb[np] = midq + (uint64_t)pair[np] * midq_blocks;
}
if (midq_blocks <= 8u) {
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
uint32_t p = i / midq_blocks;
uint32_t b = i - p * midq_blocks;
sxq[p][b] = xqb[p][b];
}
__syncthreads();
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) {
uint32_t row = blockIdx.x * ROW_SPAN + row_lane + rr * 32u;
if (row >= out_dim) continue;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)expert * down_expert_bytes + (uint64_t)row * down_row_bytes);
float acc[16] = {0.0f};
for (uint32_t b = lane; b < midq_blocks; b += 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[0] ? xqb[0] + b : NULL, xqb[1] ? xqb[1] + b : NULL,
xqb[2] ? xqb[2] + b : NULL, xqb[3] ? xqb[3] + b : NULL,
xqb[4] ? xqb[4] + b : NULL, xqb[5] ? xqb[5] + b : NULL,
xqb[6] ? xqb[6] + b : NULL, xqb[7] ? xqb[7] + b : NULL, np < 8u ? np : 8u, acc);
if (np > 8u) {
dev_dot_q2_K_q8_K_block8(wr + b, xqb[8] ? xqb[8] + b : NULL, xqb[9] ? xqb[9] + b : NULL,
xqb[10] ? xqb[10] + b : NULL, xqb[11] ? xqb[11] + b : NULL,
xqb[12] ? xqb[12] + b : NULL, xqb[13] ? xqb[13] + b : NULL,
xqb[14] ? xqb[14] + b : NULL, xqb[15] ? xqb[15] + b : NULL, np - 8u, acc + 8);
}
}
for (uint32_t p = 0; p < np; p++) {
acc[p] = quarter_warp_sum_f32(acc[p], lane);
if (lane == 0) {
if (atomic_out) {
uint32_t tok = pair[p] / n_expert;
atomicAdd(down_out + (uint64_t)tok * out_dim + row, acc[p]);
} else {
down_out[(uint64_t)pair[p] * out_dim + row] = acc[p];
}
}
}
}
}
__global__ static void moe_down_sorted_p2_qwarp32_kernel(
float *down_out,
const char *down_base,
const cuda_block_q8_K *midq,
const uint32_t *sorted_pairs,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t pair_count) {
uint32_t lane = threadIdx.x & 7u;
uint32_t pair_lane = (threadIdx.x >> 3u) & 1u;
uint32_t row = blockIdx.x * 16u + (threadIdx.x >> 4u);
uint32_t sorted_idx = blockIdx.y * 2u + pair_lane;
if (row >= out_dim || sorted_idx >= pair_count) return;
uint32_t pair = sorted_pairs[sorted_idx];
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const cuda_block_q8_K *xq = midq + (uint64_t)pair * midq_blocks;
float acc = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 8u) acc += dev_dot_q2_K_q8_K_block(wr + b, xq + b);
acc = quarter_warp_sum_f32(acc, lane);
if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc;
}
__global__ static void moe_sum_kernel(float *out, const float *down, uint32_t out_dim, uint32_t n_expert, uint32_t n_tokens) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * out_dim;
if (gid >= n) return;
uint32_t tok = gid / out_dim;
uint32_t row = gid - (uint64_t)tok * out_dim;
float acc = 0.0f;
for (uint32_t e = 0; e < n_expert; e++) acc += down[((uint64_t)tok * n_expert + e) * out_dim + row];
out[gid] = acc;
}
__global__ static void moe_sum_owned_kernel(
float *out,
const float *down,
const int32_t *selected,
uint32_t out_dim,
uint32_t n_expert,
uint32_t n_tokens) {
const uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t n = (uint64_t)n_tokens * out_dim;
if (gid >= n) return;
const uint32_t tok = (uint32_t)(gid / out_dim);
const uint32_t row = (uint32_t)(gid - (uint64_t)tok * out_dim);
float acc = 0.0f;
#pragma unroll
for (uint32_t slot = 0; slot < 6u; slot++) {
if (slot >= n_expert) break;
const uint64_t pair = (uint64_t)tok * n_expert + slot;
const float value = selected[pair] >= 0
? down[pair * out_dim + row] : 0.0f;
acc = __fadd_rn(acc, value);
}
out[gid] = acc;
}
__device__ static float dev_iq2_xxs_dot_f32(const cuda_block_iq2_xxs *row, const float *x, uint32_t nb) {
float acc = 0.0f;
for (uint32_t b = 0; b < nb; b++) {
const cuda_block_iq2_xxs *xb = row + b;
const float d = dev_f16_to_f32(xb->d);
const uint16_t *q2 = xb->qs;
const float *xf = x + (uint64_t)b * CUDA_QK_K;
for (uint32_t ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) {
const uint32_t aux_g = (uint32_t)q2[0] | ((uint32_t)q2[1] << 16);
const uint32_t aux_s = (uint32_t)q2[2] | ((uint32_t)q2[3] << 16);
q2 += 4;
const float dl = d * (0.5f + (float)(aux_s >> 28)) * 0.25f;
const uint8_t grids[4] = {
(uint8_t)(aux_g & 0xffu),
(uint8_t)((aux_g >> 8) & 0xffu),
(uint8_t)((aux_g >> 16) & 0xffu),
(uint8_t)((aux_g >> 24) & 0xffu),
};
for (uint32_t half = 0; half < 2; half++) {
for (uint32_t g = 0; g < 2; g++) {
const uint32_t gi = half * 2 + g;
const uint64_t grid = cuda_iq2xxs_grid[grids[gi]];
const uint8_t signs = cuda_ksigns_iq2xs[(aux_s >> (14u * half + 7u * g)) & 127u];
for (uint32_t i = 0; i < 8; i++) {
float w = (float)((grid >> (8u * i)) & 0xffu);
if (signs & (1u << i)) w = -w;
acc += dl * w * xf[ib32 * 32u + half * 16u + g * 8u + i];
}
}
}
}
}
return acc;
}
__device__ static float dev_q2_K_dot_f32(const cuda_block_q2_K *row, const float *x, uint32_t nb) {
float acc = 0.0f;
for (uint32_t b = 0; b < nb; b++) {
const cuda_block_q2_K *xb = row + b;
const float d = dev_f16_to_f32(xb->d);
const float dmin = dev_f16_to_f32(xb->dmin);
for (uint32_t il = 0; il < 16; il++) {
const uint32_t chunk = il / 8u;
const uint32_t pair = il & 1u;
const uint32_t shift = ((il / 2u) & 3u) * 2u;
const uint8_t sc = xb->scales[il];
const float dl = d * (float)(sc & 0x0fu);
const float ml = dmin * (float)(sc >> 4);
const uint8_t *q = xb->qs + 32u * chunk + 16u * pair;
const float *xf = x + (uint64_t)b * CUDA_QK_K + chunk * 128u + ((il % 8u) / 2u) * 32u + pair * 16u;
for (uint32_t i = 0; i < 16; i++) {
const float w = dl * (float)((q[i] >> shift) & 3u) - ml;
acc += w * xf[i];
}
}
}
return acc;
}
__global__ static void moe_gate_up_mid_f32_kernel(
float *gate_out,
float *up_out,
float *mid_out,
const char *gate_base,
const char *up_base,
const float *x,
const int32_t *selected,
const float *weights,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t n_expert,
float clamp) {
uint32_t row = blockIdx.x;
uint32_t pair = blockIdx.y;
if (row >= expert_mid_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
uint32_t expert = (uint32_t)expert_i;
const uint32_t nb = expert_in_dim / CUDA_QK_K;
const cuda_block_iq2_xxs *gr = (const cuda_block_iq2_xxs *)(gate_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const cuda_block_iq2_xxs *ur = (const cuda_block_iq2_xxs *)(up_base + (uint64_t)expert * gate_expert_bytes + (uint64_t)row * gate_row_bytes);
const float *xr = x + (uint64_t)tok * expert_in_dim;
float gate = 0.0f;
float up = 0.0f;
for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) {
gate += dev_iq2_xxs_dot_f32(gr + b, xr + (uint64_t)b * CUDA_QK_K, 1);
up += dev_iq2_xxs_dot_f32(ur + b, xr + (uint64_t)b * CUDA_QK_K, 1);
}
__shared__ float partial_gate[256];
__shared__ float partial_up[256];
partial_gate[threadIdx.x] = gate;
partial_up[threadIdx.x] = up;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) {
partial_gate[threadIdx.x] += partial_gate[threadIdx.x + stride];
partial_up[threadIdx.x] += partial_up[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
gate = partial_gate[0];
up = partial_up[0];
if (clamp > 1.0e-6f) {
if (gate > clamp) gate = clamp;
if (up > clamp) up = clamp;
if (up < -clamp) up = -clamp;
}
const uint64_t off = (uint64_t)pair * expert_mid_dim + row;
gate_out[off] = gate;
up_out[off] = up;
mid_out[off] = (gate / (1.0f + expf(-gate))) * up * weights[(uint64_t)tok * n_expert + slot];
}
}
__global__ static void moe_down_f32_kernel(
float *down_out,
const char *down_base,
const float *mid,
const int32_t *selected,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_mid_dim,
uint32_t out_dim,
uint32_t n_expert) {
uint32_t row = blockIdx.x;
uint32_t pair = blockIdx.y;
if (row >= out_dim) return;
uint32_t tok = pair / n_expert;
uint32_t slot = pair - tok * n_expert;
int32_t expert_i = selected[(uint64_t)tok * n_expert + slot];
if (expert_i < 0) expert_i = 0;
const uint32_t nb = expert_mid_dim / CUDA_QK_K;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base + (uint64_t)(uint32_t)expert_i * down_expert_bytes + (uint64_t)row * down_row_bytes);
const float *xr = mid + (uint64_t)pair * expert_mid_dim;
float acc = 0.0f;
for (uint32_t b = threadIdx.x; b < nb; b += blockDim.x) acc += dev_q2_K_dot_f32(wr + b, xr + (uint64_t)b * CUDA_QK_K, 1);
__shared__ float partial[256];
partial[threadIdx.x] = acc;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
__syncthreads();
}
if (threadIdx.x == 0) down_out[(uint64_t)pair * out_dim + row] = partial[0];
}
static int routed_moe_launch(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *down,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float clamp,
const ds4_gpu_tensor *x,
uint32_t layer_index,
uint32_t n_tokens,
int allow_streaming,
int owned_filtered) {
if (!out || !gate || !up || !mid || !down || !model_map || !selected || !weights || !x ||
n_tokens == 0 || n_total_expert == 0 || n_expert == 0 ||
expert_in_dim % CUDA_QK_K != 0 || expert_mid_dim % CUDA_QK_K != 0 ||
gate_offset > model_size || up_offset > model_size || down_offset > model_size ||
x->bytes < (uint64_t)n_tokens * expert_in_dim * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * n_expert * sizeof(int32_t) ||
weights->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) ||
gate->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) ||
up->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) ||
mid->bytes < (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float) ||
down->bytes < (uint64_t)n_tokens * n_expert * out_dim * sizeof(float) ||
out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float)) {
return 0;
}
const int q4k_path = (gate_type == 12u && down_type == 12u);
const int iq2_path = (gate_type == 16u && down_type == 10u);
const int mxfp4_path = (gate_type == 39u && down_type == 39u);
if (!q4k_path && !iq2_path && !mxfp4_path) return 0;
if (g_ssd_streaming_mode && !owned_filtered) {
const ds4_gpu_stream_expert_table table = {
model_map, model_size, layer_index, n_total_expert,
gate_offset, up_offset, down_offset, gate_expert_bytes, down_expert_bytes};
if ((uint64_t)n_tokens * n_expert > UINT32_MAX ||
!ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
&table, selected, n_tokens * n_expert)) return 0;
allow_streaming = 1;
}
/* The aligned artifacts replace the raw expert tensors on integrated
* CUDA systems. Route both prefill and decode before resolving a raw
* pointer, otherwise the fallback cache would duplicate tens of GiB. */
if (iq2_path && !g_ssd_streaming_mode &&
cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled()) {
const uint64_t gate_total = (uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_total = (uint64_t)n_total_expert * down_expert_bytes;
const uint64_t gate_aligned_bytes = ds4_mmq_iq2_xxs_aligned_bytes(
(int)expert_mid_dim, (int)expert_in_dim, (int)n_total_expert);
const uint64_t down_aligned_bytes = ds4_mmq_q2_k_aligned_bytes(
(int)out_dim, (int)expert_mid_dim, (int)n_total_expert);
const char *gate_aligned = cuda_derived_weight_ptr(
model_map, gate_offset, gate_total,
CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE,
expert_in_dim, expert_mid_dim, n_total_expert,
gate_aligned_bytes);
const char *up_aligned = cuda_derived_weight_ptr(
model_map, up_offset, gate_total,
CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE,
expert_in_dim, expert_mid_dim, n_total_expert,
gate_aligned_bytes);
const char *down_aligned = cuda_derived_weight_ptr(
model_map, down_offset, down_total,
CUDA_DERIVED_Q2_K_ALIGNED_MOE,
expert_mid_dim, out_dim, n_total_expert,
down_aligned_bytes);
if (gate_aligned && up_aligned && down_aligned) {
const cudaStream_t aligned_stream =
n_tokens == 1u ? cuda_decode_stream() : (cudaStream_t)0;
if (owned_filtered && n_tokens <= 8u && (n_expert == 3u || n_expert == 6u)) {
/* Keep TP decode and short appends on their Q8_K arithmetic.
* Only the lossless weight layout changes here. */
const uint32_t xb = expert_in_dim / CUDA_QK_K;
const uint32_t mb = expert_mid_dim / CUDA_QK_K;
if (down->bytes < (uint64_t)n_tokens * xb * sizeof(cuda_block_q8_K) ||
gate->bytes < (uint64_t)n_tokens * n_expert * mb * sizeof(cuda_block_q8_K)) return 0;
auto *xq = (cuda_block_q8_K *)down->ptr;
auto *mq = (cuda_block_q8_K *)gate->ptr;
const uint64_t scales = ((uint64_t)n_total_expert * expert_mid_dim * xb * 2u + 63u) & ~63ull;
q8_K_quantize_kernel<<<dim3(xb, n_tokens), 256, 0, aligned_stream>>>(
xq, (const float *)x->ptr, expert_in_dim, n_tokens);
if (!cuda_ok(cudaGetLastError(), "aligned owned x quantize")) return 0;
moe_gate_up_mid_aligned_q8k_kernel<<<
dim3((expert_mid_dim + 7u) / 8u, n_tokens * n_expert), 256, 0, aligned_stream>>>(
(float *)mid->ptr, gate_aligned, up_aligned, xq,
(const int32_t *)selected->ptr, (const float *)weights->ptr,
xb, expert_mid_dim, n_expert, scales, clamp);
if (!cuda_ok(cudaGetLastError(), "aligned owned gate/up")) return 0;
q8_K_quantize_owned_kernel<<<dim3(mb, n_tokens * n_expert), 256, 0, aligned_stream>>>(
mq, (const float *)mid->ptr, (const int32_t *)selected->ptr,
expert_mid_dim, n_tokens * n_expert, 0, n_total_expert);
if (!cuda_ok(cudaGetLastError(), "aligned owned mid quantize")) return 0;
const dim3 grid((out_dim + 31u) / 32u, n_tokens);
if (n_expert == 6u) {
moe_down_sum_qwarp32_kernel<6, true><<<grid, 256, 0, aligned_stream>>>(
(float *)out->ptr, down_aligned, mq, (const int32_t *)selected->ptr,
down_expert_bytes, down_row_bytes, mb, out_dim, n_total_expert);
} else {
moe_down_sum_qwarp32_kernel<3, true><<<grid, 256, 0, aligned_stream>>>(
(float *)out->ptr, down_aligned, mq, (const int32_t *)selected->ptr,
down_expert_bytes, down_row_bytes, mb, out_dim, n_total_expert);
}
return cuda_ok(cudaGetLastError(), "aligned owned down");
}
int rc;
if (n_tokens == 1u ||
(n_tokens <= 8u && ds4_gpu_device_is_spark() &&
!g_q8_dequant_gemm_enabled && expert_in_dim % 1024u == 0)) {
rc = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec(
gate_aligned, up_aligned,
(const float *)x->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
(float *)mid->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)n_tokens, (int)n_total_expert, (int)n_expert,
clamp, aligned_stream);
if (rc == 0) {
const uint32_t assignments = n_tokens * n_expert;
rc = ds4_mmq_q2_K_aligned_moe_vec(
down_aligned, (const float *)mid->ptr,
(const int32_t *)selected->ptr,
(float *)down->ptr,
(int)out_dim, (int)expert_mid_dim,
(int)assignments, (int)n_total_expert,
/*n_expert_used=*/1,
aligned_stream);
}
} else {
rc = 1;
const uint64_t assignments =
(uint64_t)n_tokens * n_expert;
if (!owned_filtered && assignments >= 1024u) {
size_t input_q8_bytes = 0;
size_t down_q8_bytes = 0;
size_t work_bytes = 0;
if (ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_scratch_sizes(
(int)expert_mid_dim, (int)expert_in_dim,
(int)n_tokens, (int)n_total_expert,
(int)n_expert, &input_q8_bytes,
&down_q8_bytes, &work_bytes) == 0 &&
gate->bytes >= input_q8_bytes &&
up->bytes >= down_q8_bytes &&
mid->bytes >= work_bytes) {
rc = ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_soa(
gate_aligned, up_aligned, down_aligned,
(const float *)x->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate->ptr, input_q8_bytes,
up->ptr, down_q8_bytes,
mid->ptr, work_bytes,
NULL, 0,
(float *)down->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)out_dim, (int)n_tokens,
(int)n_total_expert, (int)n_expert,
clamp, aligned_stream);
if (rc == 0) {
static int direct_logged = 0;
if (!direct_logged) {
direct_logged = 1;
fprintf(stderr,
"ds4: routed MoE using fused-direct CUDA prefill\n");
}
}
}
}
if (rc != 0) {
rc = ds4_mmq_iq2_xxs_q2_K_moe_fused_soa(
gate_aligned, up_aligned, down_aligned,
(const float *)x->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
(float *)gate->ptr, (float *)up->ptr,
(float *)mid->ptr, (float *)down->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)out_dim, (int)n_tokens,
(int)n_total_expert, (int)n_expert,
clamp, owned_filtered, aligned_stream);
}
}
if (rc == 0) {
const uint64_t n = (uint64_t)n_tokens * out_dim;
moe_mmq_sum_kernel<<<
(uint32_t)((n + 255u) / 256u), 256, 0, aligned_stream>>>(
(float *)out->ptr, (const float *)down->ptr,
owned_filtered ? (const int32_t *)selected->ptr : NULL,
out_dim, n_expert, n_tokens,
/*guard_nonfinite=*/1);
if (cuda_ok(cudaGetLastError(), "aligned moe sum launch")) {
static int logged = 0;
if (!logged) {
logged = 1;
fprintf(stderr,
"ds4: routed MoE using aligned CUDA artifacts\n");
}
return 1;
}
rc = -1;
}
fprintf(stderr,
"ds4: aligned routed-MoE returned %d "
"(layer=%u n_tokens=%u)\n",
rc, layer_index, n_tokens);
if (cuda_model_map_replaces_complete(model_map)) return 0;
}
}
/* Native MXFP4 routed experts use the vendored MMVQ decode kernels and
* MMQ matrix kernels. On Blackwell the latter dispatch to FP4 MMA; older
* CUDA devices use the mathematically equivalent DP4A implementation.
* SSD streaming presents a compact table containing only the selected
* experts, together with IDs remapped into that table. */
if (mxfp4_path) {
if (!cuda_use_mxfp4_mmq()) {
fprintf(stderr, "ds4: CUDA MXFP4 requires the MMQ backend\n");
return 0;
}
const uint64_t gate_total =
(uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_total =
(uint64_t)n_total_expert * down_expert_bytes;
if (gate_total > model_size - gate_offset ||
gate_total > model_size - up_offset ||
down_total > model_size - down_offset) {
return 0;
}
const uint64_t slot_count = (uint64_t)n_tokens * n_expert;
const int logical_tier = ds4_tensor_device_idx(out);
const int use_stream_selected_cache =
allow_streaming &&
g_ssd_streaming_mode &&
g_stream_selected_cache.valid &&
g_stream_selected_cache.logical_tier == logical_tier &&
g_stream_selected_cache.model_map == model_map &&
g_stream_selected_cache.layer == layer_index &&
g_stream_selected_cache.n_total_expert == n_total_expert &&
g_stream_selected_cache.slot_count >= slot_count &&
g_stream_selected_cache.gate_offset == gate_offset &&
g_stream_selected_cache.up_offset == up_offset &&
g_stream_selected_cache.down_offset == down_offset &&
g_stream_selected_cache.gate_expert_bytes == gate_expert_bytes &&
g_stream_selected_cache.down_expert_bytes == down_expert_bytes &&
g_stream_selected_cache.gate_ptr &&
g_stream_selected_cache.up_ptr &&
g_stream_selected_cache.down_ptr &&
g_stream_selected_cache.slot_selected_tensor.ptr &&
g_stream_selected_cache.slot_selected_tensor.bytes >=
slot_count * sizeof(int32_t);
if (g_ssd_streaming_mode && allow_streaming &&
!use_stream_selected_cache) {
fprintf(stderr,
"ds4: CUDA streaming MXFP4 experts are unavailable for layer %u\n",
layer_index);
return 0;
}
const ds4_gpu_tensor *mx_selected = use_stream_selected_cache ?
&g_stream_selected_cache.slot_selected_tensor : selected;
uint32_t weight_experts = use_stream_selected_cache ?
g_stream_selected_cache.compact_count : n_total_expert;
const char *gate_w = use_stream_selected_cache ?
g_stream_selected_cache.gate_ptr :
cuda_resolve_weight_ptr(model_map, gate_offset, gate_total,
logical_tier, "mxfp4 moe gate");
const char *up_w = use_stream_selected_cache ?
g_stream_selected_cache.up_ptr :
cuda_resolve_weight_ptr(model_map, up_offset, gate_total,
logical_tier, "mxfp4 moe up");
const char *down_w = use_stream_selected_cache ?
g_stream_selected_cache.down_ptr :
cuda_resolve_weight_ptr(model_map, down_offset, down_total,
logical_tier, "mxfp4 moe down");
if (!gate_w || !up_w || !down_w || weight_experts == 0u) return 0;
ds4_gpu_tensor compact_selected = *mx_selected;
if (use_stream_selected_cache && n_tokens > 1u) {
const int32_t *ids = NULL;
if (!cuda_stream_compact_prefill(&gate_w, &up_w, &down_w,
&ids, &weight_experts)) return 0;
compact_selected.ptr = (void *)ids;
mx_selected = &compact_selected;
}
const cudaStream_t stream =
n_tokens == 1u ? cuda_decode_stream() : (cudaStream_t)0;
int rc = -1;
if (n_tokens == 1u && n_expert == 6u) {
rc = ds4_mmq_mxfp4_moe_gate_up_mid_vec(
gate_w, up_w, (const float *)x->ptr,
(const int32_t *)mx_selected->ptr,
(const float *)weights->ptr, (float *)mid->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)n_tokens, (int)weight_experts, (int)n_expert,
clamp, stream);
if (rc == 0) {
rc = ds4_mmq_mxfp4_moe_down_sum6_vec(
down_w, (const float *)mid->ptr,
(const int32_t *)mx_selected->ptr, (float *)out->ptr,
(int)out_dim, (int)expert_mid_dim,
(int)n_tokens, (int)weight_experts, (int)n_expert,
stream);
}
} else {
rc = ds4_mmq_mxfp4_moe_pair(
gate_w, up_w, (const float *)x->ptr,
(const int32_t *)mx_selected->ptr,
(float *)gate->ptr, (float *)up->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)n_tokens, (int)weight_experts, (int)n_expert,
stream);
if (rc == 0) {
const uint64_t mid_floats =
slot_count * expert_mid_dim;
moe_mmq_swiglu_weighted_clamp_kernel<<<
(uint32_t)((mid_floats + 255u) / 256u), 256, 0, stream>>>(
(float *)mid->ptr,
(const float *)gate->ptr, (const float *)up->ptr,
(const float *)weights->ptr,
expert_mid_dim, n_tokens, n_expert, clamp);
rc = cuda_ok(cudaGetLastError(),
"mxfp4 moe swiglu launch") ? 0 : -1;
}
if (rc == 0) {
rc = ds4_mmq_mxfp4_moe(
down_w, (const float *)mid->ptr,
(const int32_t *)mx_selected->ptr,
(float *)down->ptr,
(int)out_dim, (int)expert_mid_dim,
(int)slot_count, (int)weight_experts,
/*n_expert_used=*/1, stream);
}
if (rc == 0) {
const uint64_t n = (uint64_t)n_tokens * out_dim;
moe_mmq_sum_kernel<<<
(uint32_t)((n + 255u) / 256u), 256, 0, stream>>>(
(float *)out->ptr, (const float *)down->ptr,
owned_filtered
? (const int32_t *)mx_selected->ptr
: NULL,
out_dim, n_expert, n_tokens, /*guard_nonfinite=*/1);
rc = cuda_ok(cudaGetLastError(),
"mxfp4 moe sum launch") ? 0 : -1;
}
}
if (rc == 0) return 1;
fprintf(stderr,
"ds4: CUDA MXFP4 routed-MoE returned %d "
"(layer=%u n_tokens=%u)\n",
rc, layer_index, n_tokens);
return 0;
}
/* Q4_K routed-MoE dispatch:
* n_tokens == 1 and n_expert == 6:
* use_direct_down_sum + moe_gate_up_mid_decode_q4K_qwarp32
* + moe_down_q4K_sum6_qwarp32.
* n_tokens == 1 and n_expert == 3:
* use the same direct path with moe_down_q4K_sum3_qwarp32.
* Decode TP relies on this for splitting the six selected
* experts into two groups.
* n_tokens == 1 and other n_expert:
* use the same per-pair gate/up kernel plus the generic
* q4K down + sum path.
* n_tokens > 1: default sorted-pairs expert-tile path groups token/expert
* pairs by expert and uses Q4_K tile8 gate/up + down kernels
* (`DS4_CUDA_MOE_NO_Q4_SORTED=1` restores the older
* token-indexed decode-style prefill kernels). */
const uint64_t gate_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_bytes = (uint64_t)n_total_expert * down_expert_bytes;
if (gate_bytes > model_size - gate_offset ||
gate_bytes > model_size - up_offset ||
down_bytes > model_size - down_offset) {
return 0;
}
const uint64_t required_slot_count = (uint64_t)n_tokens * n_expert;
const int logical_tier = ds4_tensor_device_idx(out);
const int use_stream_selected_cache =
allow_streaming &&
g_ssd_streaming_mode &&
g_stream_selected_cache.valid &&
g_stream_selected_cache.logical_tier == logical_tier &&
g_stream_selected_cache.model_map == model_map &&
g_stream_selected_cache.layer == layer_index &&
g_stream_selected_cache.n_total_expert == n_total_expert &&
g_stream_selected_cache.slot_count >= required_slot_count &&
g_stream_selected_cache.gate_offset == gate_offset &&
g_stream_selected_cache.up_offset == up_offset &&
g_stream_selected_cache.down_offset == down_offset &&
g_stream_selected_cache.gate_expert_bytes == gate_expert_bytes &&
g_stream_selected_cache.down_expert_bytes == down_expert_bytes &&
g_stream_selected_cache.gate_ptr &&
g_stream_selected_cache.up_ptr &&
g_stream_selected_cache.down_ptr &&
g_stream_selected_cache.slot_selected_tensor.ptr &&
g_stream_selected_cache.slot_selected_tensor.bytes >=
required_slot_count * sizeof(int32_t);
if (g_ssd_streaming_mode && allow_streaming &&
!use_stream_selected_cache) {
fprintf(stderr,
"ds4: CUDA streaming selected experts are unavailable for layer %u\n",
layer_index);
return 0;
}
if (use_stream_selected_cache) {
selected = &g_stream_selected_cache.slot_selected_tensor;
n_total_expert = g_stream_selected_cache.compact_count;
}
const char *gate_w = use_stream_selected_cache ?
g_stream_selected_cache.gate_ptr :
cuda_resolve_weight_ptr(model_map, gate_offset, gate_bytes,
logical_tier, "moe_gate");
const char *up_w = use_stream_selected_cache ?
g_stream_selected_cache.up_ptr :
cuda_resolve_weight_ptr(model_map, up_offset, gate_bytes,
logical_tier, "moe_up");
const char *down_w = use_stream_selected_cache ?
g_stream_selected_cache.down_ptr :
cuda_resolve_weight_ptr(model_map, down_offset, down_bytes,
logical_tier, "moe_down");
if (!gate_w || !up_w || !down_w) return 0;
/* V4.1's short TP prefills need scalar-row arithmetic. Keep existing
* multi-GPU batching dispatch unchanged for other model geometries. */
const bool exact_owned_rows = owned_filtered && expert_in_dim == 5120u &&
expert_mid_dim == 2304u;
const bool small_exact_batch = iq2_path && (use_stream_selected_cache || exact_owned_rows) &&
n_tokens <= 8u && (n_expert == 3u || n_expert == 6u) &&
getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6") == NULL;
/* mmq routed-MoE prefill tier (ported from the Entrpi/ds4 fork).
* IQ2_XXS gate/up pair (one shared activation quantize + routing
* pass) -> SwiGLU + clamp + router weight -> Q2_K down, treating
* each (token, slot) assignment as its own single-expert row ->
* guarded slot sum. Buffers gate/up/mid/down are already sized to
* [n_tokens, n_expert, *] by the validation above. Any entry
* failure falls through to the legacy sorted-pairs path (the
* buffers are scratch there too). Streaming uses the same bounded cache
* and remapped IDs as the other routed kernels. */
const bool owned_mmq = owned_filtered && n_tokens >= 128u && n_total_expert >= n_expert &&
getenv("DS4_CUDA_MOE_NO_OWNED_MMQ") == NULL;
if (iq2_path && n_tokens > 1u && (!owned_filtered || owned_mmq) &&
!small_exact_batch && cuda_use_mmq()) {
const uint64_t n_assignments = (uint64_t)n_tokens * n_expert;
const char *mmq_gate = gate_w, *mmq_up = up_w, *mmq_down = down_w;
const int32_t *mmq_ids = (const int32_t *)selected->ptr;
uint32_t mmq_experts = n_total_expert;
const bool aligned = use_stream_selected_cache && n_tokens >= 128u &&
n_assignments >= 1024u && expert_in_dim % 1024u == 0u &&
cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled();
if (use_stream_selected_cache &&
!cuda_stream_compact_prefill(&mmq_gate, &mmq_up, &mmq_down,
&mmq_ids, &mmq_experts,
aligned ? expert_in_dim : 0,
expert_mid_dim, out_dim)) return 0;
if (aligned) {
size_t input_bytes = 0, down_bytes = 0, work_bytes = 0;
int rc = 1;
if (ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_scratch_sizes(
expert_mid_dim, expert_in_dim, n_tokens, mmq_experts, n_expert,
&input_bytes, &down_bytes, &work_bytes) == 0 &&
gate->bytes >= input_bytes && up->bytes >= down_bytes && mid->bytes >= work_bytes) {
rc = ds4_mmq_iq2_xxs_q2_K_moe_fused_direct_soa(
mmq_gate, mmq_up, mmq_down, (const float *)x->ptr, mmq_ids,
(const float *)weights->ptr, gate->ptr, input_bytes,
up->ptr, down_bytes, mid->ptr, work_bytes, NULL, 0,
(float *)down->ptr, expert_mid_dim, expert_in_dim, out_dim,
n_tokens, mmq_experts, n_expert, clamp, (cudaStream_t)0);
}
if (rc != 0) rc = ds4_mmq_iq2_xxs_q2_K_moe_fused_soa(
mmq_gate, mmq_up, mmq_down, (const float *)x->ptr, mmq_ids,
(const float *)weights->ptr, (float *)gate->ptr, (float *)up->ptr,
(float *)mid->ptr, (float *)down->ptr, expert_mid_dim, expert_in_dim,
out_dim, n_tokens, mmq_experts, n_expert, clamp, 0, (cudaStream_t)0);
if (rc != 0) return 0;
const uint64_t n = (uint64_t)n_tokens * out_dim;
moe_mmq_sum_kernel<<<(unsigned)((n + 255u) / 256u), 256>>>(
(float *)out->ptr, (const float *)down->ptr, NULL,
out_dim, n_expert, n_tokens, 1);
return cuda_ok(cudaGetLastError(), "aligned streaming moe sum");
}
int rc = ds4_mmq_iq2_xxs_moe_pair(
mmq_gate, mmq_up, (const float *)x->ptr, mmq_ids,
(float *)gate->ptr, (float *)up->ptr,
(int)expert_mid_dim, (int)expert_in_dim,
(int)n_tokens, (int)mmq_experts, (int)n_expert,
(cudaStream_t)0);
if (rc == 0) {
const uint64_t mid_floats = n_assignments * expert_mid_dim;
moe_mmq_swiglu_weighted_clamp_kernel<<<(uint32_t)((mid_floats + 255) / 256), 256>>>(
(float *)mid->ptr,
(const float *)gate->ptr, (const float *)up->ptr,
(const float *)weights->ptr,
expert_mid_dim, n_tokens, n_expert, clamp);
rc = cuda_ok(cudaGetLastError(), "mmq moe swiglu launch") ? 0 : -1;
}
if (rc == 0) {
rc = ds4_mmq_q2_K_moe(
mmq_down, (const float *)mid->ptr, mmq_ids,
(float *)down->ptr,
(int)out_dim, (int)expert_mid_dim,
(int)n_assignments, (int)mmq_experts,
/*n_expert_used=*/1,
(cudaStream_t)0);
}
if (rc == 0) {
const uint64_t n = (uint64_t)n_tokens * out_dim;
moe_mmq_sum_kernel<<<(uint32_t)((n + 255) / 256), 256>>>(
(float *)out->ptr, (const float *)down->ptr,
owned_filtered ? mmq_ids : NULL, out_dim, n_expert, n_tokens,
/*guard_nonfinite=*/1);
if (cuda_ok(cudaGetLastError(), "mmq moe sum launch")) return 1;
rc = -1;
}
fprintf(stderr, "ds4: mmq routed-MoE tier rc=%d (layer=%u n_tokens=%u); falling back\n",
rc, layer_index, n_tokens);
}
int ok = 1;
const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K;
const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K;
const uint64_t xq_count = (uint64_t)n_tokens * xq_blocks;
const uint64_t midq_count = (uint64_t)n_tokens * n_expert * midq_blocks;
const uint64_t xq_bytes = xq_count * sizeof(cuda_block_q8_K);
const uint64_t midq_bytes = midq_count * sizeof(cuda_block_q8_K);
if (down->bytes >= xq_bytes && gate->bytes >= midq_bytes) {
cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr;
cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr;
const uint32_t profile_moe = getenv("DS4_CUDA_MOE_PROFILE") != NULL;
cudaEvent_t prof_ev[7] = {NULL, NULL, NULL, NULL, NULL, NULL, NULL};
if (profile_moe) {
for (uint32_t i = 0; i < 7u; i++) {
if (cudaEventCreate(&prof_ev[i]) != cudaSuccess) {
for (uint32_t j = 0; j < i; j++) (void)cudaEventDestroy(prof_ev[j]);
memset(prof_ev, 0, sizeof(prof_ev));
break;
}
}
if (prof_ev[0]) (void)cudaEventRecord(prof_ev[0], 0);
}
const uint32_t pair_count = n_tokens * n_expert;
const uint32_t use_q4_sorted_pairs =
q4k_path && n_tokens > 1u &&
(owned_filtered ||
(getenv("DS4_CUDA_MOE_NO_Q4_SORTED") == NULL &&
getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL &&
getenv("DS4_CUDA_MOE_TILE4") == NULL));
const uint32_t use_sorted_pairs =
n_tokens > 1u && !small_exact_batch &&
(owned_filtered || !q4k_path || use_q4_sorted_pairs);
const uint32_t use_expert_tiles =
use_sorted_pairs &&
(owned_filtered || getenv("DS4_CUDA_MOE_NO_EXPERT_TILES") == NULL);
const uint32_t q4_owned_batch =
owned_filtered && q4k_path && n_tokens >= 4u && n_tokens <= 16u;
/* Small batches (DSpark stage chain / verify, n<=8) leave most of an
* 8-slot expert tile empty (1-2 rows per expert): tile4 halves the
* wasted dot-slots and measures ~2x faster there. Q4 TP batches use
* tile8 because its exact tensor-core row-span kernels recover more
* than the empty slots cost. Large prefill also keeps tile8. */
const uint32_t expert_tile_m =
getenv("DS4_CUDA_MOE_TILE4") ? 4u :
(getenv("DS4_CUDA_MOE_TILE8") ? 8u :
(q4_owned_batch || n_tokens > 8u ? 8u : 4u));
const uint32_t write_gate_up = getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL;
const uint32_t use_p2_sorted =
use_sorted_pairs && !owned_filtered &&
getenv("DS4_CUDA_MOE_NO_P2") == NULL;
const uint32_t use_atomic_down = !q4k_path && use_expert_tiles &&
(getenv("DS4_CUDA_MOE_ATOMIC_DOWN") != NULL ||
(n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_ATOMIC_DOWN") == NULL));
const uint32_t use_owned_sparse_buffers = owned_filtered &&
getenv("DS4_CUDA_MOE_NO_OWNED_SPARSE_BUFFERS") == NULL;
const uint32_t use_gate_row2048 = use_expert_tiles && expert_tile_m == 8u &&
(getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL ||
getenv("DS4_CUDA_MOE_GATE_ROW256") != NULL ||
getenv("DS4_CUDA_MOE_GATE_ROW128") != NULL ||
((q4_owned_batch || n_tokens >= 128u) &&
getenv("DS4_CUDA_MOE_NO_GATE_ROW2048") == NULL &&
getenv("DS4_CUDA_MOE_NO_GATE_ROW256") == NULL &&
getenv("DS4_CUDA_MOE_NO_GATE_ROW128") == NULL));
const uint32_t use_q4_mma_tiles16 = q4k_path && use_expert_tiles &&
expert_tile_m == 8u && cuda_q4_mma_ok() &&
getenv("DS4_CUDA_MOE_NO_Q4_MMA_TILE16") == NULL;
const uint32_t use_down_tile16 = !q4k_path && use_atomic_down && expert_tile_m == 8u &&
n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_DOWN_TILE16") == NULL;
const uint32_t use_small_sorted_prep =
owned_filtered && q4k_path && n_tokens <= 16u && pair_count <= 96u &&
n_total_expert <= 128u && use_sorted_pairs && use_expert_tiles &&
getenv("DS4_CUDA_MOE_NO_SMALL_SORTED_PREP") == NULL;
const uint32_t force_q4_down_rowspan =
getenv("DS4_CUDA_MOE_DOWN_ROW512") != NULL ||
getenv("DS4_CUDA_MOE_DOWN_ROW1024") != NULL ||
getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL;
const uint32_t use_q4_down_rowspan =
q4k_path && use_expert_tiles && expert_tile_m == 8u &&
(q4_owned_batch || n_tokens >= 128u || force_q4_down_rowspan) &&
getenv("DS4_CUDA_MOE_NO_Q4_DOWN_ROWSPAN") == NULL;
const uint32_t use_decode_lut_gate =
(n_tokens == 1u || small_exact_batch) &&
getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") == NULL;
const uint32_t gate_row_span =
getenv("DS4_CUDA_MOE_GATE_ROW2048") != NULL ? 2048u :
getenv("DS4_CUDA_MOE_GATE_ROW1024") != NULL ? 1024u : 512u;
const uint32_t down_row_span =
getenv("DS4_CUDA_MOE_DOWN_ROW512") != NULL ? 512u :
getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL ? 2048u :
getenv("DS4_CUDA_MOE_DOWN_ROW1024") != NULL ? 1024u : 512u;
const uint32_t use_down_row2048 = !q4k_path && use_atomic_down && expert_tile_m == 8u &&
(getenv("DS4_CUDA_MOE_DOWN_ROW2048") != NULL ||
getenv("DS4_CUDA_MOE_DOWN_ROW256") != NULL ||
getenv("DS4_CUDA_MOE_DOWN_ROW128") != NULL ||
getenv("DS4_CUDA_MOE_DOWN_ROW64") != NULL ||
(use_down_tile16 &&
getenv("DS4_CUDA_MOE_NO_DOWN_ROW2048") == NULL &&
getenv("DS4_CUDA_MOE_NO_DOWN_ROW256") == NULL &&
getenv("DS4_CUDA_MOE_NO_DOWN_ROW128") == NULL &&
getenv("DS4_CUDA_MOE_NO_DOWN_ROW64") == NULL));
const uint32_t use_direct_down_sum =
(n_tokens == 1u || small_exact_batch) && (n_expert == 6u || n_expert == 3u) &&
getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6") == NULL;
const uint32_t use_direct_midq =
q4k_path && use_direct_down_sum && !write_gate_up &&
getenv("DS4_CUDA_MOE_DIRECT_MIDQ") != NULL &&
getenv("DS4_CUDA_MOE_NO_DIRECT_MIDQ") == NULL;
const uint32_t use_q4_gate_h16r8 =
q4k_path && !use_direct_midq &&
getenv("DS4_CUDA_MOE_Q4_GATE_H16R8") != NULL &&
getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16R8") == NULL;
const uint32_t use_q4_gate_h16 =
q4k_path && !use_direct_midq && !use_q4_gate_h16r8 &&
getenv("DS4_CUDA_MOE_Q4_GATE_H16") != NULL &&
getenv("DS4_CUDA_MOE_NO_Q4_GATE_H16") == NULL;
const uint32_t use_q4_gate_w32r16 =
q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 &&
getenv("DS4_CUDA_MOE_Q4_GATE_W32R16") != NULL &&
getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32R16") == NULL;
const uint32_t use_q4_gate_w32 =
q4k_path && !use_direct_midq && !use_q4_gate_h16r8 && !use_q4_gate_h16 &&
!use_q4_gate_w32r16 &&
getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32") == NULL;
const uint32_t use_q4_gate_w32_noaux =
use_q4_gate_w32 && !write_gate_up &&
getenv("DS4_CUDA_MOE_NO_Q4_GATE_W32_NOAUX") == NULL;
const uint32_t use_q4_down_slot3 =
q4k_path && use_direct_down_sum && n_expert == 3u &&
getenv("DS4_CUDA_MOE_Q4_DOWN_SLOT3") != NULL &&
getenv("DS4_CUDA_MOE_NO_Q4_DOWN_SLOT3") == NULL;
const uint32_t use_q4_midq_sidecar =
q4k_path && use_direct_down_sum && use_q4_gate_w32_noaux &&
!use_direct_midq && !write_gate_up &&
(expert_mid_dim % CUDA_QK_K) == 0u &&
getenv("DS4_CUDA_MOE_MIDQ_SIDECAR") != NULL &&
getenv("DS4_CUDA_MOE_NO_MIDQ_SIDECAR") == NULL;
float *midq_sidecar = use_q4_midq_sidecar ? (float *)up->ptr : NULL;
if (g_cuda_moe_decode_graph &&
!owned_filtered &&
!profile_moe &&
q4k_path &&
n_tokens == 1u &&
use_direct_down_sum &&
use_q4_gate_w32 &&
!use_q4_gate_w32r16 &&
!use_q4_down_slot3 &&
!use_direct_midq &&
!use_q4_midq_sidecar &&
(n_expert == 3u || n_expert == 6u)) {
int grc = routed_moe_decode_q4_graph_launch(
logical_tier,
(float *)out->ptr,
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
down_w,
xq,
midq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
n_expert,
write_gate_up,
clamp,
(const float *)x->ptr);
if (grc == 1) return 1;
if (grc < 0) return 0;
}
uint32_t *sorted_pairs = NULL;
uint32_t *sorted_offsets = NULL;
uint32_t *sorted_counts = NULL;
uint32_t *tile_total = NULL;
uint32_t *tile_experts = NULL;
uint32_t *tile_starts = NULL;
uint32_t *tile16_total = NULL;
uint32_t *tile16_experts = NULL;
uint32_t *tile16_starts = NULL;
uint32_t tile_capacity = 0;
uint32_t tile16_capacity = 0;
dim3 xq_grid(xq_blocks, n_tokens, 1);
q8_K_quantize_kernel<<<xq_grid, 256, 0, cuda_decode_stream()>>>(xq, (const float *)x->ptr, expert_in_dim, n_tokens);
ok = cuda_ok(cudaGetLastError(), "routed_moe x quantize launch");
if (prof_ev[1]) (void)cudaEventRecord(prof_ev[1], 0);
if (ok && use_sorted_pairs) {
const uint64_t counts_bytes = (uint64_t)n_total_expert * sizeof(uint32_t);
const uint64_t offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t);
const uint64_t cursors_bytes = (uint64_t)n_total_expert * sizeof(uint32_t);
const uint64_t sorted_bytes = (uint64_t)pair_count * sizeof(uint32_t);
tile_capacity = (pair_count + expert_tile_m - 1u) / expert_tile_m + n_total_expert;
tile16_capacity = (use_down_tile16 || use_q4_mma_tiles16) ? ((pair_count + 15u) / 16u + n_total_expert) : 0u;
const uint64_t tile_offsets_bytes = ((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t);
const uint64_t tile_total_bytes = sizeof(uint32_t);
const uint64_t tile_experts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t);
const uint64_t tile_starts_bytes = (uint64_t)tile_capacity * sizeof(uint32_t);
const uint64_t tile16_offsets_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? (((uint64_t)n_total_expert + 1ull) * sizeof(uint32_t)) : 0u;
const uint64_t tile16_total_bytes = (use_down_tile16 || use_q4_mma_tiles16) ? sizeof(uint32_t) : 0u;
const uint64_t tile16_experts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t);
const uint64_t tile16_starts_bytes = (uint64_t)tile16_capacity * sizeof(uint32_t);
const uint64_t tile_offsets_off = counts_bytes + offsets_bytes + cursors_bytes + sorted_bytes;
const uint64_t tile_total_off = tile_offsets_off + tile_offsets_bytes;
const uint64_t tile_experts_off = tile_total_off + tile_total_bytes;
const uint64_t tile_starts_off = tile_experts_off + tile_experts_bytes;
const uint64_t tile16_offsets_off = tile_starts_off + tile_starts_bytes;
const uint64_t tile16_total_off = tile16_offsets_off + tile16_offsets_bytes;
const uint64_t tile16_experts_off = tile16_total_off + tile16_total_bytes;
const uint64_t tile16_starts_off = tile16_experts_off + tile16_experts_bytes;
const uint64_t scratch_bytes = tile16_starts_off + tile16_starts_bytes;
uint8_t *scratch = (uint8_t *)cuda_tmp_alloc_on(logical_tier, scratch_bytes,
"routed_moe sorted pairs");
if (!scratch) {
ok = 0;
} else {
uint32_t *counts = (uint32_t *)scratch;
uint32_t *offsets = (uint32_t *)(scratch + counts_bytes);
uint32_t *cursors = (uint32_t *)(scratch + counts_bytes + offsets_bytes);
sorted_pairs = (uint32_t *)(scratch + counts_bytes + offsets_bytes + cursors_bytes);
sorted_offsets = offsets;
sorted_counts = counts;
uint32_t *tile_offsets = (uint32_t *)(scratch + tile_offsets_off);
tile_total = (uint32_t *)(scratch + tile_total_off);
tile_experts = (uint32_t *)(scratch + tile_experts_off);
tile_starts = (uint32_t *)(scratch + tile_starts_off);
uint32_t *tile16_offsets = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_offsets_off) : NULL;
tile16_total = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_total_off) : NULL;
tile16_experts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_experts_off) : NULL;
tile16_starts = (use_down_tile16 || use_q4_mma_tiles16) ? (uint32_t *)(scratch + tile16_starts_off) : NULL;
if (use_small_sorted_prep) {
moe_prepare_sorted_tiles_small_kernel<<<1, 128, 0, cuda_decode_stream()>>>(
counts, offsets, cursors, sorted_pairs,
tile_offsets, tile_total, tile_experts, tile_starts,
tile16_offsets, tile16_total, tile16_experts, tile16_starts,
(const int32_t *)selected->ptr, pair_count, n_total_expert,
expert_tile_m, use_down_tile16 || use_q4_mma_tiles16);
ok = cuda_ok(cudaGetLastError(),
"routed_moe small sorted setup launch");
} else {
ok = cuda_ok(cudaMemset(counts, 0, counts_bytes),
"routed_moe sorted counts clear");
}
if (ok && !use_small_sorted_prep) {
moe_count_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
counts,
(const int32_t *)selected->ptr,
pair_count,
n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe sorted count launch");
}
if (ok && !use_small_sorted_prep) {
moe_prefix_sorted_pairs_kernel<<<1, 1, 0, cuda_decode_stream()>>>(offsets, cursors, counts, n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe sorted prefix launch");
}
if (ok && !use_small_sorted_prep) {
moe_scatter_sorted_pairs_kernel<<<(pair_count + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
sorted_pairs,
cursors,
(const int32_t *)selected->ptr,
pair_count,
n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe sorted scatter launch");
}
if (ok && use_expert_tiles && !use_small_sorted_prep) {
moe_build_expert_tile_offsets_kernel<<<1, 1, 0, cuda_decode_stream()>>>(tile_offsets, tile_total, counts, expert_tile_m, n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile offsets launch");
}
if (ok && use_expert_tiles && !use_small_sorted_prep) {
moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
tile_experts, tile_starts, tile_offsets, counts, expert_tile_m, n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe expert tiles launch");
}
if (ok && use_expert_tiles && !use_small_sorted_prep &&
(use_down_tile16 || use_q4_mma_tiles16)) {
moe_build_expert_tile_offsets_kernel<<<1, 1, 0, cuda_decode_stream()>>>(tile16_offsets, tile16_total, counts, 16u, n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 offsets launch");
}
if (ok && use_expert_tiles && !use_small_sorted_prep &&
(use_down_tile16 || use_q4_mma_tiles16)) {
moe_build_expert_tiles_kernel<<<(n_total_expert + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
tile16_experts, tile16_starts, tile16_offsets, counts, 16u, n_total_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe expert tile16 launch");
}
}
}
if (prof_ev[2]) (void)cudaEventRecord(prof_ev[2], 0);
if (ok && owned_filtered && use_sorted_pairs &&
!use_owned_sparse_buffers) {
const uint64_t mid_bytes =
(uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float);
ok = cuda_ok(cudaMemset(mid->ptr, 0, (size_t)mid_bytes),
"owned routed_moe mid clear");
}
if (ok) {
dim3 mgrid((expert_mid_dim + 31u) / 32u, n_tokens * n_expert, 1);
if (ok && sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && tile_total && tile_experts && tile_starts) {
if (q4k_path) {
const int use_q4_mma = cuda_q4_mma_ok() &&
((((uintptr_t)gate_w | (uintptr_t)up_w |
gate_row_bytes | gate_expert_bytes) & 15u) == 0u) &&
xq_blocks <= 16u && (expert_mid_dim & 7u) == 0u;
const int use_q4_mma_t16 = use_q4_mma && use_q4_mma_tiles16 &&
tile16_total && tile16_experts && tile16_starts &&
xq_blocks == 16u && cuda_q4_mma_tile16_shmem_ok(0);
if (use_q4_mma_t16 && use_gate_row2048) {
const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert);
const size_t t16sh = 16u * 16u * sizeof(cuda_block_q8_K);
if (gate_row_span == 512u) {
dim3 tgrid((expert_mid_dim + 511u) / 512u, t16cap, 1);
moe_gate_up_mid_q4K_tile16_mma_kernel<512><<<tgrid, 256, t16sh, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else if (gate_row_span == 1024u) {
dim3 tgrid((expert_mid_dim + 1023u) / 1024u, t16cap, 1);
moe_gate_up_mid_q4K_tile16_mma_kernel<1024><<<tgrid, 256, t16sh, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else {
dim3 tgrid((expert_mid_dim + 2047u) / 2048u, t16cap, 1);
moe_gate_up_mid_q4K_tile16_mma_kernel<2048><<<tgrid, 256, t16sh, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else if (use_q4_mma && use_gate_row2048) {
if (gate_row_span == 512u) {
dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1);
moe_gate_up_mid_q4K_tile8_mma_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else if (gate_row_span == 1024u) {
dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1);
moe_gate_up_mid_q4K_tile8_mma_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else {
dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1);
moe_gate_up_mid_q4K_tile8_mma_kernel<2048><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else if (use_gate_row2048) {
if (gate_row_span == 512u) {
dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1);
moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else if (gate_row_span == 1024u) {
dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1);
moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else {
dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1);
moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<2048><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else {
dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1);
moe_gate_up_mid_q4K_expert_tile8_rowspan_kernel<32><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else if (use_gate_row2048) {
if (gate_row_span == 512u) {
dim3 tgrid((expert_mid_dim + 511u) / 512u, tile_capacity, 1);
moe_gate_up_mid_expert_tile8_rowspan_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else if (gate_row_span == 1024u) {
dim3 tgrid((expert_mid_dim + 1023u) / 1024u, tile_capacity, 1);
moe_gate_up_mid_expert_tile8_rowspan_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else {
dim3 tgrid((expert_mid_dim + 2047u) / 2048u, tile_capacity, 1);
moe_gate_up_mid_expert_tile8_row2048_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else if (expert_tile_m == 8u) {
dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1);
moe_gate_up_mid_expert_tile8_row32_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
} else {
dim3 tgrid((expert_mid_dim + 31u) / 32u, tile_capacity, 1);
moe_gate_up_mid_expert_tile4_row32_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr,
gate_w, up_w, xq, sorted_pairs, sorted_offsets, sorted_counts,
tile_total, tile_experts, tile_starts, (const float *)weights->ptr,
gate_expert_bytes, gate_row_bytes, xq_blocks, expert_mid_dim, n_expert,
write_gate_up, clamp);
}
} else if (ok && sorted_pairs && use_p2_sorted) {
dim3 p2_mgrid((expert_mid_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1);
moe_gate_up_mid_sorted_p2_qwarp32_kernel<<<p2_mgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
sorted_pairs,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
pair_count,
clamp);
} else if (ok && sorted_pairs) {
moe_gate_up_mid_sorted_qwarp32_kernel<<<mgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
sorted_pairs,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
clamp);
} else if (ok) {
dim3 qgrid((expert_mid_dim + MOE_DECODE_ROWS_PER_BLOCK - 1u) / MOE_DECODE_ROWS_PER_BLOCK, n_tokens * n_expert, 1);
if (q4k_path) {
/* Q4_K gate/up: the decode kernel is token-indexed via
* pair = blockIdx.y; tok = pair / n_expert, so the same
* launch covers both n_tokens == 1 (decode) and n_tokens > 1
* (prefill). q4k_path is steered here by use_sorted_pairs = 0
* cascading the IQ2 sorted/expert-tile branches off. */
if (use_direct_midq) {
dim3 mqgrid(midq_blocks, n_tokens * n_expert, 1);
moe_gate_up_midq_decode_q4K_qwarp32_kernel<<<mqgrid, 256, 0, cuda_decode_stream()>>>(
(float *)mid->ptr,
midq,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
clamp);
} else if (use_q4_gate_h16r8) {
dim3 h8grid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1);
moe_gate_up_mid_decode_q4K_hwarp16_row8_kernel<<<h8grid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
} else if (use_q4_gate_w32r16) {
dim3 w16grid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1);
moe_gate_up_mid_decode_q4K_warp32_row16_kernel<<<w16grid, 512, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
} else if (use_q4_gate_w32) {
dim3 wgrid((expert_mid_dim + 7u) / 8u, n_tokens * n_expert, 1);
if (use_q4_gate_w32_noaux) {
if (use_q4_midq_sidecar) {
moe_gate_up_mid_decode_q4K_warp32_noaux_sidecar_kernel<<<wgrid, 256, 0, cuda_decode_stream()>>>(
(float *)mid->ptr,
midq_sidecar,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
clamp);
} else {
moe_gate_up_mid_decode_q4K_warp32_noaux_kernel<<<wgrid, 256, 0, cuda_decode_stream()>>>(
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
clamp);
}
} else {
moe_gate_up_mid_decode_q4K_warp32_kernel<<<wgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
}
} else if (use_q4_gate_h16) {
dim3 hgrid((expert_mid_dim + 15u) / 16u, n_tokens * n_expert, 1);
moe_gate_up_mid_decode_q4K_hwarp16_kernel<<<hgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
} else {
moe_gate_up_mid_decode_q4K_qwarp32_kernel<<<qgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
}
} else if (use_decode_lut_gate) {
moe_gate_up_mid_decode_lut_qwarp32_kernel<<<qgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
write_gate_up,
clamp);
} else {
moe_gate_up_mid_qwarp32_kernel<<<qgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
n_expert,
clamp);
}
}
ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch");
}
if (prof_ev[3]) (void)cudaEventRecord(prof_ev[3], 0);
if (ok && !use_direct_midq) {
dim3 midq_grid(midq_blocks, n_tokens * n_expert, 1);
if (use_q4_midq_sidecar) {
q8_K_quantize_sidecar_kernel<<<midq_grid, 256, 0, cuda_decode_stream()>>>(
midq,
(const float *)mid->ptr,
midq_sidecar,
expert_mid_dim,
n_tokens * n_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe mid sidecar quantize launch");
} else if (use_owned_sparse_buffers) {
q8_K_quantize_owned_kernel<<<midq_grid, 256, 0, cuda_decode_stream()>>>(
midq,
(const float *)mid->ptr,
(const int32_t *)selected->ptr,
expert_mid_dim,
n_tokens * n_expert,
0u,
n_total_expert);
ok = cuda_ok(cudaGetLastError(),
"owned routed_moe active mid quantize launch");
} else {
q8_K_quantize_kernel<<<midq_grid, 256, 0, cuda_decode_stream()>>>(midq, (const float *)mid->ptr, expert_mid_dim, n_tokens * n_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe mid quantize launch");
}
}
if (prof_ev[4]) (void)cudaEventRecord(prof_ev[4], 0);
if (ok && owned_filtered && use_sorted_pairs && !use_atomic_down &&
!use_owned_sparse_buffers) {
const uint64_t down_clear_bytes =
(uint64_t)n_tokens * n_expert * out_dim * sizeof(float);
ok = cuda_ok(cudaMemset(down->ptr, 0, (size_t)down_clear_bytes),
"owned routed_moe down clear");
}
if (ok) {
dim3 dgrid((out_dim + 31u) / 32u, n_tokens * n_expert, 1);
uint32_t *down_tile_total = tile_total;
uint32_t *down_tile_experts = tile_experts;
uint32_t *down_tile_starts = tile_starts;
uint32_t down_tile_capacity = tile_capacity;
if (use_down_tile16 && tile16_total && tile16_experts && tile16_starts) {
down_tile_total = tile16_total;
down_tile_experts = tile16_experts;
down_tile_starts = tile16_starts;
down_tile_capacity = tile16_capacity;
}
if (use_direct_down_sum) {
dim3 sgrid((out_dim + 31u) / 32u, n_tokens, 1);
if (q4k_path) {
if (n_expert == 6u) {
moe_down_q4K_sum6_qwarp32_kernel<<<sgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim);
} else {
if (use_q4_down_slot3) {
dim3 swgrid((out_dim + 7u) / 8u, 1, 1);
moe_down_q4K_sum3_slotwarp_kernel<<<swgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim);
} else {
moe_down_q4K_sum3_qwarp32_kernel<<<sgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim);
}
}
} else {
if (n_expert == 6u) {
moe_down_sum_qwarp32_kernel<6><<<sgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim);
} else {
moe_down_sum_qwarp32_kernel<3><<<sgrid, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim);
}
}
} else if (use_atomic_down) {
uint64_t n = (uint64_t)n_tokens * out_dim;
zero_kernel<<<(n + 255u) / 256u, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, n);
ok = cuda_ok(cudaGetLastError(), "routed_moe atomic zero launch");
}
if (use_direct_down_sum) {
/* The direct decode kernel writes the final token row. */
} else if (sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts &&
down_tile_total && down_tile_experts && down_tile_starts) {
if (q4k_path) {
const int use_q4_down_mma = cuda_q4_mma_ok() &&
((((uintptr_t)down_w | down_row_bytes | down_expert_bytes) & 15u) == 0u) &&
midq_blocks <= 8u && (out_dim & 7u) == 0u;
const int use_q4_down_t16 = use_q4_down_mma && use_q4_mma_tiles16 &&
tile16_total && tile16_experts && tile16_starts &&
midq_blocks <= 16u && cuda_q4_mma_tile16_shmem_ok(1);
if (use_q4_down_t16 && use_q4_down_rowspan) {
const unsigned t16cap = (unsigned)((pair_count + 15u) / 16u + n_total_expert);
const size_t dt16sh = 16u * (size_t)midq_blocks * sizeof(cuda_block_q8_K);
if (down_row_span == 512u) {
dim3 tgrid((out_dim + 511u) / 512u, t16cap, 1);
moe_down_q4K_tile16_mma_kernel<512><<<tgrid, 256, dt16sh, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else if (down_row_span == 1024u) {
dim3 tgrid((out_dim + 1023u) / 1024u, t16cap, 1);
moe_down_q4K_tile16_mma_kernel<1024><<<tgrid, 256, dt16sh, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else {
dim3 tgrid((out_dim + 2047u) / 2048u, t16cap, 1);
moe_down_q4K_tile16_mma_kernel<2048><<<tgrid, 256, dt16sh, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
tile16_total, tile16_experts, tile16_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
}
} else if (use_q4_down_mma && use_q4_down_rowspan) {
if (down_row_span == 512u) {
dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1);
moe_down_q4K_tile8_mma_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else if (down_row_span == 1024u) {
dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1);
moe_down_q4K_tile8_mma_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else {
dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1);
moe_down_q4K_tile8_mma_kernel<2048><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
}
} else if (use_q4_down_rowspan) {
if (down_row_span == 512u) {
dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1);
moe_down_q4K_expert_tile8_rowspan_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else if (down_row_span == 1024u) {
dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1);
moe_down_q4K_expert_tile8_rowspan_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
} else {
dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1);
moe_down_q4K_expert_tile8_rowspan_kernel<2048><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
}
} else {
dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1);
moe_down_q4K_expert_tile8_rowspan_kernel<32><<<tgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert);
}
} else if (use_down_row2048) {
if (down_row_span == 512u) {
dim3 tgrid((out_dim + 511u) / 512u, down_tile_capacity, 1);
moe_down_expert_tile16_rowspan_kernel<512><<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
} else if (down_row_span == 1024u) {
dim3 tgrid((out_dim + 1023u) / 1024u, down_tile_capacity, 1);
moe_down_expert_tile16_rowspan_kernel<1024><<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
} else {
dim3 tgrid((out_dim + 2047u) / 2048u, down_tile_capacity, 1);
moe_down_expert_tile16_row2048_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
}
} else if (use_down_tile16) {
dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1);
moe_down_expert_tile16_row32_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
} else if (expert_tile_m == 8u) {
dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1);
moe_down_expert_tile8_row32_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
} else {
dim3 tgrid((out_dim + 31u) / 32u, down_tile_capacity, 1);
moe_down_expert_tile4_row32_kernel<<<tgrid, 256, 0, cuda_decode_stream()>>>(
use_atomic_down ? (float *)out->ptr : (float *)down->ptr,
down_w, midq, sorted_pairs, sorted_offsets, sorted_counts,
down_tile_total, down_tile_experts, down_tile_starts, down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, use_atomic_down);
}
} else if (sorted_pairs && use_p2_sorted) {
dim3 p2_dgrid((out_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1);
moe_down_sorted_p2_qwarp32_kernel<<<p2_dgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w,
midq,
sorted_pairs,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
n_expert,
pair_count);
} else if (sorted_pairs) {
moe_down_sorted_qwarp32_kernel<<<dgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w,
midq,
sorted_pairs,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
n_expert);
} else if (q4k_path) {
/* Q4_K prefill down. New kernel mirrors moe_down_qwarp32_kernel
* grid/geometry, swapping the weight block type to cuda_block_q4_K
* and the dot helper to dev_dot_q4_K_q8_K_block. Writes per-pair
* outputs into down->ptr; moe_sum_kernel below sums them across
* experts into out->ptr. */
moe_down_q4K_qwarp32_kernel<<<dgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
n_expert);
} else {
moe_down_qwarp32_kernel<<<dgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
n_expert);
}
ok = cuda_ok(cudaGetLastError(), "routed_moe down launch");
}
if (prof_ev[5]) (void)cudaEventRecord(prof_ev[5], 0);
if (ok && !use_atomic_down && !use_direct_down_sum) {
uint64_t n = (uint64_t)n_tokens * out_dim;
if (use_owned_sparse_buffers) {
moe_sum_owned_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
(const float *)down->ptr,
(const int32_t *)selected->ptr,
out_dim,
n_expert,
n_tokens);
} else {
moe_sum_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
(const float *)down->ptr,
out_dim,
n_expert,
n_tokens);
}
ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch");
}
if (prof_ev[6]) {
(void)cudaEventRecord(prof_ev[6], 0);
if (cudaEventSynchronize(prof_ev[6]) == cudaSuccess) {
float ms_xq = 0.0f, ms_sort = 0.0f, ms_gate = 0.0f, ms_midq = 0.0f, ms_down = 0.0f, ms_sum = 0.0f, ms_total = 0.0f;
(void)cudaEventElapsedTime(&ms_xq, prof_ev[0], prof_ev[1]);
(void)cudaEventElapsedTime(&ms_sort, prof_ev[1], prof_ev[2]);
(void)cudaEventElapsedTime(&ms_gate, prof_ev[2], prof_ev[3]);
(void)cudaEventElapsedTime(&ms_midq, prof_ev[3], prof_ev[4]);
(void)cudaEventElapsedTime(&ms_down, prof_ev[4], prof_ev[5]);
(void)cudaEventElapsedTime(&ms_sum, prof_ev[5], prof_ev[6]);
(void)cudaEventElapsedTime(&ms_total, prof_ev[0], prof_ev[6]);
fprintf(stderr,
"ds4: CUDA MoE profile tokens=%u pairs=%u xq=%.3f sort=%.3f gateup=%.3f midq=%.3f down=%.3f sum=%.3f total=%.3f ms\n",
n_tokens, pair_count, ms_xq, ms_sort, ms_gate, ms_midq, ms_down, ms_sum, ms_total);
}
for (uint32_t i = 0; i < 7u; i++) (void)cudaEventDestroy(prof_ev[i]);
}
return ok;
}
if (ok) {
dim3 mgrid(expert_mid_dim, n_tokens * n_expert, 1);
moe_gate_up_mid_f32_kernel<<<mgrid, 256, 0, cuda_decode_stream()>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
(const float *)x->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
expert_in_dim,
expert_mid_dim,
n_expert,
clamp);
ok = cuda_ok(cudaGetLastError(), "routed_moe gate/up launch");
}
if (ok) {
dim3 dgrid(out_dim, n_tokens * n_expert, 1);
moe_down_f32_kernel<<<dgrid, 256, 0, cuda_decode_stream()>>>(
(float *)down->ptr,
down_w,
(const float *)mid->ptr,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
expert_mid_dim,
out_dim,
n_expert);
ok = cuda_ok(cudaGetLastError(), "routed_moe down launch");
}
if (ok) {
uint64_t n = (uint64_t)n_tokens * out_dim;
moe_sum_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out->ptr, (const float *)down->ptr, out_dim, n_expert, n_tokens);
ok = cuda_ok(cudaGetLastError(), "routed_moe sum launch");
}
return ok;
}
extern "C" int ds4_gpu_routed_moe_one_owned_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *down,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t resident_expert_base,
uint32_t resident_expert_count,
float clamp,
const ds4_gpu_tensor *x,
ds4_gpu_tensor *down_output,
bool pack_fixed3,
ds4_gpu_tensor *shared_prequant) {
if (!out || !gate || !up || !mid || !down || !model_map ||
!selected || !weights || !x || n_expert != 6u ||
n_total_expert == 0u || resident_expert_count == 0u ||
gate_expert_bytes == 0u || gate_row_bytes == 0u ||
down_expert_bytes == 0u || down_row_bytes == 0u ||
resident_expert_base >= n_total_expert ||
resident_expert_count > n_total_expert - resident_expert_base ||
expert_in_dim % CUDA_QK_K != 0u ||
expert_mid_dim % CUDA_QK_K != 0u ||
selected->bytes < 6u * sizeof(int32_t) ||
weights->bytes < 6u * sizeof(float) ||
x->bytes < (uint64_t)expert_in_dim * sizeof(float) ||
mid->bytes < 6ull * expert_mid_dim * sizeof(float) ||
out->bytes < (uint64_t)out_dim * sizeof(float)) {
return 0;
}
if (pack_fixed3 && resident_expert_base == 0u) return 0;
const bool q4k_path = gate_type == 12u && down_type == 12u;
const bool mxfp4_path = gate_type == 39u && down_type == 39u;
if (!q4k_path && !mxfp4_path &&
(gate_type != 16u || down_type != 10u)) {
return 0;
}
if (q4k_path && getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL) {
fprintf(stderr, "ds4: CUDA owned Q4 decode does not support gate/up auxiliary output\n");
return 0;
}
if (!q4k_path && !mxfp4_path &&
getenv("DS4_CUDA_MOE_NO_DECODE_LUT_GATE") != NULL) {
fprintf(stderr, "ds4: CUDA owned IQ2 decode requires the LUT gate path\n");
return 0;
}
const bool write_aux =
!q4k_path && !mxfp4_path &&
getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL;
if (resident_expert_base > UINT64_MAX / gate_expert_bytes ||
resident_expert_count > UINT64_MAX / gate_expert_bytes ||
resident_expert_base > UINT64_MAX / down_expert_bytes ||
resident_expert_count > UINT64_MAX / down_expert_bytes) {
return 0;
}
const uint64_t gate_shift = (uint64_t)resident_expert_base * gate_expert_bytes;
const uint64_t down_shift = (uint64_t)resident_expert_base * down_expert_bytes;
const uint64_t gate_bytes = (uint64_t)resident_expert_count * gate_expert_bytes;
const uint64_t down_bytes = (uint64_t)resident_expert_count * down_expert_bytes;
if (gate_offset > model_size || gate_shift > model_size - gate_offset ||
gate_bytes > model_size - gate_offset - gate_shift ||
up_offset > model_size || gate_shift > model_size - up_offset ||
gate_bytes > model_size - up_offset - gate_shift ||
down_offset > model_size || down_shift > model_size - down_offset ||
down_bytes > model_size - down_offset - down_shift) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const char *gate_w = (const char *)cuda_resolve_weight_ptr(
model_map, gate_offset + gate_shift, gate_bytes,
logical_tier, "moe_owned_gate");
const char *up_w = (const char *)cuda_resolve_weight_ptr(
model_map, up_offset + gate_shift, gate_bytes,
logical_tier, "moe_owned_up");
const char *down_w = (const char *)cuda_resolve_weight_ptr(
model_map, down_offset + down_shift, down_bytes,
logical_tier, "moe_owned_down");
if (!gate_w || !up_w || !down_w) return 0;
const uint64_t down_output_bytes =
(uint64_t)(pack_fixed3 ? 4u : 6u) * out_dim * sizeof(float);
if (mxfp4_path) {
const uint64_t slot_bytes = 6ull * out_dim * sizeof(float);
if (!cuda_use_mxfp4_mmq() ||
gate->bytes < 6u * sizeof(int32_t) ||
up->bytes < 6u * sizeof(float) ||
down->bytes < slot_bytes ||
(down_output && down_output->bytes < down_output_bytes)) {
return 0;
}
cudaStream_t stream = cuda_decode_stream();
int32_t *local_ids = (int32_t *)gate->ptr;
float *local_weights = (float *)up->ptr;
mxfp4_prepare_owned_assignments_kernel<<<1, 32, 0, stream>>>(
local_ids,
local_weights,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
resident_expert_base,
resident_expert_count);
if (!cuda_ok(cudaGetLastError(),
"owned MXFP4 assignment prepare launch")) {
return 0;
}
if (ds4_mmq_mxfp4_moe_gate_up_mid_vec(
gate_w,
up_w,
(const float *)x->ptr,
local_ids,
local_weights,
(float *)mid->ptr,
expert_mid_dim,
expert_in_dim,
1,
resident_expert_count,
6,
clamp,
stream) != 0) {
return 0;
}
if (ds4_mmq_mxfp4_moe_vec(
down_w,
(const float *)mid->ptr,
local_ids,
(float *)down->ptr,
out_dim,
expert_mid_dim,
6,
resident_expert_count,
1,
stream) != 0) {
return 0;
}
if (pack_fixed3) {
float *packed_dst = down_output
? (float *)down_output->ptr
: (float *)down->ptr;
moe_down_owned_pack_f32_slots_kernel<<<
(out_dim + 255u) / 256u, 256, 0, stream>>>(
packed_dst,
(const float *)down->ptr,
(const int32_t *)selected->ptr,
out_dim,
resident_expert_base,
resident_expert_count);
if (!cuda_ok(cudaGetLastError(),
"owned MXFP4 down pack launch")) {
return 0;
}
} else if (down_output) {
const uint64_t slot_count = 6ull * out_dim;
moe_down_owned_copy_f32_slots_kernel<<<
(slot_count + 255u) / 256u, 256, 0, stream>>>(
(float *)down_output->ptr,
(const float *)down->ptr,
slot_count);
if (!cuda_ok(cudaGetLastError(),
"owned MXFP4 down peer copy launch")) {
return 0;
}
}
return 1;
}
const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K;
const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K;
const uint64_t xq_bytes = (uint64_t)xq_blocks * sizeof(cuda_block_q8_K);
const uint64_t midq_bytes = 6ull * midq_blocks * sizeof(cuda_block_q8_K);
const uint64_t aux_bytes = 6ull * expert_mid_dim * sizeof(float);
const uint64_t shared_q8_blocks = expert_in_dim / 32u;
const uint64_t shared_q8_bytes = shared_q8_blocks * 32u;
const uint64_t shared_scale_offset =
(shared_q8_bytes + 15u) & ~15ull;
const uint64_t shared_prequant_bytes =
shared_scale_offset + shared_q8_blocks * sizeof(float);
if (down->bytes < xq_bytes || down->bytes < down_output_bytes ||
(down_output && down_output->bytes < down_output_bytes) ||
gate->bytes < midq_bytes ||
(shared_prequant &&
(shared_prequant->bytes < shared_prequant_bytes ||
ds4_tensor_device_idx(shared_prequant) != logical_tier)) ||
(write_aux && (gate->bytes < aux_bytes || up->bytes < aux_bytes))) {
return 0;
}
float *down_dst = (float *)(down_output ? down_output->ptr : down->ptr);
cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr;
cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr;
dim3 xq_grid(xq_blocks, 1, 1);
if (shared_prequant) {
int8_t *shared_xq = (int8_t *)shared_prequant->ptr;
float *shared_scale = (float *)((char *)shared_prequant->ptr +
shared_scale_offset);
q8_K_q8_0_quantize_kernel<<<xq_grid, 256>>>(
xq, shared_xq, shared_scale, (const float *)x->ptr,
expert_in_dim, 1u);
} else {
q8_K_quantize_kernel<<<xq_grid, 256>>>(
xq, (const float *)x->ptr, expert_in_dim, 1u);
}
if (!cuda_ok(cudaGetLastError(), "owned routed_moe x quantize launch")) return 0;
if (q4k_path) {
dim3 gate_grid((expert_mid_dim + 7u) / 8u, 6u, 1u);
moe_gate_up_mid_decode_q4K_owned_warp32_noaux_kernel<<<gate_grid, 256>>>(
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
6u,
resident_expert_base,
resident_expert_count,
clamp);
} else {
dim3 gate_grid((expert_mid_dim + 31u) / 32u, 6u, 1u);
moe_gate_up_mid_decode_lut_owned_qwarp32_kernel<<<gate_grid, 256>>>(
(float *)gate->ptr,
(float *)up->ptr,
(float *)mid->ptr,
gate_w,
up_w,
xq,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
gate_expert_bytes,
gate_row_bytes,
xq_blocks,
expert_mid_dim,
6u,
resident_expert_base,
resident_expert_count,
write_aux,
clamp);
}
if (!cuda_ok(cudaGetLastError(), "owned routed_moe gate/up launch")) return 0;
dim3 midq_grid(midq_blocks, 6u, 1u);
q8_K_quantize_owned_kernel<<<midq_grid, 256>>>(
midq,
(const float *)mid->ptr,
(const int32_t *)selected->ptr,
expert_mid_dim,
6u,
resident_expert_base,
resident_expert_count);
if (!cuda_ok(cudaGetLastError(), "owned routed_moe mid quantize launch")) return 0;
dim3 down_grid((out_dim + 31u) / 32u, pack_fixed3 ? 4u : 6u, 1u);
if (q4k_path && pack_fixed3) {
moe_down_q4K_owned_packed_qwarp32_kernel<<<down_grid, 256>>>(
down_dst,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
resident_expert_base,
resident_expert_count);
} else if (q4k_path) {
moe_down_q4K_owned_slots_qwarp32_kernel<<<down_grid, 256>>>(
down_dst,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
resident_expert_base,
resident_expert_count);
} else if (pack_fixed3) {
moe_down_owned_packed_qwarp32_kernel<<<down_grid, 256>>>(
down_dst,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
resident_expert_base,
resident_expert_count);
} else {
moe_down_owned_slots_qwarp32_kernel<<<down_grid, 256>>>(
down_dst,
down_w,
midq,
(const int32_t *)selected->ptr,
down_expert_bytes,
down_row_bytes,
midq_blocks,
out_dim,
resident_expert_base,
resident_expert_count);
}
return cuda_ok(cudaGetLastError(), "owned routed_moe down launch");
}
extern "C" int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_slots,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split,
uint32_t rows) {
if (!out || !home_slots || !peer_slots || !selected || out_dim == 0u ||
rows == 0u || rows > 65535u) {
return 0;
}
const uint64_t row_elems = (uint64_t)rows * out_dim;
if (row_elems > UINT64_MAX / (6u * sizeof(float))) return 0;
const uint64_t out_bytes = row_elems * sizeof(float);
const uint64_t slots_bytes = row_elems * 6u * sizeof(float);
const uint64_t selected_bytes = (uint64_t)rows * 6u * sizeof(int32_t);
if (out->bytes < out_bytes || home_slots->bytes < slots_bytes ||
peer_slots->bytes < slots_bytes || selected->bytes < selected_bytes) {
return 0;
}
const dim3 grid((out_dim + 255u) / 256u, rows, 1u);
moe_owned_slots_combine_fixed3_kernel<<<grid, 256>>>(
(float *)out->ptr,
(const float *)home_slots->ptr,
(const float *)peer_slots->ptr,
(const int32_t *)selected->ptr,
out_dim,
expert_split);
return cuda_ok(cudaGetLastError(),
"owned routed_moe slot rows combine launch");
}
extern "C" int ds4_gpu_routed_moe_owned_slots_combine_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_slots,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split) {
return ds4_gpu_routed_moe_owned_slots_combine_rows_tensor(
out, home_slots, peer_slots, selected,
out_dim, expert_split, 1u);
}
extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_packed,
const ds4_gpu_tensor *selected,
uint32_t out_dim,
uint32_t expert_split) {
const uint64_t home_bytes = 6ull * out_dim * sizeof(float);
const uint64_t peer_bytes = 4ull * out_dim * sizeof(float);
if (!out || !home_slots || !peer_packed || !selected || out_dim == 0u ||
out->bytes < (uint64_t)out_dim * sizeof(float) ||
home_slots->bytes < home_bytes || peer_packed->bytes < peer_bytes ||
selected->bytes < 6u * sizeof(int32_t)) {
return 0;
}
moe_owned_packed_combine_fixed3_kernel<<<
(out_dim + 255u) / 256u, 256>>>(
(float *)out->ptr,
(const float *)home_slots->ptr,
(const float *)peer_packed->ptr,
(const int32_t *)selected->ptr,
out_dim,
expert_split);
return cuda_ok(cudaGetLastError(),
"owned routed_moe packed combine launch");
}
extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x,
const ds4_gpu_tensor *add_in,
uint32_t layer_index,
bool force_resident) {
if (add_in) {
if (!ds4_gpu_add_tensor(out, out, add_in,
(uint32_t)(out->bytes / sizeof(float)))) return 0;
}
return routed_moe_launch(out, gate, up, mid, down, model_map, model_size,
gate_offset, up_offset, down_offset,
gate_type, down_type,
gate_expert_bytes, gate_row_bytes,
down_expert_bytes, down_row_bytes,
expert_in_dim, expert_mid_dim, out_dim,
selected, weights, n_total_expert, n_expert, clamp, x,
layer_index, 1, force_resident ? 0 : 1, 0);
}
extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) {
(void)force_resident;
if (mid_is_f16) *mid_is_f16 = false;
return routed_moe_launch(out, gate, up, mid, down, model_map, model_size,
gate_offset, up_offset, down_offset,
gate_type, down_type,
gate_expert_bytes, gate_row_bytes,
down_expert_bytes, down_row_bytes,
expert_in_dim, expert_mid_dim, out_dim,
selected, weights, n_total_expert, n_expert, clamp, x,
layer_index, n_tokens, 1, 0);
}
extern "C" int ds4_gpu_routed_moe_batch_owned_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *down,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t resident_expert_base,
uint32_t resident_expert_count,
float clamp,
const ds4_gpu_tensor *x,
uint32_t layer_index,
uint32_t n_tokens,
bool *mid_is_f16) {
if (mid_is_f16) *mid_is_f16 = false;
if (!selected || !weights || n_tokens == 0u || n_expert == 0u ||
n_total_expert == 0u || resident_expert_count == 0u ||
gate_expert_bytes == 0u || gate_row_bytes == 0u ||
down_expert_bytes == 0u || down_row_bytes == 0u ||
resident_expert_base >= n_total_expert ||
resident_expert_count > n_total_expert - resident_expert_base ||
resident_expert_base > UINT64_MAX / gate_expert_bytes ||
resident_expert_base > UINT64_MAX / down_expert_bytes) {
return 0;
}
const uint64_t pair_count = (uint64_t)n_tokens * n_expert;
if (pair_count > UINT32_MAX ||
selected->bytes < pair_count * sizeof(int32_t) ||
weights->bytes < pair_count * sizeof(float)) {
return 0;
}
const uint64_t gate_shift =
(uint64_t)resident_expert_base * gate_expert_bytes;
const uint64_t down_shift =
(uint64_t)resident_expert_base * down_expert_bytes;
if (gate_offset > model_size || gate_shift > model_size - gate_offset ||
up_offset > model_size || gate_shift > model_size - up_offset ||
down_offset > model_size || down_shift > model_size - down_offset) {
return 0;
}
moe_filter_owned_pairs_kernel<<<(pair_count + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
(int32_t *)selected->ptr,
(float *)weights->ptr,
pair_count,
n_total_expert,
resident_expert_base,
resident_expert_count);
if (!cuda_ok(cudaGetLastError(), "owned routed_moe pair filter launch")) return 0;
return routed_moe_launch(
out, gate, up, mid, down, model_map, model_size,
gate_offset + gate_shift,
up_offset + gate_shift,
down_offset + down_shift,
gate_type, down_type,
gate_expert_bytes, gate_row_bytes,
down_expert_bytes, down_row_bytes,
expert_in_dim, expert_mid_dim, out_dim,
selected, weights, resident_expert_count, n_expert,
clamp, x, layer_index, n_tokens, 0, 1);
}
extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) {
if (!out || !mix || !model_map || n_hc != 4) return 0;
const uint64_t mix_bytes = 24ull * sizeof(float);
if (scale_offset > model_size || model_size - scale_offset < 3ull * sizeof(float) ||
base_offset > model_size || model_size - base_offset < mix_bytes ||
mix->bytes < mix_bytes || out->bytes < mix_bytes) return 0;
const int logical_tier = ds4_tensor_device_idx(out);
const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale");
const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base");
if (!scale || !base) return 0;
uint32_t n_rows = (uint32_t)(mix->bytes / mix_bytes);
if (out->bytes / mix_bytes < n_rows) n_rows = (uint32_t)(out->bytes / mix_bytes);
hc_split_sinkhorn_kernel<<<(n_rows + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr, (const float *)mix->ptr,
scale,
base,
n_rows, sinkhorn_iters, eps);
return cuda_ok(cudaGetLastError(), "hc_split_sinkhorn launch");
}
extern "C" int ds4_gpu_hc_weighted_sum_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *weights, uint32_t n_embd, uint32_t n_hc) {
if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float)));
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)weights->ptr,
n_embd, n_hc, n_tokens, n_hc);
return cuda_ok(cudaGetLastError(), "hc_weighted_sum launch");
}
extern "C" int ds4_gpu_hc_weighted_sum_split_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
if (!out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out->bytes / ((uint64_t)n_embd * sizeof(float)));
uint32_t stride = (uint32_t)(2u * n_hc + n_hc * n_hc);
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr,
n_embd, n_hc, n_tokens, stride);
return cuda_ok(cudaGetLastError(), "hc_weighted_sum_split launch");
}
extern "C" int ds4_gpu_hc_split_weighted_sum_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps) {
if (!out || !split || !mix || !residual_hc || !model_map ||
n_embd == 0 || n_hc != 4) {
return 0;
}
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset) {
return 0;
}
uint64_t n_rows = out->bytes / out_row_bytes;
if (mix->bytes < n_rows * mix_bytes ||
split->bytes < n_rows * mix_bytes ||
residual_hc->bytes < n_rows * residual_row_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, 3ull * sizeof(float), logical_tier, "hc_scale");
const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, mix_bytes, logical_tier, "hc_base");
if (!scale || !base) return 0;
hc_split_weighted_sum_fused_kernel<<<(uint32_t)n_rows, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
(float *)split->ptr,
(const float *)mix->ptr,
(const float *)residual_hc->ptr,
scale,
base,
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps);
return cuda_ok(cudaGetLastError(), "hc split weighted sum launch");
}
extern "C" int ds4_gpu_hc_split_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float norm_eps) {
if (getenv("DS4_CUDA_DISABLE_HC_SPLIT_NORM_FUSED") == NULL) {
if (!out || !norm_out || !split || !mix || !residual_hc || !model_map ||
n_embd == 0 || n_hc != 4) {
return 0;
}
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
norm_out->bytes < out->bytes ||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset ||
norm_weight_offset > model_size ||
(uint64_t)n_embd * sizeof(float) > model_size - norm_weight_offset) {
return 0;
}
uint64_t n_rows = out->bytes / out_row_bytes;
if (n_rows == 1) {
if (mix->bytes < n_rows * mix_bytes ||
split->bytes < n_rows * mix_bytes ||
residual_hc->bytes < n_rows * residual_row_bytes) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset,
3ull * sizeof(float), logical_tier, "hc_scale");
const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset,
mix_bytes, logical_tier, "hc_base");
const float *norm_w = (const float *)cuda_resolve_weight_ptr(model_map, norm_weight_offset,
(uint64_t)n_embd * sizeof(float), logical_tier, "hc_norm_weight");
if (!scale || !base || !norm_w) return 0;
hc_split_weighted_sum_norm_fused_kernel<<<(uint32_t)n_rows, 256, 0, cuda_decode_stream()>>>(
(float *)out->ptr,
(float *)norm_out->ptr,
(float *)split->ptr,
(const float *)mix->ptr,
(const float *)residual_hc->ptr,
scale,
base,
norm_w,
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps, norm_eps);
return cuda_ok(cudaGetLastError(), "hc split weighted sum norm launch");
}
}
/* Multi-row fallback: norm EVERY row (rms_norm_weight_tensor is the
* single-row entry and would leave rows 1..n-1 of norm_out untouched). */
if (!out || n_embd == 0) return 0;
return ds4_gpu_hc_split_weighted_sum_tensor(out, split, mix, residual_hc,
model_map, model_size,
scale_offset, base_offset,
n_embd, n_hc,
sinkhorn_iters, eps) &&
ds4_gpu_rms_norm_weight_rows_tensor(
norm_out, out, model_map, model_size,
norm_weight_offset, n_embd,
(uint32_t)(out->bytes /
((uint64_t)n_embd * sizeof(float))),
norm_eps);
}
extern "C" int ds4_gpu_output_hc_weights_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *pre,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_hc,
float eps) {
if (!out || !pre || !model_map || n_hc == 0) return 0;
const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float);
if (row_bytes == 0 || out->bytes < row_bytes || out->bytes % row_bytes != 0 ||
pre->bytes < out->bytes ||
scale_offset > model_size || sizeof(float) > model_size - scale_offset ||
base_offset > model_size || row_bytes > model_size - base_offset) {
return 0;
}
const uint64_t n_tokens = out->bytes / row_bytes;
const int logical_tier = ds4_tensor_device_idx(out);
const float *scale = (const float *)cuda_resolve_weight_ptr(model_map, scale_offset, sizeof(float), logical_tier, "output_hc_scale");
const float *base = (const float *)cuda_resolve_weight_ptr(model_map, base_offset, row_bytes, logical_tier, "output_hc_base");
if (!scale || !base) return 0;
uint64_t n = n_tokens * n_hc;
output_hc_weights_kernel<<<(n + 255) / 256, 256>>>(
(float *)out->ptr,
(const float *)pre->ptr,
scale,
base,
n_hc,
(uint32_t)n_tokens,
eps);
return cuda_ok(cudaGetLastError(), "output hc weights launch");
}
extern "C" int ds4_gpu_hc_expand_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) {
if (!out_hc || !block_out || !residual_hc || !post || !comb || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float)));
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
hc_expand_kernel<<<(n_elem + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out_hc->ptr,
(const float *)block_out->ptr,
(const float *)block_out->ptr,
(const float *)block_out->ptr,
(const float *)residual_hc->ptr,
(const float *)post->ptr,
(const float *)comb->ptr,
n_embd, n_hc, n_tokens,
n_hc, n_hc * n_hc, 0, 0);
return cuda_ok(cudaGetLastError(), "hc_expand launch");
}
extern "C" int ds4_gpu_hc_expand_add_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) {
if (!out_hc || !block_out || !block_add || !residual_hc || !post || !comb ||
n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float)));
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
hc_expand_kernel<<<(n_elem + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out_hc->ptr,
(const float *)block_out->ptr,
(const float *)block_add->ptr,
(const float *)block_out->ptr,
(const float *)residual_hc->ptr,
(const float *)post->ptr,
(const float *)comb->ptr,
n_embd, n_hc, n_tokens,
n_hc, n_hc * n_hc, 1, 0);
return cuda_ok(cudaGetLastError(), "hc_expand_add launch");
}
extern "C" int ds4_gpu_hc_expand_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float)));
uint32_t mix_hc = 2u * n_hc + n_hc * n_hc;
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
const float *base = (const float *)split->ptr;
hc_expand_kernel<<<(n_elem + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out_hc->ptr,
(const float *)block_out->ptr,
(const float *)block_out->ptr,
(const float *)block_out->ptr,
(const float *)residual_hc->ptr,
base + n_hc,
base + 2u * n_hc,
n_embd, n_hc, n_tokens,
mix_hc, mix_hc, 0, 0);
return cuda_ok(cudaGetLastError(), "hc_expand_split launch");
}
extern "C" int ds4_gpu_hc_expand_add_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float)));
uint32_t mix_hc = 2u * n_hc + n_hc * n_hc;
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
const float *base = (const float *)split->ptr;
hc_expand_kernel<<<(n_elem + 255) / 256, 256, 0, cuda_decode_stream()>>>((float *)out_hc->ptr,
(const float *)block_out->ptr,
(const float *)block_add->ptr,
(const float *)block_out->ptr,
(const float *)residual_hc->ptr,
base + n_hc,
base + 2u * n_hc,
n_embd, n_hc, n_tokens,
mix_hc, mix_hc, 1, 0);
return cuda_ok(cudaGetLastError(), "hc_expand_add_split launch");
}
extern "C" int ds4_gpu_hc_expand_add2_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *block_add2, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
if (!out_hc || !block_out || !block_add || !block_add2 ||
!residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;
uint32_t n_tokens = (uint32_t)(out_hc->bytes / ((uint64_t)n_hc * n_embd * sizeof(float)));
uint32_t mix_hc = 2u * n_hc + n_hc * n_hc;
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
const float *base = (const float *)split->ptr;
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
(const float *)block_out->ptr,
(const float *)block_add->ptr,
(const float *)block_add2->ptr,
(const float *)residual_hc->ptr,
base + n_hc,
base + 2u * n_hc,
n_embd, n_hc, n_tokens,
mix_hc, mix_hc, 1, 1);
return cuda_ok(cudaGetLastError(), "hc_expand_add2_split launch");
}
extern "C" int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *routed_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) {
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out,
model_map, model_size,
weight_offset,
in_dim, out_dim,
shared_mid,
routed_out,
NULL,
NULL, NULL, NULL, 0,
residual_hc,
split,
n_embd, n_hc,
"shared_down_hc_expand");
}
return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size,
weight_offset, in_dim, out_dim,
shared_mid, 1) &&
ds4_gpu_hc_expand_add_split_tensor(out_hc, shared_out, routed_out,
residual_hc, split, n_embd, n_hc);
}
extern "C" int ds4_gpu_shared_down_hc_expand_add_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *routed_out,
const ds4_gpu_tensor *routed_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) {
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out,
model_map, model_size,
weight_offset,
in_dim, out_dim,
shared_mid,
routed_out,
routed_add,
NULL, NULL, NULL, 0,
residual_hc,
split,
n_embd, n_hc,
"shared_down_hc_expand_add");
}
return ds4_gpu_matmul_q8_0_tensor(shared_out, model_map, model_size,
weight_offset, in_dim, out_dim,
shared_mid, 1) &&
ds4_gpu_hc_expand_add2_split_tensor(out_hc, shared_out, routed_out,
routed_add, residual_hc, split,
n_embd, n_hc);
}
extern "C" int ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *home_slots,
const ds4_gpu_tensor *peer_packed,
const ds4_gpu_tensor *selected,
uint32_t expert_split,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") != NULL) return 0;
return cuda_matmul_q8_0_hc_expand_tensor_labeled(
out_hc,
shared_out,
model_map,
model_size,
weight_offset,
in_dim,
out_dim,
shared_mid,
NULL,
NULL,
home_slots,
peer_packed,
selected,
expert_split,
residual_hc,
split,
n_embd,
n_hc,
"shared_down_hc_expand_owned");
}
extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (getenv("DS4_CUDA_DISABLE_Q8_HC_EXPAND_FUSED") == NULL) {
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, block_out,
model_map, model_size,
weight_offset,
in_dim, out_dim,
x,
NULL,
NULL,
NULL, NULL, NULL, 0,
residual_hc,
split,
n_embd, n_hc,
"q8_hc_expand");
}
return ds4_gpu_matmul_q8_0_tensor(block_out, model_map, model_size,
weight_offset, in_dim, out_dim, x, 1) &&
ds4_gpu_hc_expand_split_tensor(out_hc, block_out, residual_hc,
split, n_embd, n_hc);
}
/* --gpu-vram auto probe. Defined here (in the .cu unit) so the
* C-side parser (ds4_gpu_args.c) does not need to include
* <cuda_runtime.h>. Returns 0 on success, nonzero on error
* (errbuf populated). See ds4_gpu_args.h.
*
* Side-effect-light: changes cudaSetDevice during probing; callers
* that care about the active device should reset it themselves
* before continuing. (The mgpu init path resets it anyway.) */
extern "C" int ds4_gpu_args_probe_auto_cuda(const int *device_filter,
int filter_len,
ds4_gpu_config *out,
size_t safety_margin_bytes,
char *errbuf,
size_t errbuflen) {
if (!out) {
if (errbuf && errbuflen) snprintf(errbuf, errbuflen, "internal: NULL out");
return 1;
}
int visible = 0;
cudaError_t rc = cudaGetDeviceCount(&visible);
if (rc != cudaSuccess || visible <= 0) {
if (errbuf && errbuflen) {
snprintf(errbuf, errbuflen,
"cudaGetDeviceCount failed: %s",
rc == cudaSuccess ? "no devices" : cudaGetErrorString(rc));
}
return 1;
}
/* Build the device list: either the explicit filter or 0..visible-1. */
int devs[DS4_MAX_GPUS];
int n_dev = 0;
if (device_filter && filter_len > 0) {
if (filter_len > DS4_MAX_GPUS) {
if (errbuf && errbuflen) {
snprintf(errbuf, errbuflen,
"--gpu-devices filter has %d entries (max %d)",
filter_len, DS4_MAX_GPUS);
}
return 1;
}
for (int i = 0; i < filter_len; i++) {
int d = device_filter[i];
if (d < 0 || d >= visible) {
if (errbuf && errbuflen) {
snprintf(errbuf, errbuflen,
"--gpu-devices: device %d not in 0..%d",
d, visible - 1);
}
return 1;
}
devs[n_dev++] = d;
}
} else {
int cap = visible < DS4_MAX_GPUS ? visible : DS4_MAX_GPUS;
for (int i = 0; i < cap; i++) devs[n_dev++] = i;
}
out->n_gpus = n_dev;
out->safety_margin_bytes = safety_margin_bytes;
for (int i = 0; i < n_dev; i++) {
int d = devs[i];
rc = cudaSetDevice(d);
if (rc != cudaSuccess) {
if (errbuf && errbuflen) {
snprintf(errbuf, errbuflen,
"cudaSetDevice(%d) failed: %s",
d, cudaGetErrorString(rc));
}
return 1;
}
size_t free_b = 0, total_b = 0;
rc = cudaMemGetInfo(&free_b, &total_b);
if (rc != cudaSuccess) {
if (errbuf && errbuflen) {
snprintf(errbuf, errbuflen,
"cudaMemGetInfo on device %d failed: %s",
d, cudaGetErrorString(rc));
}
return 1;
}
/* Auto-mode reserve. Auto-probe is the only place we override
* the user's stated budget, so this is where the conservative-
* on-the-user's-behalf reserve belongs. The
* engine path (engine_classify_multi_tier) still subtracts the
* user-supplied safety_margin_bytes + the cuBLAS workspace from
* whatever budget we hand back; that math is unchanged and
* applies on top of the reserve we trim here.
*
* Reserve = max(2 GiB, 5 % of free). Why these numbers:
* - 2 GiB floor covers runtime scratch / Q8 dequant caches /
* MTP optional state on small GPUs (8-12 GB cards) where
* 5 % is < 1 GiB and not enough headroom.
* - 5 % of free scales the reserve up on larger cards where
* workspace + KV growth needs proportionally more room.
* Explicit --gpu-vram 47,37 budgets do not go through this
* probe and are unaffected. */
const size_t reserve_floor = (size_t)2ull * 1024ull * 1024ull * 1024ull;
const size_t reserve_pct = free_b / 20u;
const size_t reserve = reserve_floor > reserve_pct ? reserve_floor : reserve_pct;
const size_t budget = free_b > reserve ? (free_b - reserve) : 0;
(void)safety_margin_bytes;
out->device_indices[i] = d;
out->vram_bytes[i] = budget;
}
return 0;
}
static int cuda_stream_selected_ensure_bytes(
char **ptr, uint64_t *capacity, uint64_t bytes, const char *label) {
if (*ptr && *capacity >= bytes) return 1;
if (*ptr) {
(void)cudaFree(*ptr);
*ptr = NULL;
*capacity = 0;
}
if (bytes == 0 || bytes > (uint64_t)SIZE_MAX) return 0;
cudaError_t err = cudaMalloc((void **)ptr, (size_t)bytes);
if (err != cudaSuccess) {
fprintf(stderr, "ds4: CUDA streaming %s allocation failed for %.2f MiB: %s\n",
label, (double)bytes / 1048576.0, cudaGetErrorString(err));
(void)cudaGetLastError();
return 0;
}
*capacity = bytes;
return 1;
}
static int cuda_stream_selected_ensure_i32(uint64_t count) {
if (count == 0 || count > UINT64_MAX / sizeof(int32_t)) return 0;
const uint64_t bytes = count * sizeof(int32_t);
return cuda_stream_selected_ensure_bytes(
(char **)&g_stream_selected_cache.slot_selected_ptr,
&g_stream_selected_cache.slot_selected_capacity,
bytes,
"selected-id remap");
}
/* MMQ schedules a rectangular expert grid. A global decode cache can contain
* thousands of mostly inactive slots, overflowing that grid. Stage at most one
* layer, byte-for-byte, and leave the persistent cache and its IDs untouched. */
static int cuda_stream_compact_prefill(const char **gate, const char **up,
const char **down, const int32_t **ids,
uint32_t *experts, uint32_t in_dim,
uint32_t mid_dim, uint32_t out_dim) {
auto &cache = g_stream_selected_cache;
const uint64_t count = g_stream_prefill_slots.size();
if (!cache.valid || !count || g_stream_prefill_ids.size() != cache.slot_count)
return 0;
const uint64_t gate_bytes = in_dim ? ds4_mmq_iq2_xxs_aligned_bytes(mid_dim, in_dim, count) :
count * cache.gate_expert_bytes;
const uint64_t down_bytes = in_dim ? ds4_mmq_q2_k_aligned_bytes(out_dim, mid_dim, count) :
count * cache.down_expert_bytes;
if (!gate_bytes || !down_bytes) return 0;
const uint64_t weights_bytes = 2 * gate_bytes + down_bytes;
const uint64_t ids_offset = (weights_bytes + 3u) & ~UINT64_C(3);
const uint64_t slots_offset = ids_offset + (uint64_t)cache.slot_count * sizeof(int32_t);
const uint64_t bytes = slots_offset + (in_dim ? count * sizeof(int32_t) : 0);
if (!cuda_stream_selected_ensure_bytes(&cache.prefill_ptr,
&cache.prefill_capacity, bytes, "compact prefill experts")) return 0;
char *g = cache.prefill_ptr, *u = g + gate_bytes, *d = u + gate_bytes;
int32_t *compact_ids = (int32_t *)(cache.prefill_ptr + ids_offset);
if (!cuda_ok(cudaMemcpy(compact_ids, g_stream_prefill_ids.data(),
(size_t)cache.slot_count * sizeof(int32_t), cudaMemcpyHostToDevice),
"compact prefill IDs")) return 0;
if (in_dim) {
int32_t *slots = (int32_t *)(cache.prefill_ptr + slots_offset);
if (!cuda_ok(cudaMemcpy(slots, g_stream_prefill_slots.data(), count * sizeof(int32_t),
cudaMemcpyHostToDevice), "aligned prefill slots") ||
!ds4_repack_selected_experts(g, cache.gate_ptr, slots, DS4_REPACK_IQ2_XXS_ALIGNED_MOE,
mid_dim, in_dim, count) ||
!ds4_repack_selected_experts(u, cache.up_ptr, slots, DS4_REPACK_IQ2_XXS_ALIGNED_MOE,
mid_dim, in_dim, count) ||
!ds4_repack_selected_experts(d, cache.down_ptr, slots, DS4_REPACK_Q2_K_ALIGNED_MOE,
out_dim, mid_dim, count)) return 0;
} else for (uint64_t i = 0; i < count; i++) {
const uint64_t slot = (uint32_t)g_stream_prefill_slots[i];
if (!cuda_ok(cudaMemcpyAsync(g + i * cache.gate_expert_bytes,
cache.gate_ptr + slot * cache.gate_expert_bytes,
cache.gate_expert_bytes, cudaMemcpyDeviceToDevice, 0), "compact gate") ||
!cuda_ok(cudaMemcpyAsync(u + i * cache.gate_expert_bytes,
cache.up_ptr + slot * cache.gate_expert_bytes,
cache.gate_expert_bytes, cudaMemcpyDeviceToDevice, 0), "compact up") ||
!cuda_ok(cudaMemcpyAsync(d + i * cache.down_expert_bytes,
cache.down_ptr + slot * cache.down_expert_bytes,
cache.down_expert_bytes, cudaMemcpyDeviceToDevice, 0), "compact down")) return 0;
}
*gate = g; *up = u; *down = d;
*ids = compact_ids;
*experts = (uint32_t)count;
return 1;
}
static int cuda_stream_selected_ranges_valid(
const ds4_gpu_stream_expert_table *table) {
if (!table || !table->model_map || table->model_size == 0 ||
table->n_total_expert == 0 || table->gate_expert_bytes == 0 ||
table->down_expert_bytes == 0) {
return 0;
}
if ((uint64_t)table->n_total_expert >
UINT64_MAX / table->gate_expert_bytes ||
(uint64_t)table->n_total_expert >
UINT64_MAX / table->down_expert_bytes) {
return 0;
}
const uint64_t gate_bytes =
(uint64_t)table->n_total_expert * table->gate_expert_bytes;
const uint64_t down_bytes =
(uint64_t)table->n_total_expert * table->down_expert_bytes;
return table->gate_offset <= table->model_size &&
gate_bytes <= table->model_size - table->gate_offset &&
table->up_offset <= table->model_size &&
gate_bytes <= table->model_size - table->up_offset &&
table->down_offset <= table->model_size &&
down_bytes <= table->model_size - table->down_offset;
}
/* The staging ring spans all misses in a request. Drain it before publishing
* the request, or before any error return can release/reuse its destinations. */
struct cuda_stream_upload_batch {
uint64_t chunks = 0;
bool active = false;
int finish() {
if (!active) return 1;
active = false;
if (!g_stream_selected_upload_stream || cuda_ok(
cudaStreamSynchronize(g_stream_selected_upload_stream), "stream expert batch upload"))
return 1;
/* A failed asynchronous copy must not turn into a cache hit later. */
ds4_gpu_stream_expert_cache_prefetch_finish(true);
g_stream_expert_by_gate.clear();
for (auto &slot : g_stream_expert_slots) slot.used = 0;
return 0;
}
~cuda_stream_upload_batch() { (void)finish(); }
};
static int cuda_stream_selected_cache_begin_load(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t slot_count) {
cuda_stream_prefetch_before_load(table);
cuda_stream_selected_cache_invalidate();
if (!g_ssd_streaming_mode) return 1;
if (!cuda_stream_selected_ranges_valid(table) || !selected_ids || !slot_count)
return 0;
if (g_n_gpus != 1) {
fprintf(stderr, "ds4: CUDA SSD streaming requires single-GPU placement\n");
return 0;
}
if (!cuda_ok(cudaStreamSynchronize(cuda_decode_stream()), "stream expert reuse wait"))
return 0;
try {
std::vector<int32_t> expert_to_slot(table->n_total_expert, -1);
std::vector<int32_t> unique, remap(slot_count);
for (uint32_t i = 0; i < slot_count; i++) {
const int32_t expert = selected_ids[i];
if (expert < 0 || (uint32_t)expert >= table->n_total_expert) {
fprintf(stderr, "ds4: CUDA streaming expert id %d is outside 0..%u at layer %u\n",
expert, table->n_total_expert, table->layer);
return 0;
}
if (expert_to_slot[expert] < 0) {
expert_to_slot[expert] = (int32_t)unique.size();
unique.push_back(expert);
}
remap[i] = expert_to_slot[expert];
}
auto &cache = g_stream_selected_cache;
if (cache.model_map != table->model_map ||
cache.gate_expert_bytes != table->gate_expert_bytes ||
cache.down_expert_bytes != table->down_expert_bytes ||
g_stream_expert_slots.size() < unique.size()) {
cuda_stream_selected_cache_release();
if (ds4_gpu_set_current_device(0) != 0) return 0;
if (table->gate_expert_bytes > (UINT64_MAX - table->down_expert_bytes) / 2u)
return 0;
const uint64_t expert_bytes = 2u * table->gate_expert_bytes + table->down_expert_bytes;
uint64_t capacity = ds4_gpu_stream_expert_cache_budget_for_expert_size(
table->gate_expert_bytes, table->down_expert_bytes);
if (capacity < unique.size()) capacity = unique.size();
size_t free_bytes = 0, total_bytes = 0;
if (!cuda_ok(cudaMemGetInfo(&free_bytes, &total_bytes), "stream expert memory budget"))
return 0;
int integrated = 0;
uint64_t host_available = 0;
if (cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated, g_gpu[0].device_id) == cudaSuccess &&
integrated && ds4_linux_nonmovable_memory(&host_available))
free_bytes = (size_t)std::min(host_available, (uint64_t)total_bytes);
const uint64_t reserve = UINT64_C(8) << 30;
const uint64_t available = free_bytes > reserve ? free_bytes - reserve : 0;
capacity = std::min(capacity, available / expert_bytes);
if (capacity < unique.size()) {
fprintf(stderr, "ds4: CUDA SSD cache cannot stage %zu experts with system headroom\n",
unique.size());
return 0;
}
/* A budget is a hint. Back off before publishing any live slots. */
while (capacity >= unique.size()) {
if (cuda_stream_selected_ensure_bytes(&cache.gate_ptr, &cache.gate_capacity,
capacity * table->gate_expert_bytes, "cached gate experts") &&
cuda_stream_selected_ensure_bytes(&cache.up_ptr, &cache.up_capacity,
capacity * table->gate_expert_bytes, "cached up experts") &&
cuda_stream_selected_ensure_bytes(&cache.down_ptr, &cache.down_capacity,
capacity * table->down_expert_bytes, "cached down experts"))
break;
cuda_stream_selected_cache_release();
if (capacity == unique.size()) return 0;
capacity = std::max((uint64_t)unique.size(), capacity * 3u / 4u);
}
g_stream_expert_slots.resize(capacity);
cache.model_map = table->model_map;
cache.gate_expert_bytes = table->gate_expert_bytes;
cache.down_expert_bytes = table->down_expert_bytes;
cache.logical_tier = 0;
fprintf(stderr, "ds4: CUDA SSD expert cache: %llu slots, %.2f GiB\n",
(unsigned long long)capacity,
(double)(capacity * expert_bytes) / 1073741824.0);
}
if (!cuda_stream_selected_ensure_i32(slot_count)) return 0;
if (!g_stream_expert_budget || g_stream_expert_clock == UINT64_MAX) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
g_stream_expert_by_gate.clear();
for (auto &slot : g_stream_expert_slots) slot.used = 0;
g_stream_expert_clock = 1;
}
const uint64_t stamp = ++g_stream_expert_clock;
std::vector<int32_t> slots(unique.size(), -1);
/* Protect every hit first: a miss must not evict a later request's hit. */
for (size_t i = 0; i < unique.size(); i++) {
const uint64_t expert = (uint32_t)unique[i];
const uint64_t gate = table->gate_offset + expert * table->gate_expert_bytes;
const auto found = g_stream_expert_by_gate.find(gate);
if (found == g_stream_expert_by_gate.end()) continue;
auto &slot = g_stream_expert_slots[found->second];
if (slot.up != table->up_offset + expert * table->gate_expert_bytes ||
slot.down != table->down_offset + expert * table->down_expert_bytes) {
slot.used = 0;
g_stream_expert_by_gate.erase(found);
continue;
}
slots[i] = (int32_t)found->second;
slot.used = stamp;
}
cuda_stream_upload_batch uploads;
for (size_t i = 0; i < unique.size(); i++) {
if (slots[i] >= 0) continue;
uint32_t victim = UINT32_MAX;
uint64_t oldest = stamp;
for (uint32_t j = 0; j < g_stream_expert_slots.size(); j++) {
if (!cuda_stream_prefetch_protects(g_stream_expert_slots[j]) &&
g_stream_expert_slots[j].used < oldest) {
oldest = g_stream_expert_slots[j].used;
victim = j;
if (!oldest) break;
}
}
if (victim == UINT32_MAX) return 0;
auto &slot = g_stream_expert_slots[victim];
if (slot.used) g_stream_expert_by_gate.erase(slot.gate);
slot.used = 0;
const uint64_t expert = (uint32_t)unique[i];
const uint64_t gate = table->gate_offset + expert * table->gate_expert_bytes;
const uint64_t up = table->up_offset + expert * table->gate_expert_bytes;
const uint64_t down = table->down_offset + expert * table->down_expert_bytes;
uploads.active = true;
if (!cuda_model_copy_to_device_streamed(
cache.gate_ptr + (uint64_t)victim * table->gate_expert_bytes,
table->model_map, table->model_size, gate, table->gate_expert_bytes, "stream gate", uploads.chunks) ||
!cuda_model_copy_to_device_streamed(
cache.up_ptr + (uint64_t)victim * table->gate_expert_bytes,
table->model_map, table->model_size, up, table->gate_expert_bytes, "stream up", uploads.chunks) ||
!cuda_model_copy_to_device_streamed(
cache.down_ptr + (uint64_t)victim * table->down_expert_bytes,
table->model_map, table->model_size, down, table->down_expert_bytes, "stream down", uploads.chunks))
return 0;
slot = {gate, up, down, stamp};
g_stream_expert_by_gate[gate] = victim;
slots[i] = (int32_t)victim;
}
if (!uploads.finish()) return 0;
g_stream_prefill_ids = remap;
g_stream_prefill_slots = slots;
for (auto &id : remap) id = slots[id];
if (!cuda_ok(cudaMemcpy(cache.slot_selected_ptr, remap.data(),
(size_t)slot_count * sizeof(int32_t), cudaMemcpyHostToDevice),
"stream selected-id remap copy")) return 0;
cache.layer = table->layer;
cache.n_total_expert = table->n_total_expert;
cache.slot_count = slot_count;
cache.compact_count = (uint32_t)g_stream_expert_slots.size();
cache.gate_offset = table->gate_offset;
cache.up_offset = table->up_offset;
cache.down_offset = table->down_offset;
cache.slot_selected_tensor.ptr = cache.slot_selected_ptr;
cache.slot_selected_tensor.bytes = (uint64_t)slot_count * sizeof(int32_t);
cache.slot_selected_tensor.owner = 0;
cache.slot_selected_tensor.device_id = 0;
cache.valid = 1;
return 1;
} catch (...) {
cuda_stream_selected_cache_release();
return 0;
}
}
/* One reader, one next layer, no second expert arena. Only the foreground
* changes the cache index. Reserved slots are invisible until the reader has
* finished every copy, and cannot be reused while it is running. */
struct cuda_stream_prefetch_slot {
uint32_t index;
cuda_stream_expert_slot value;
};
struct cuda_stream_prefetch_copy {
char *destination;
uint64_t offset, bytes;
};
static struct {
pthread_t thread;
bool active = false, ok = false;
std::atomic<bool> cancel{false};
ds4_gpu_stream_expert_table table = {};
std::vector<cuda_stream_prefetch_slot> slots;
std::vector<cuda_stream_prefetch_copy> copies;
int fd = -1, direct_fd = -1, device = 0;
uint64_t align = 1, file_size = 0, bytes = 0;
void *stage_raw[2] = {};
void *stage[2] = {};
uint64_t stage_bytes = 0;
cudaStream_t stream = NULL;
cudaEvent_t ready[2] = {};
double started = 0, elapsed = 0;
} g_stream_prefetch;
static bool cuda_stream_slot_in_table(const cuda_stream_expert_slot &slot,
const ds4_gpu_stream_expert_table &table) {
if (!slot.used || slot.gate < table.gate_offset) return false;
const uint64_t delta = slot.gate - table.gate_offset;
return table.gate_expert_bytes && delta % table.gate_expert_bytes == 0 &&
delta / table.gate_expert_bytes < table.n_total_expert;
}
static bool cuda_stream_prefetch_protects(const cuda_stream_expert_slot &slot) {
return g_stream_prefetch.active && cuda_stream_slot_in_table(slot, g_stream_prefetch.table);
}
static void cuda_stream_prefetch_before_load(const ds4_gpu_stream_expert_table *table) {
auto &p = g_stream_prefetch;
if (!p.active) return;
const bool same = table && table->model_map == p.table.model_map &&
table->model_size == p.table.model_size &&
table->gate_expert_bytes == p.table.gate_expert_bytes &&
table->down_expert_bytes == p.table.down_expert_bytes;
if (!same || table->gate_offset == p.table.gate_offset)
ds4_gpu_stream_expert_cache_prefetch_finish(!same);
}
static void *cuda_stream_prefetch_read(void *) {
auto &p = g_stream_prefetch;
p.ok = cudaSetDevice(p.device) == cudaSuccess;
uint64_t chunk_index = 0;
const uint64_t chunk = UINT64_C(8) << 20;
for (const auto &copy : p.copies) {
for (uint64_t offset = 0; p.ok && offset < copy.bytes; offset += chunk) {
if (p.cancel.load(std::memory_order_relaxed)) { p.ok = false; break; }
const unsigned ring = chunk_index % 2u;
if (chunk_index >= 2 && cudaEventSynchronize(p.ready[ring]) != cudaSuccess) {
p.ok = false; break;
}
const uint64_t bytes = std::min(chunk, copy.bytes - offset);
const char *payload = NULL;
p.ok = cuda_model_stage_read_from(p.fd, &p.direct_fd, p.align, p.file_size,
p.stage[ring], p.stage_bytes, copy.offset + offset, bytes, &payload) != 0;
if (!p.ok || p.cancel.load(std::memory_order_relaxed)) { p.ok = false; break; }
p.ok = cudaMemcpyAsync(copy.destination + offset, payload, bytes,
cudaMemcpyHostToDevice, p.stream) == cudaSuccess &&
cudaEventRecord(p.ready[ring], p.stream) == cudaSuccess;
#if defined(POSIX_FADV_DONTNEED)
if (p.direct_fd < 0)
(void)posix_fadvise(p.fd, copy.offset + offset, bytes, POSIX_FADV_DONTNEED);
#endif
p.bytes += bytes;
chunk_index++;
}
if (!p.ok) break;
}
/* Even an interrupted read may have queued copies into reserved slots. */
if (cudaStreamSynchronize(p.stream) != cudaSuccess) p.ok = false;
p.elapsed = cuda_wall_sec() - p.started;
return NULL;
}
extern "C" void ds4_gpu_stream_expert_cache_prefetch_finish(bool cancel) {
auto &p = g_stream_prefetch;
const double start = cuda_wall_sec();
if (cancel) p.cancel.store(true, std::memory_order_relaxed);
if (p.active) {
if (pthread_join(p.thread, NULL) != 0) {
fprintf(stderr, "ds4: cannot join CUDA SSD reader safely\n");
abort();
}
p.active = false;
}
const bool publish = p.ok && !cancel;
for (const auto &reserved : p.slots) {
auto &slot = g_stream_expert_slots[reserved.index];
slot.used = 0;
if (publish) {
try {
g_stream_expert_by_gate[reserved.value.gate] = reserved.index;
slot = reserved.value;
/* Look-ahead is not evidence of reuse. Only an actual router
* hit promotes these slots, so unused reads cannot displace
* the experts that should stay hot for following decode. */
slot.used = 1;
} catch (...) {
/* A failed optional index insertion remains an ordinary miss. */
}
}
}
if (!p.slots.empty() && getenv("DS4_CUDA_SSD_PREFETCH_PROFILE"))
fprintf(stderr, "ds4: CUDA SSD prefetch layer=%u experts=%zu bytes=%llu "
"read=%.3f wait=%.3f ms published=%u\n", p.table.layer, p.slots.size(),
(unsigned long long)p.bytes, p.elapsed * 1000,
(cuda_wall_sec() - start) * 1000, publish);
p.slots.clear();
p.copies.clear();
p.ok = false;
if (p.fd >= 0) close(p.fd);
if (p.direct_fd >= 0) close(p.direct_fd);
p.fd = p.direct_fd = -1;
if (cancel) {
for (unsigned i = 0; i < 2; i++) {
if (p.ready[i]) (void)cudaEventDestroy(p.ready[i]);
if (p.stage_raw[i]) (void)cudaFreeHost(p.stage_raw[i]);
p.ready[i] = NULL;
p.stage_raw[i] = p.stage[i] = NULL;
}
if (p.stream) (void)cudaStreamDestroy(p.stream);
p.stream = NULL;
p.stage_bytes = 0;
}
}
static void cuda_stream_prefetch_exit(void) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
}
extern "C" int ds4_gpu_stream_expert_cache_prefetch(
const ds4_gpu_stream_expert_table *current,
const ds4_gpu_stream_expert_table *next) {
ds4_gpu_stream_expert_cache_prefetch_finish(false);
auto &p = g_stream_prefetch;
auto &cache = g_stream_selected_cache;
if (!g_ssd_streaming_mode || g_n_gpus != 1 || !g_stream_expert_budget ||
getenv("DS4_CUDA_DISABLE_SSD_PREFETCH") || g_model_fd < 0 ||
!cuda_stream_selected_ranges_valid(current) || !cuda_stream_selected_ranges_valid(next) ||
current->model_map != next->model_map || current->model_map != cache.model_map ||
current->model_map != g_model_fd_host_base || current->model_size != next->model_size ||
current->gate_offset == next->gate_offset ||
current->gate_expert_bytes != next->gate_expert_bytes ||
current->down_expert_bytes != next->down_expert_bytes ||
next->gate_expert_bytes != cache.gate_expert_bytes ||
next->down_expert_bytes != cache.down_expert_bytes ||
g_stream_expert_slots.size() < (uint64_t)current->n_total_expert + next->n_total_expert ||
g_stream_expert_clock >= UINT64_MAX - 2 || g_model_direct_align > (UINT64_C(1) << 20))
return 0;
try {
/* Registered after CUDA initialization and the cache's static objects:
* an early exit must join before either is destroyed. */
static bool exit_registered = false;
if (!exit_registered) {
if (atexit(cuda_stream_prefetch_exit) != 0) return 0;
exit_registered = true;
}
p.table = *next;
p.cancel.store(false, std::memory_order_relaxed);
p.bytes = 0;
p.elapsed = 0;
p.align = g_model_direct_align;
p.file_size = g_model_file_size;
if (cudaGetDevice(&p.device) != cudaSuccess) return 0;
if (!p.stream) {
p.stage_bytes = (UINT64_C(8) << 20) + p.align;
if (cudaStreamCreateWithFlags(&p.stream, cudaStreamNonBlocking) != cudaSuccess)
throw 0;
for (unsigned i = 0; i < 2; i++) {
if (cudaMallocHost(&p.stage_raw[i], p.stage_bytes + p.align) != cudaSuccess ||
cudaEventCreateWithFlags(&p.ready[i], cudaEventDisableTiming) != cudaSuccess)
throw 0;
p.stage[i] = cuda_align_ptr(p.stage_raw[i], p.align);
}
}
p.fd = dup(g_model_fd);
if (p.fd < 0) throw 0;
p.direct_fd = g_model_direct_fd >= 0 ? dup(g_model_direct_fd) : -1;
std::vector<uint32_t> victims;
for (uint32_t i = 0; i < g_stream_expert_slots.size(); i++) {
const auto &slot = g_stream_expert_slots[i];
if (!cuda_stream_slot_in_table(slot, *current) && !cuda_stream_slot_in_table(slot, *next))
victims.push_back(i);
}
std::stable_sort(victims.begin(), victims.end(), [](uint32_t a, uint32_t b) {
return g_stream_expert_slots[a].used < g_stream_expert_slots[b].used;
});
p.slots.reserve(next->n_total_expert);
p.copies.reserve(3u * next->n_total_expert);
for (uint32_t expert = 0; expert < next->n_total_expert; expert++) {
const cuda_stream_expert_slot value = {
next->gate_offset + expert * next->gate_expert_bytes,
next->up_offset + expert * next->gate_expert_bytes,
next->down_offset + expert * next->down_expert_bytes, UINT64_MAX};
const auto hit = g_stream_expert_by_gate.find(value.gate);
if (hit != g_stream_expert_by_gate.end()) {
const auto &slot = g_stream_expert_slots[hit->second];
if (slot.up == value.up && slot.down == value.down) continue;
throw 0;
}
if (p.slots.size() >= victims.size()) throw 0;
const uint32_t index = victims[p.slots.size()];
p.slots.push_back({index, value});
auto &slot = g_stream_expert_slots[index];
if (slot.used) g_stream_expert_by_gate.erase(slot.gate);
slot = value;
}
/* Read in file order and merge adjacent misses when their cache slots
* are contiguous. The private two-buffer ring overlaps reads/copies. */
for (unsigned part = 0; part < 3; part++) {
for (const auto &slot : p.slots) {
const uint64_t bytes = part == 2 ? next->down_expert_bytes : next->gate_expert_bytes;
const uint64_t offset = part == 0 ? slot.value.gate : part == 1 ? slot.value.up : slot.value.down;
char *dst = (part == 0 ? cache.gate_ptr : part == 1 ? cache.up_ptr : cache.down_ptr) + slot.index * bytes;
if (!p.copies.empty() && p.copies.back().offset + p.copies.back().bytes == offset &&
p.copies.back().destination + p.copies.back().bytes == dst)
p.copies.back().bytes += bytes;
else p.copies.push_back({dst, offset, bytes});
}
}
if (p.slots.empty()) {
ds4_gpu_stream_expert_cache_prefetch_finish(false);
return 0;
}
p.started = cuda_wall_sec();
if (pthread_create(&p.thread, NULL, cuda_stream_prefetch_read, NULL) != 0) throw 0;
p.active = true;
return 1;
} catch (...) {
ds4_gpu_stream_expert_cache_prefetch_finish(true);
(void)cudaGetLastError();
return 0;
}
}
__device__ __forceinline__ static float glm_rope_yarn_corr_factor_dev(
int n_dims, int n_ctx_orig, float n_rot, float base) {
return n_dims * logf(n_ctx_orig / (n_rot * 2.0f * (float)M_PI)) /
(2.0f * logf(base));
}
__device__ __forceinline__ static float glm_rope_yarn_ramp_dev(
float low, float high, int i0) {
const float y = (i0 / 2 - low) / fmaxf(0.001f, high - low);
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
}
__device__ __forceinline__ static void glm_rope_yarn_dev(
float theta_extrap, float freq_scale, const float corr_dims[2],
int i0, float ext_factor, float mscale,
float *cos_theta, float *sin_theta) {
float theta_interp = freq_scale * theta_extrap;
float theta = theta_interp;
if (ext_factor != 0.0f) {
float ramp_mix = glm_rope_yarn_ramp_dev(corr_dims[0], corr_dims[1], i0) *
ext_factor;
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
}
*cos_theta = cosf(theta) * mscale;
*sin_theta = sinf(theta) * mscale;
}
static int cuda_current_tier(void) {
int dev = 0;
if (cudaGetDevice(&dev) != cudaSuccess) return 0;
return dev;
}
/* ===== GLM 5.2 stubs (to be implemented; fail loudly) ===== */
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-parameter"
__global__ static void add3_kernel(float *out, const float *a,
const float *b, const float *c,
uint32_t n) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = a[i] + b[i] + c[i];
}
extern "C" int ds4_gpu_add3_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const ds4_gpu_tensor *c,
uint32_t n) {
if (!out || !a || !b || !c || n == 0 ||
out->bytes < (uint64_t)n * sizeof(float) ||
a->bytes < (uint64_t)n * sizeof(float) ||
b->bytes < (uint64_t)n * sizeof(float) ||
c->bytes < (uint64_t)n * sizeof(float)) {
return 0;
}
add3_kernel<<<(n + 255) / 256, 256>>>(
(float *)out->ptr, (const float *)a->ptr,
(const float *)b->ptr, (const float *)c->ptr, n);
return cuda_ok(cudaGetLastError(), "add3 launch");
}
/* Fused decode residual: sum_out = a + b; norm_out = rmsnorm(sum) * w.
* Single row, one block (two-pass over n with a shared reduction). */
__global__ static void glm_add_rms_norm_weight_kernel(
float *norm_out,
float *sum_out,
const float *a,
const float *b,
const float *w,
uint32_t n,
float eps) {
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
__shared__ float sh[32];
float sumsq = 0.0f;
for (uint32_t i = tid; i < n; i += nth) {
const float v = a[i] + b[i];
sum_out[i] = v;
sumsq += v * v;
}
for (int off = 16; off > 0; off >>= 1) {
sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off);
}
if ((tid & 31u) == 0u) sh[tid >> 5] = sumsq;
__syncthreads();
if (tid < 32u) {
sumsq = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f;
for (int off = 16; off > 0; off >>= 1) {
sumsq += __shfl_xor_sync(0xffffffffu, sumsq, off);
}
if (tid == 0u) sh[0] = sumsq;
}
__syncthreads();
const float scale = rsqrtf(sh[0] / (float)n + eps);
for (uint32_t i = tid; i < n; i += nth) {
norm_out[i] = (sum_out[i] * scale) * w[i];
}
}
extern "C" int ds4_gpu_add_rms_norm_weight_tensor(
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *sum_out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps) {
if (!norm_out || !sum_out || !a || !b || !model_map || n == 0 ||
norm_out->bytes < (uint64_t)n * sizeof(float) ||
sum_out->bytes < (uint64_t)n * sizeof(float) ||
a->bytes < (uint64_t)n * sizeof(float) ||
b->bytes < (uint64_t)n * sizeof(float) ||
weight_offset > model_size ||
(uint64_t)n * sizeof(float) > model_size - weight_offset) {
return 0;
}
const int logical_tier = cuda_current_tier();
const float *w = (const float *)cuda_resolve_weight_ptr(
model_map, weight_offset, (uint64_t)n * sizeof(float),
logical_tier, "rms_weight");
if (!w) return 0;
glm_add_rms_norm_weight_kernel<<<1, 1024>>>(
(float *)norm_out->ptr, (float *)sum_out->ptr,
(const float *)a->ptr, (const float *)b->ptr, w, n, eps);
return cuda_ok(cudaGetLastError(), "add rms norm weight");
}
extern "C" bool ds4_gpu_commands_active(void) {
return false;
}
__global__ static void glm_embed_token_q8_0_kernel(
float *out,
const unsigned char *w,
uint32_t token,
uint32_t n_embd) {
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
if (d >= n_embd) return;
const uint64_t row_blocks = n_embd / 32u;
const unsigned char *blk =
w + ((uint64_t)token * row_blocks + (d >> 5)) * 34u;
const float scale = __half2float(*(const __half *)blk);
out[d] = scale * (float)((const int8_t *)(blk + 2))[d & 31u];
}
extern "C" int ds4_gpu_embed_token_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd) {
if (!out || !model_map || n_embd == 0 || (n_embd & 31u) != 0u ||
token >= n_vocab) {
return 0;
}
if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */
fprintf(stderr, "ds4: embed_token_quant: unsupported type %u\n",
weight_type);
return 0;
}
const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u;
if (weight_offset > model_size ||
(uint64_t)n_vocab * row_bytes > model_size - weight_offset ||
out->bytes < (uint64_t)n_embd * sizeof(float)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr(
model_map, weight_offset, (uint64_t)n_vocab * row_bytes,
logical_tier, "glm_token_embd");
if (!w) return 0;
glm_embed_token_q8_0_kernel<<<(n_embd + 255) / 256, 256>>>(
(float *)out->ptr, w, token, n_embd);
return cuda_ok(cudaGetLastError(), "glm embed token launch");
}
__global__ static void glm_embed_tokens_q8_0_kernel(
float *out,
const int32_t *tokens,
const unsigned char *w,
uint32_t n_tokens,
uint32_t n_embd) {
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
uint64_t n = (uint64_t)n_tokens * n_embd;
if (gid >= n) return;
uint32_t t = gid / n_embd;
uint32_t d = gid - (uint64_t)t * n_embd;
int32_t tok = tokens[t];
const uint64_t row_blocks = n_embd / 32u;
const unsigned char *blk = w + ((uint64_t)tok * row_blocks + (d >> 5)) * 34u;
const float scale = __half2float(*(const __half *)blk);
out[gid] = scale * (float)((const int8_t *)(blk + 2))[d & 31u];
}
extern "C" int ds4_gpu_embed_tokens_quant_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (!out || !tokens || !model_map || n_tokens == 0 || n_embd == 0 ||
(n_embd & 31u) != 0u) {
return 0;
}
if (weight_type != 8u) { /* DS4_TENSOR_Q8_0 */
fprintf(stderr, "ds4: embed_tokens_quant: unsupported type %u\n",
weight_type);
return 0;
}
const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u;
if (weight_offset > model_size ||
(uint64_t)n_vocab * row_bytes > model_size - weight_offset ||
out->bytes < (uint64_t)n_tokens * n_embd * sizeof(float) ||
tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const unsigned char *w = (const unsigned char *)cuda_resolve_weight_ptr(
model_map, weight_offset, (uint64_t)n_vocab * row_bytes,
logical_tier, "glm_token_embd");
if (!w) return 0;
uint64_t n = (uint64_t)n_tokens * n_embd;
glm_embed_tokens_q8_0_kernel<<<(n + 255) / 256, 256>>>(
(float *)out->ptr,
(const int32_t *)tokens->ptr,
w, n_tokens, n_embd);
return cuda_ok(cudaGetLastError(), "glm embed tokens launch");
}
__global__ static void glm53_embedding_bf16_kernel(
float *out,
const uint16_t *weights,
const int32_t *tokens,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_vocab) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t n = (uint64_t)n_tokens * n_embd;
if (i >= n) return;
const uint32_t token_row = (uint32_t)(i / n_embd);
const uint32_t col = (uint32_t)(i - (uint64_t)token_row * n_embd);
const int32_t token = tokens[token_row];
if (token < 0 || (uint32_t)token >= n_vocab) {
out[i] = 0.0f;
return;
}
const uint16_t bits = weights[(uint64_t)(uint32_t)token * n_embd + col];
out[i] = __uint_as_float((uint32_t)bits << 16);
}
__global__ static void glm53_f32_to_bf16_kernel(
__nv_bfloat16 *out,
const float *x,
uint64_t n) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) out[i] = __float2bfloat16(x[i]);
}
__global__ static void glm53_matvec_bf16_f32_kernel(
float *out,
const uint16_t *weights,
const float *x,
uint32_t in_dim,
uint32_t out_dim) {
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t col = blockIdx.x * 8u + warp;
const uint32_t row = blockIdx.y;
float sum = 0.0f;
if (col < out_dim) {
const uint16_t *wrow = weights + (uint64_t)col * in_dim;
const float *xrow = x + (uint64_t)row * in_dim;
for (uint32_t i = lane; i < in_dim; i += 32u) {
const float w = __uint_as_float((uint32_t)wrow[i] << 16);
sum = fmaf(w, xrow[i], sum);
}
}
sum = warp_sum_f32(sum);
if (lane == 0u && col < out_dim) {
out[(uint64_t)row * out_dim + col] = sum;
}
}
extern "C" int ds4_gpu_glm53_embedding_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
const ds4_gpu_tensor *token_ids,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_vocab) {
if (!out || !token_ids || !model_map || n_tokens == 0u ||
n_embd == 0u || n_vocab == 0u ||
(uint64_t)n_vocab > UINT64_MAX / n_embd ||
weight_offset > model_size) {
return 0;
}
const uint64_t weight_elements = (uint64_t)n_vocab * n_embd;
if (weight_elements > UINT64_MAX / sizeof(uint16_t)) return 0;
const uint64_t weight_bytes = weight_elements * sizeof(uint16_t);
const uint64_t output_elements = (uint64_t)n_tokens * n_embd;
if (weight_bytes > model_size - weight_offset ||
token_ids->bytes < (uint64_t)n_tokens * sizeof(int32_t) ||
out->bytes < output_elements * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const char *weights = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, logical_tier,
"GLM-5.3 BF16 embedding");
if (!weights) return 0;
glm53_embedding_bf16_kernel<<<
(unsigned)((output_elements + 255u) / 256u), 256, 0,
cuda_decode_stream()>>>(
(float *)out->ptr, (const uint16_t *)weights,
(const int32_t *)token_ids->ptr, n_tokens, n_embd, n_vocab);
return cuda_ok(cudaGetLastError(), "GLM-5.3 BF16 embedding launch");
}
extern "C" int ds4_gpu_glm53_matmul_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
if (!out || !x || !model_map || !g_cublas_ready || in_dim == 0u ||
out_dim == 0u || n_rows == 0u ||
(uint64_t)out_dim > UINT64_MAX / in_dim ||
weight_offset > model_size) {
return 0;
}
const uint64_t weight_elements = (uint64_t)out_dim * in_dim;
const uint64_t weight_bytes = weight_elements * sizeof(uint16_t);
const uint64_t input_elements = (uint64_t)n_rows * in_dim;
const uint64_t output_elements = (uint64_t)n_rows * out_dim;
if (weight_bytes > model_size - weight_offset ||
x->bytes < input_elements * sizeof(float) ||
out->bytes < output_elements * sizeof(float)) {
return 0;
}
const int logical_tier = ds4_tensor_device_idx(out);
const char *weights = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, logical_tier,
"GLM-5.3 BF16 matrix");
if (!weights) return 0;
if (n_rows <= 8u) {
const dim3 grid((out_dim + 7u) / 8u, n_rows, 1u);
glm53_matvec_bf16_f32_kernel<<<grid, 256u, 0,
cuda_decode_stream()>>>(
(float *)out->ptr, (const uint16_t *)weights,
(const float *)x->ptr, in_dim, out_dim);
return cuda_ok(cudaGetLastError(),
"GLM-5.3 BF16/F32 matvec launch");
}
__nv_bfloat16 *xb = (__nv_bfloat16 *)cuda_tmp_alloc_on(
logical_tier, input_elements * sizeof(__nv_bfloat16),
"GLM-5.3 BF16 matmul activations");
if (!xb) return 0;
glm53_f32_to_bf16_kernel<<<
(unsigned)((input_elements + 255u) / 256u), 256, 0,
cuda_decode_stream()>>>(xb, (const float *)x->ptr, input_elements);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 BF16 activation conversion launch")) {
return 0;
}
const float alpha = 1.0f;
const float beta = 0.0f;
const cublasStatus_t status = cublasGemmEx(
cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T, CUBLAS_OP_N,
(int)out_dim, (int)n_rows, (int)in_dim,
&alpha,
weights, CUDA_R_16BF, (int)in_dim,
xb, CUDA_R_16BF, (int)in_dim,
&beta,
out->ptr, CUDA_R_32F, (int)out_dim,
CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
return cublas_ok(status, "GLM-5.3 BF16 matmul");
}
enum {
GLM53_CUDA_KDA_DIM = 128,
GLM53_CUDA_KDA_HISTORY = 3,
};
__device__ __forceinline__ static float glm53_cuda_silu(float x) {
return x / (1.0f + expf(-x));
}
__device__ __forceinline__ static float glm53_cuda_sigmoid(float x) {
return 1.0f / (1.0f + expf(-x));
}
__global__ static void glm53_kda_decode_kernel(
float *out,
float *conv_state,
float *state,
const float *q_in,
const float *k_in,
const float *v_in,
const float *raw_gate,
const float *raw_beta,
const float *output_gate,
const float *q_conv,
const float *k_conv,
const float *v_conv,
const float *a_log,
const float *dt_bias,
const float *output_norm,
uint32_t n_heads,
uint32_t n_rows,
float lower_bound,
float norm_eps) {
const uint32_t row = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
if (row >= n_rows || head >= n_heads || tid >= GLM53_CUDA_KDA_DIM) {
return;
}
__shared__ float scratch[656];
float *sq = scratch;
float *sk = sq + GLM53_CUDA_KDA_DIM;
float *sd = sk + GLM53_CUDA_KDA_DIM;
float *sv = sd + GLM53_CUDA_KDA_DIM;
float *so = sv + GLM53_CUDA_KDA_DIM;
float *reduce_q = so + GLM53_CUDA_KDA_DIM;
float *reduce_k = reduce_q + 4;
float *reduce_o = reduce_k + 4;
float *beta_shared = reduce_o + 4;
const uint32_t projection = n_heads * GLM53_CUDA_KDA_DIM;
const uint32_t channel = head * GLM53_CUDA_KDA_DIM + tid;
const uint64_t input_base =
(uint64_t)row * projection + head * GLM53_CUDA_KDA_DIM;
const uint64_t conv_row_stride =
(uint64_t)3 * GLM53_CUDA_KDA_HISTORY * projection;
float *q_state = conv_state + (uint64_t)row * conv_row_stride;
float *k_state = q_state + GLM53_CUDA_KDA_HISTORY * projection;
float *v_state = k_state + GLM53_CUDA_KDA_HISTORY * projection;
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
for (uint32_t w = 0; w < GLM53_CUDA_KDA_HISTORY; w++) {
q_acc = fmaf(q_state[(uint64_t)w * projection + channel],
q_conv[(uint64_t)channel * 4u + w], q_acc);
k_acc = fmaf(k_state[(uint64_t)w * projection + channel],
k_conv[(uint64_t)channel * 4u + w], k_acc);
v_acc = fmaf(v_state[(uint64_t)w * projection + channel],
v_conv[(uint64_t)channel * 4u + w], v_acc);
}
const float q_new = q_in[input_base + tid];
const float k_new = k_in[input_base + tid];
const float v_new = v_in[input_base + tid];
q_acc = fmaf(q_new, q_conv[(uint64_t)channel * 4u + 3u], q_acc);
k_acc = fmaf(k_new, k_conv[(uint64_t)channel * 4u + 3u], k_acc);
v_acc = fmaf(v_new, v_conv[(uint64_t)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ull * projection + channel];
q_state[2ull * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ull * projection + channel];
k_state[2ull * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ull * projection + channel];
v_state[2ull * projection + channel] = v_new;
sq[tid] = glm53_cuda_silu(q_acc);
sk[tid] = glm53_cuda_silu(k_acc);
sv[tid] = glm53_cuda_silu(v_acc);
const float gate = raw_gate[input_base + tid] + dt_bias[channel];
sd[tid] = expf(lower_bound *
glm53_cuda_sigmoid(expf(a_log[head]) * gate));
if (tid == 0u) {
beta_shared[0] = glm53_cuda_sigmoid(
raw_beta[(uint64_t)row * n_heads + head]);
}
__syncthreads();
float q_sumsq = warp_sum_f32(sq[tid] * sq[tid]);
float k_sumsq = warp_sum_f32(sk[tid] * sk[tid]);
if (lane == 0u) {
reduce_q[warp] = q_sumsq;
reduce_k[warp] = k_sumsq;
}
__syncthreads();
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = __shfl_sync(0xffffffffu, warp_sum_f32(q_total), 0);
k_total = __shfl_sync(0xffffffffu, warp_sum_f32(k_total), 0);
sq[tid] *= rsqrtf(q_total + 1.0e-6f) * 0.08838834764831845f;
sk[tid] *= rsqrtf(k_total + 1.0e-6f);
__syncthreads();
const uint32_t k0 = lane * 4u;
const float4 q4 = *(const float4 *)(sq + k0);
const float4 k4 = *(const float4 *)(sk + k0);
const float4 decay4 = *(const float4 *)(sd + k0);
const uint64_t state_head =
((uint64_t)row * n_heads + head) *
GLM53_CUDA_KDA_DIM * GLM53_CUDA_KDA_DIM;
for (uint32_t value = warp; value < GLM53_CUDA_KDA_DIM; value += 4u) {
float4 *hptr = (float4 *)(state + state_head +
(uint64_t)value * GLM53_CUDA_KDA_DIM + k0);
float4 h = *hptr;
h.x *= decay4.x;
h.y *= decay4.y;
h.z *= decay4.z;
h.w *= decay4.w;
const float hk = __shfl_sync(
0xffffffffu, warp_sum_f32(dot4_f32(h, k4)), 0);
const float delta_v = (sv[value] - hk) * beta_shared[0];
h.x = fmaf(k4.x, delta_v, h.x);
h.y = fmaf(k4.y, delta_v, h.y);
h.z = fmaf(k4.z, delta_v, h.z);
h.w = fmaf(k4.w, delta_v, h.w);
*hptr = h;
const float hq = __shfl_sync(
0xffffffffu, warp_sum_f32(dot4_f32(h, q4)), 0);
if (lane == 0u) so[value] = hq;
}
__syncthreads();
float o_sumsq = warp_sum_f32(so[tid] * so[tid]);
if (lane == 0u) reduce_o[warp] = o_sumsq;
__syncthreads();
float o_total = lane < 4u ? reduce_o[lane] : 0.0f;
o_total = __shfl_sync(0xffffffffu, warp_sum_f32(o_total), 0);
const float o_scale =
rsqrtf(o_total / (float)GLM53_CUDA_KDA_DIM + norm_eps);
const float out_gate = glm53_cuda_sigmoid(output_gate[input_base + tid]);
out[input_base + tid] =
so[tid] * o_scale * output_norm[tid] * out_gate;
}
__global__ static void glm53_kda_prefill_prepare_kernel(
float *q,
float *k,
float *v,
float *raw_gate,
float *conv_state,
const float *q_conv,
const float *k_conv,
const float *v_conv,
const float *a_log,
const float *dt_bias,
uint32_t n_heads,
uint32_t n_tokens,
float lower_bound) {
const uint32_t head = blockIdx.x;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
if (head >= n_heads || tid >= GLM53_CUDA_KDA_DIM) return;
__shared__ float sq[GLM53_CUDA_KDA_DIM];
__shared__ float sk[GLM53_CUDA_KDA_DIM];
__shared__ float reduce_q[4];
__shared__ float reduce_k[4];
const uint32_t projection = n_heads * GLM53_CUDA_KDA_DIM;
const uint32_t channel = head * GLM53_CUDA_KDA_DIM + tid;
float *q_state = conv_state;
float *k_state = q_state + GLM53_CUDA_KDA_HISTORY * projection;
float *v_state = k_state + GLM53_CUDA_KDA_HISTORY * projection;
for (uint32_t token = 0; token < n_tokens; token++) {
const uint64_t index = (uint64_t)token * projection + channel;
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
for (uint32_t w = 0; w < GLM53_CUDA_KDA_HISTORY; w++) {
q_acc = fmaf(q_state[(uint64_t)w * projection + channel],
q_conv[(uint64_t)channel * 4u + w], q_acc);
k_acc = fmaf(k_state[(uint64_t)w * projection + channel],
k_conv[(uint64_t)channel * 4u + w], k_acc);
v_acc = fmaf(v_state[(uint64_t)w * projection + channel],
v_conv[(uint64_t)channel * 4u + w], v_acc);
}
const float q_new = q[index];
const float k_new = k[index];
const float v_new = v[index];
q_acc = fmaf(q_new, q_conv[(uint64_t)channel * 4u + 3u], q_acc);
k_acc = fmaf(k_new, k_conv[(uint64_t)channel * 4u + 3u], k_acc);
v_acc = fmaf(v_new, v_conv[(uint64_t)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ull * projection + channel];
q_state[2ull * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ull * projection + channel];
k_state[2ull * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ull * projection + channel];
v_state[2ull * projection + channel] = v_new;
sq[tid] = glm53_cuda_silu(q_acc);
sk[tid] = glm53_cuda_silu(k_acc);
v[index] = glm53_cuda_silu(v_acc);
raw_gate[index] = expf(lower_bound * glm53_cuda_sigmoid(
expf(a_log[head]) * (raw_gate[index] + dt_bias[channel])));
__syncthreads();
float q_total = warp_sum_f32(sq[tid] * sq[tid]);
float k_total = warp_sum_f32(sk[tid] * sk[tid]);
if (lane == 0u) {
reduce_q[warp] = q_total;
reduce_k[warp] = k_total;
}
__syncthreads();
q_total = lane < 4u ? reduce_q[lane] : 0.0f;
k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = __shfl_sync(0xffffffffu, warp_sum_f32(q_total), 0);
k_total = __shfl_sync(0xffffffffu, warp_sum_f32(k_total), 0);
q[index] = sq[tid] * rsqrtf(q_total + 1.0e-6f) *
0.08838834764831845f;
k[index] = sk[tid] * rsqrtf(k_total + 1.0e-6f);
__syncthreads();
}
}
__global__ static void glm53_kda_prefill_prepare_parallel_kernel(
float *q,
float *k,
float *v,
float *raw_gate,
const float *raw_q,
const float *raw_k,
const float *raw_v,
const float *conv_state,
const float *q_conv,
const float *k_conv,
const float *v_conv,
const float *a_log,
const float *dt_bias,
uint32_t n_heads,
uint32_t n_tokens,
float lower_bound) {
const uint32_t token = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
if (token >= n_tokens || head >= n_heads ||
tid >= GLM53_CUDA_KDA_DIM) {
return;
}
__shared__ float sq[GLM53_CUDA_KDA_DIM];
__shared__ float sk[GLM53_CUDA_KDA_DIM];
__shared__ float reduce_q[4];
__shared__ float reduce_k[4];
const uint32_t projection = n_heads * GLM53_CUDA_KDA_DIM;
const uint32_t channel = head * GLM53_CUDA_KDA_DIM + tid;
const float *q_state = conv_state;
const float *k_state = q_state + GLM53_CUDA_KDA_HISTORY * projection;
const float *v_state = k_state + GLM53_CUDA_KDA_HISTORY * projection;
const uint64_t index = (uint64_t)token * projection + channel;
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
#pragma unroll
for (uint32_t w = 0; w < GLM53_CUDA_KDA_HISTORY; w++) {
const int32_t source_token =
(int32_t)token + (int32_t)w - GLM53_CUDA_KDA_HISTORY;
float q_value;
float k_value;
float v_value;
if (source_token >= 0) {
const uint64_t source =
(uint64_t)(uint32_t)source_token * projection + channel;
q_value = raw_q[source];
k_value = raw_k[source];
v_value = raw_v[source];
} else {
const uint32_t state_row = token + w;
const uint64_t source =
(uint64_t)state_row * projection + channel;
q_value = q_state[source];
k_value = k_state[source];
v_value = v_state[source];
}
q_acc = fmaf(q_value, q_conv[(uint64_t)channel * 4u + w], q_acc);
k_acc = fmaf(k_value, k_conv[(uint64_t)channel * 4u + w], k_acc);
v_acc = fmaf(v_value, v_conv[(uint64_t)channel * 4u + w], v_acc);
}
q_acc = fmaf(raw_q[index],
q_conv[(uint64_t)channel * 4u + 3u], q_acc);
k_acc = fmaf(raw_k[index],
k_conv[(uint64_t)channel * 4u + 3u], k_acc);
v_acc = fmaf(raw_v[index],
v_conv[(uint64_t)channel * 4u + 3u], v_acc);
sq[tid] = glm53_cuda_silu(q_acc);
sk[tid] = glm53_cuda_silu(k_acc);
v[index] = glm53_cuda_silu(v_acc);
raw_gate[index] = expf(lower_bound * glm53_cuda_sigmoid(
expf(a_log[head]) * (raw_gate[index] + dt_bias[channel])));
__syncthreads();
float q_total = warp_sum_f32(sq[tid] * sq[tid]);
float k_total = warp_sum_f32(sk[tid] * sk[tid]);
if (lane == 0u) {
reduce_q[warp] = q_total;
reduce_k[warp] = k_total;
}
__syncthreads();
q_total = lane < 4u ? reduce_q[lane] : 0.0f;
k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = __shfl_sync(0xffffffffu, warp_sum_f32(q_total), 0);
k_total = __shfl_sync(0xffffffffu, warp_sum_f32(k_total), 0);
q[index] = sq[tid] * rsqrtf(q_total + 1.0e-6f) *
0.08838834764831845f;
k[index] = sk[tid] * rsqrtf(k_total + 1.0e-6f);
}
__global__ static void glm53_kda_prefill_update_conv_state_kernel(
float *conv_state,
const float *raw_q,
const float *raw_k,
const float *raw_v,
uint32_t projection,
uint32_t n_tokens) {
const uint32_t channel = blockIdx.x * blockDim.x + threadIdx.x;
if (channel >= projection) return;
float *q_state = conv_state;
float *k_state = q_state + GLM53_CUDA_KDA_HISTORY * projection;
float *v_state = k_state + GLM53_CUDA_KDA_HISTORY * projection;
#pragma unroll
for (uint32_t w = 0; w < GLM53_CUDA_KDA_HISTORY; w++) {
const uint64_t source =
(uint64_t)(n_tokens - GLM53_CUDA_KDA_HISTORY + w) *
projection + channel;
const uint64_t target = (uint64_t)w * projection + channel;
q_state[target] = raw_q[source];
k_state[target] = raw_k[source];
v_state[target] = raw_v[source];
}
}
__global__ static void glm53_kda_prefill_recurrence_kernel(
float *out,
float *state,
const float *q,
const float *k,
const float *v,
const float *decay,
const float *raw_beta,
uint32_t n_heads,
uint32_t n_tokens) {
const uint32_t head = blockIdx.x;
const uint32_t value = blockIdx.y * 4u + (threadIdx.x >> 5u);
const uint32_t lane = threadIdx.x & 31u;
if (head >= n_heads || value >= GLM53_CUDA_KDA_DIM) return;
const uint32_t projection = n_heads * GLM53_CUDA_KDA_DIM;
const uint32_t k0 = lane * 4u;
float4 *state_ptr = (float4 *)(state +
((uint64_t)head * GLM53_CUDA_KDA_DIM + value) *
GLM53_CUDA_KDA_DIM + k0);
float4 h = *state_ptr;
for (uint32_t token = 0; token < n_tokens; token++) {
const uint64_t base =
(uint64_t)token * projection + head * GLM53_CUDA_KDA_DIM;
const float4 q4 = *(const float4 *)(q + base + k0);
const float4 k4 = *(const float4 *)(k + base + k0);
const float4 d4 = *(const float4 *)(decay + base + k0);
h.x *= d4.x;
h.y *= d4.y;
h.z *= d4.z;
h.w *= d4.w;
const float hk = __shfl_sync(
0xffffffffu, warp_sum_f32(dot4_f32(h, k4)), 0);
const float beta = glm53_cuda_sigmoid(
raw_beta[(uint64_t)token * n_heads + head]);
const float delta_v = (v[base + value] - hk) * beta;
h.x = fmaf(k4.x, delta_v, h.x);
h.y = fmaf(k4.y, delta_v, h.y);
h.z = fmaf(k4.z, delta_v, h.z);
h.w = fmaf(k4.w, delta_v, h.w);
const float result = __shfl_sync(
0xffffffffu, warp_sum_f32(dot4_f32(h, q4)), 0);
if (lane == 0u) out[base + value] = result;
}
*state_ptr = h;
}
__global__ static void glm53_kda_prefill_output_kernel(
float *out,
const float *output_gate,
const float *output_norm,
uint32_t n_heads,
uint32_t n_tokens,
float norm_eps) {
const uint32_t token = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t tid = threadIdx.x;
const uint32_t lane = tid & 31u;
const uint32_t warp = tid >> 5u;
if (token >= n_tokens || head >= n_heads || tid >= GLM53_CUDA_KDA_DIM) {
return;
}
__shared__ float partial[4];
const uint32_t projection = n_heads * GLM53_CUDA_KDA_DIM;
const uint64_t index = (uint64_t)token * projection +
head * GLM53_CUDA_KDA_DIM + tid;
const float raw = out[index];
float total = warp_sum_f32(raw * raw);
if (lane == 0u) partial[warp] = total;
__syncthreads();
total = lane < 4u ? partial[lane] : 0.0f;
total = __shfl_sync(0xffffffffu, warp_sum_f32(total), 0);
const float scale =
rsqrtf(total / (float)GLM53_CUDA_KDA_DIM + norm_eps);
out[index] = raw * scale * output_norm[tid] *
glm53_cuda_sigmoid(output_gate[index]);
}
static bool glm53_cuda_mul_u64(uint64_t a, uint64_t b, uint64_t *out) {
if (!out || (a != 0u && b > UINT64_MAX / a)) return false;
*out = a * b;
return true;
}
static bool glm53_cuda_tensor_has(
const ds4_gpu_tensor *tensor, uint64_t elements, uint64_t elem_size) {
return tensor && tensor->ptr &&
(elements == 0u || elem_size <= UINT64_MAX / elements) &&
tensor->bytes >= elements * elem_size;
}
static const float *glm53_cuda_weight_f32(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t elements,
int logical_tier,
const char *label) {
uint64_t bytes = 0;
if (!model_map || !glm53_cuda_mul_u64(elements, sizeof(float), &bytes) ||
offset > model_size || bytes > model_size - offset) {
fprintf(stderr, "ds4: GLM-5.3 %s range is outside the mapped model\n",
label ? label : "weight");
return NULL;
}
return (const float *)cuda_resolve_weight_ptr(
model_map, offset, bytes, logical_tier, label);
}
extern "C" int ds4_gpu_glm53_kda_decode(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *k,
const ds4_gpu_tensor *v,
const ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_rows,
float gate_lower_bound,
float norm_eps) {
uint64_t projection = 0, activations = 0, conv_elements = 0;
uint64_t state_elements = 0;
if (n_heads == 0u || n_rows == 0u || gate_lower_bound >= 0.0f ||
!glm53_cuda_mul_u64(n_heads, GLM53_CUDA_KDA_DIM, &projection) ||
!glm53_cuda_mul_u64(projection, n_rows, &activations) ||
!glm53_cuda_mul_u64(activations,
3u * GLM53_CUDA_KDA_HISTORY, &conv_elements) ||
!glm53_cuda_mul_u64(activations,
GLM53_CUDA_KDA_DIM, &state_elements) ||
!glm53_cuda_tensor_has(q, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(k, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(v, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(raw_gate, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(raw_beta,
(uint64_t)n_rows * n_heads, sizeof(float)) ||
!glm53_cuda_tensor_has(output_gate, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(out, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(conv_state, conv_elements, sizeof(float)) ||
!glm53_cuda_tensor_has(recurrent_state, state_elements, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 KDA decode received invalid buffers\n");
return 0;
}
const int tier = ds4_tensor_device_idx(out);
const float *qw = glm53_cuda_weight_f32(model_map, model_size,
q_conv_offset, projection * 4u, tier, "KDA Q convolution");
const float *kw = glm53_cuda_weight_f32(model_map, model_size,
k_conv_offset, projection * 4u, tier, "KDA K convolution");
const float *vw = glm53_cuda_weight_f32(model_map, model_size,
v_conv_offset, projection * 4u, tier, "KDA V convolution");
const float *a_log = glm53_cuda_weight_f32(model_map, model_size,
a_log_offset, n_heads, tier, "KDA A_log");
const float *dt_bias = glm53_cuda_weight_f32(model_map, model_size,
dt_bias_offset, projection, tier, "KDA dt bias");
const float *output_norm = glm53_cuda_weight_f32(model_map, model_size,
output_norm_offset, GLM53_CUDA_KDA_DIM, tier, "KDA output norm");
if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm) return 0;
const dim3 grid(n_rows, n_heads, 1u);
glm53_kda_decode_kernel<<<grid, GLM53_CUDA_KDA_DIM, 0,
cuda_decode_stream()>>>(
(float *)out->ptr, (float *)conv_state->ptr,
(float *)recurrent_state->ptr, (const float *)q->ptr,
(const float *)k->ptr, (const float *)v->ptr,
(const float *)raw_gate->ptr, (const float *)raw_beta->ptr,
(const float *)output_gate->ptr, qw, kw, vw, a_log, dt_bias,
output_norm, n_heads, n_rows, gate_lower_bound, norm_eps);
return cuda_ok(cudaGetLastError(), "GLM-5.3 KDA decode launch");
}
extern "C" int ds4_gpu_glm53_kda_prefill(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
ds4_gpu_tensor *q,
ds4_gpu_tensor *k,
ds4_gpu_tensor *v,
ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_tokens,
float gate_lower_bound,
float norm_eps) {
uint64_t projection = 0, activations = 0, conv_elements = 0;
uint64_t state_elements = 0;
if (n_heads == 0u || n_tokens == 0u || gate_lower_bound >= 0.0f ||
!glm53_cuda_mul_u64(n_heads, GLM53_CUDA_KDA_DIM, &projection) ||
!glm53_cuda_mul_u64(projection, n_tokens, &activations) ||
!glm53_cuda_mul_u64(projection,
3u * GLM53_CUDA_KDA_HISTORY, &conv_elements) ||
!glm53_cuda_mul_u64(projection,
GLM53_CUDA_KDA_DIM, &state_elements) ||
!glm53_cuda_tensor_has(q, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(k, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(v, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(raw_gate, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(raw_beta,
(uint64_t)n_tokens * n_heads, sizeof(float)) ||
!glm53_cuda_tensor_has(output_gate, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(out, activations, sizeof(float)) ||
!glm53_cuda_tensor_has(conv_state, conv_elements, sizeof(float)) ||
!glm53_cuda_tensor_has(recurrent_state, state_elements, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 KDA prefill received invalid buffers\n");
return 0;
}
const int tier = ds4_tensor_device_idx(out);
const float *qw = glm53_cuda_weight_f32(model_map, model_size,
q_conv_offset, projection * 4u, tier, "KDA Q convolution");
const float *kw = glm53_cuda_weight_f32(model_map, model_size,
k_conv_offset, projection * 4u, tier, "KDA K convolution");
const float *vw = glm53_cuda_weight_f32(model_map, model_size,
v_conv_offset, projection * 4u, tier, "KDA V convolution");
const float *a_log = glm53_cuda_weight_f32(model_map, model_size,
a_log_offset, n_heads, tier, "KDA A_log");
const float *dt_bias = glm53_cuda_weight_f32(model_map, model_size,
dt_bias_offset, projection, tier, "KDA dt bias");
const float *output_norm = glm53_cuda_weight_f32(model_map, model_size,
output_norm_offset, GLM53_CUDA_KDA_DIM, tier, "KDA output norm");
if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm) return 0;
cudaStream_t stream = cuda_decode_stream();
const float *recurrence_q = (const float *)q->ptr;
const float *recurrence_k = (const float *)k->ptr;
const float *recurrence_v = (const float *)v->ptr;
bool used_parallel_prepare = false;
if (n_tokens >= GLM53_CUDA_KDA_HISTORY &&
activations <= UINT64_MAX / 3u &&
3u * activations <= SIZE_MAX / sizeof(float)) {
float *prepared = (float *)cuda_tmp_alloc_on(
tier, 3u * activations * sizeof(float),
"GLM-5.3 prepared KDA prefill activations");
if (prepared) {
float *prepared_q = prepared;
float *prepared_k = prepared_q + activations;
float *prepared_v = prepared_k + activations;
const dim3 prepare_grid(n_tokens, n_heads, 1u);
glm53_kda_prefill_prepare_parallel_kernel<<<
prepare_grid, GLM53_CUDA_KDA_DIM, 0, stream>>>(
prepared_q, prepared_k, prepared_v,
(float *)raw_gate->ptr,
(const float *)q->ptr,
(const float *)k->ptr,
(const float *)v->ptr,
(const float *)conv_state->ptr,
qw, kw, vw, a_log, dt_bias, n_heads, n_tokens,
gate_lower_bound);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 parallel KDA prefill prepare launch")) {
return 0;
}
glm53_kda_prefill_update_conv_state_kernel<<<
(projection + 255u) / 256u, 256u, 0, stream>>>(
(float *)conv_state->ptr,
(const float *)q->ptr,
(const float *)k->ptr,
(const float *)v->ptr,
(uint32_t)projection, n_tokens);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 KDA prefill state update launch")) {
return 0;
}
recurrence_q = prepared_q;
recurrence_k = prepared_k;
recurrence_v = prepared_v;
used_parallel_prepare = true;
}
}
if (!used_parallel_prepare) {
glm53_kda_prefill_prepare_kernel<<<
n_heads, GLM53_CUDA_KDA_DIM, 0, stream>>>(
(float *)q->ptr, (float *)k->ptr, (float *)v->ptr,
(float *)raw_gate->ptr, (float *)conv_state->ptr,
qw, kw, vw, a_log, dt_bias, n_heads, n_tokens,
gate_lower_bound);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 KDA prefill prepare launch")) {
return 0;
}
}
const dim3 recurrence_grid(n_heads, 32u, 1u);
glm53_kda_prefill_recurrence_kernel<<<recurrence_grid, 128u, 0,
stream>>>(
(float *)out->ptr, (float *)recurrent_state->ptr,
recurrence_q, recurrence_k, recurrence_v,
(const float *)raw_gate->ptr,
(const float *)raw_beta->ptr, n_heads, n_tokens);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 KDA prefill recurrence launch")) {
return 0;
}
const dim3 output_grid(n_tokens, n_heads, 1u);
glm53_kda_prefill_output_kernel<<<output_grid, GLM53_CUDA_KDA_DIM, 0,
stream>>>(
(float *)out->ptr, (const float *)output_gate->ptr,
output_norm, n_heads, n_tokens, norm_eps);
return cuda_ok(cudaGetLastError(), "GLM-5.3 KDA prefill output launch");
}
extern "C" int ds4_gpu_flush_encoder(void) {
/* Metal encoder flush: CUDA kernels are already queued in stream
* order, nothing to split. */
return 1;
}
template <typename CT>
__global__ static void glm_attention_full_reference_kernel(
float *heads,
const float *q,
const CT *key_cache,
const CT *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim) {
const uint32_t token = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t tid = threadIdx.x;
if (token >= n_tokens || head >= n_head) return;
extern __shared__ float reduce[];
const float *qh = q +
((uint64_t)token * n_head + head) * qk_dim;
const uint32_t visible = min(cache_len, pos0 + token + 1u);
const float scale = rsqrtf((float)qk_dim);
float max_score = -FLT_MAX;
float sum_weight = 0.0f;
float output = 0.0f;
for (uint32_t row = 0; row < visible; row++) {
const CT *krow = key_cache +
((uint64_t)row * n_head + head) * qk_dim;
float dot = 0.0f;
for (uint32_t d = tid; d < qk_dim; d += blockDim.x) {
dot = fmaf(qh[d], (float)krow[d], dot);
}
reduce[tid] = dot;
__syncthreads();
for (uint32_t step = blockDim.x >> 1u; step > 0u; step >>= 1u) {
if (tid < step) reduce[tid] += reduce[tid + step];
__syncthreads();
}
const float score = reduce[0] * scale;
const float next_max = fmaxf(max_score, score);
const float old_scale = max_score == -FLT_MAX
? 0.0f : expf(max_score - next_max);
const float row_weight = expf(score - next_max);
if (tid < value_dim) {
const CT *vrow = value_cache +
((uint64_t)row * n_head + head) * value_dim;
output = output * old_scale + row_weight * (float)vrow[tid];
}
sum_weight = sum_weight * old_scale + row_weight;
max_score = next_max;
__syncthreads();
}
if (tid < value_dim) {
heads[((uint64_t)token * n_head + head) * value_dim + tid] =
output / fmaxf(sum_weight, 1.0e-20f);
}
}
static int glm_attention_full_reference_launch(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float);
if (!heads || !q || !key_cache || !value_cache || n_tokens == 0u ||
cache_len == 0u || cache_cap == 0u || n_head == 0u ||
qk_dim == 0u || qk_dim > 256u || value_dim == 0u ||
value_dim > 256u || cache_len > cache_cap || pos0 > cache_len ||
n_tokens > cache_len - pos0 ||
q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) ||
heads->bytes <
(uint64_t)n_tokens * n_head * value_dim * sizeof(float) ||
key_cache->bytes <
(uint64_t)cache_cap * n_head * qk_dim * cache_elem ||
value_cache->bytes <
(uint64_t)cache_cap * n_head * value_dim * cache_elem) {
return 0;
}
const dim3 grid(n_tokens, n_head, 1u);
if (cache_f16) {
glm_attention_full_reference_kernel<__half><<<grid, 256u,
256u * sizeof(float), cuda_decode_stream()>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const __half *)key_cache->ptr,
(const __half *)value_cache->ptr, pos0, n_tokens, cache_len,
cache_cap, n_head, qk_dim, value_dim);
} else {
glm_attention_full_reference_kernel<float><<<grid, 256u,
256u * sizeof(float), cuda_decode_stream()>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const float *)key_cache->ptr,
(const float *)value_cache->ptr, pos0, n_tokens, cache_len,
cache_cap, n_head, qk_dim, value_dim);
}
return cuda_ok(cudaGetLastError(), "GLM full attention reference launch");
}
extern "C" int ds4_gpu_glm_attention_flash_staged_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
if (pos0 != 0u || n_tokens != cache_len) return 0;
return glm_attention_full_reference_launch(
heads, q, key_cache, value_cache, pos0, n_tokens, cache_len,
cache_cap, n_head, qk_dim, value_dim, cache_f16);
}
extern "C" int ds4_gpu_glm_attention_flash_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
return glm_attention_full_reference_launch(
heads, q, key_cache, value_cache, pos0, n_tokens, cache_len,
cache_cap, n_head, qk_dim, value_dim, cache_f16);
}
extern "C" int ds4_gpu_glm_attention_full_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
return glm_attention_full_reference_launch(
heads, q, key_cache, value_cache, pos0, n_tokens, cache_len,
cache_cap, n_head, qk_dim, value_dim, cache_f16);
}
template <typename CT>
__device__ __forceinline__ static float2 glm_cache_rope_pair_f16_dev(
const CT *rope_cache, uint64_t rope_base, uint32_t r,
uint32_t row, uint32_t qk_rope, float freq_base, float freq_scale,
float ext_factor, float attn_factor, const float corr_dims[2]) {
const float theta_base = (float)row;
const float inv_ndims = -1.0f / (float)qk_rope;
const float theta = theta_base * powf(freq_base, inv_ndims * (float)r);
float ct, st;
glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)r,
ext_factor, attn_factor, &ct, &st);
const float x0 = (float)rope_cache[rope_base + r];
const float x1 = (float)rope_cache[rope_base + r + 1u];
return make_float2(x0 * ct - x1 * st, x0 * st + x1 * ct);
}
__global__ static void glm_dense_attn_pack_q_f16_kernel(
__half *packed,
const float *q,
uint32_t n_q,
uint32_t n_head,
uint32_t head0,
uint32_t group_heads,
uint32_t dim) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t count = (uint64_t)group_heads * n_q * dim;
if (i >= count) return;
const uint32_t d = (uint32_t)(i % dim);
const uint64_t row = i / dim;
const uint32_t token = (uint32_t)(row % n_q);
const uint32_t head = head0 + (uint32_t)(row / n_q);
packed[i] = __float2half(q[((uint64_t)token * n_head + head) * dim + d]);
}
__global__ static void glm_dense_attn_causal_softmax_f16_kernel(
__half *prob,
const float *scores,
uint32_t n_q,
uint32_t n_kv,
uint32_t q_row0,
float scale) {
const uint32_t token = blockIdx.x;
const uint32_t group_head = blockIdx.y;
if (token >= n_q) return;
const uint32_t visible = min(n_kv, q_row0 + token + 1u);
const uint64_t row = ((uint64_t)group_head * n_q + token) * n_kv;
float local_max = -FLT_MAX;
for (uint32_t col = threadIdx.x; col < visible; col += blockDim.x) {
local_max = fmaxf(local_max, scores[row + col] * scale);
}
__shared__ float reduce[256];
reduce[threadIdx.x] = local_max;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
reduce[threadIdx.x] = fmaxf(reduce[threadIdx.x],
reduce[threadIdx.x + stride]);
}
__syncthreads();
}
const float max_score = reduce[0];
__syncthreads();
float local_sum = 0.0f;
for (uint32_t col = threadIdx.x; col < visible; col += blockDim.x) {
local_sum += expf(scores[row + col] * scale - max_score);
}
reduce[threadIdx.x] = local_sum;
__syncthreads();
for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) {
if (threadIdx.x < stride) {
reduce[threadIdx.x] += reduce[threadIdx.x + stride];
}
__syncthreads();
}
const float inv_sum = reduce[0] > 0.0f ? 1.0f / reduce[0] : 0.0f;
for (uint32_t col = threadIdx.x; col < n_kv; col += blockDim.x) {
const float p = col < visible
? expf(scores[row + col] * scale - max_score) * inv_sum : 0.0f;
prob[row + col] = __float2half(p);
}
}
__global__ static void glm_dense_attn_unpack_f32_kernel(
float *out,
const float *packed,
uint32_t n_q,
uint32_t n_head,
uint32_t head0,
uint32_t group_heads,
uint32_t dim) {
const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t count = (uint64_t)group_heads * n_q * dim;
if (i >= count) return;
const uint32_t d = (uint32_t)(i % dim);
const uint64_t row = i / dim;
const uint32_t token = (uint32_t)(row % n_q);
const uint32_t head = head0 + (uint32_t)(row / n_q);
out[((uint64_t)token * n_head + head) * dim + d] = packed[i];
}
extern "C" int ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_dim) {
if (!lora_out || !qk_low || !kv_lora_cache || !g_cublas_ready ||
g_n_gpus != 1 || n_q == 0u || n_kv == 0u || n_kv > cache_cap ||
q_row0 >= n_kv || n_q > n_kv - q_row0 || !cache_f16 ||
n_q > INT_MAX || n_kv > INT_MAX || n_head == 0u ||
kv_lora_dim != 512u || qk_dim == 0u ||
qk_low->bytes < (uint64_t)n_q * n_head * kv_lora_dim * sizeof(float) ||
kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * sizeof(__half) ||
lora_out->bytes < (uint64_t)n_q * n_head * kv_lora_dim * sizeof(float)) {
return 0;
}
const uint32_t visible_kv = min(n_kv, q_row0 + n_q);
const uint32_t max_group_heads = min(n_head, 8u);
if ((uint64_t)n_q > UINT64_MAX / max_group_heads / visible_kv) {
return 0;
}
const uint64_t q_count =
(uint64_t)max_group_heads * n_q * kv_lora_dim;
const uint64_t score_count =
(uint64_t)max_group_heads * n_q * visible_kv;
const uint64_t q_bytes = q_count * sizeof(__half);
const uint64_t score_off = (q_bytes + 255u) & ~255ull;
const uint64_t score_bytes = score_count * sizeof(float);
const uint64_t prob_off = (score_off + score_bytes + 255u) & ~255ull;
const uint64_t prob_bytes = score_count * sizeof(__half);
const uint64_t out_off = (prob_off + prob_bytes + 255u) & ~255ull;
const uint64_t out_bytes = q_count * sizeof(float);
if (score_off < q_bytes || prob_off < score_off || out_off < prob_off ||
out_off > UINT64_MAX - out_bytes) {
return 0;
}
unsigned char *scratch = (unsigned char *)cuda_tmp_alloc_on(
0, out_off + out_bytes, "GLM dense compact attention");
if (!scratch) return 0;
__half *q_half = (__half *)scratch;
float *scores = (float *)(scratch + score_off);
__half *prob = (__half *)(scratch + prob_off);
float *packed_out = (float *)(scratch + out_off);
const __half *kv = (const __half *)kv_lora_cache->ptr;
const float attn_scale = 1.0f / sqrtf((float)qk_dim);
cublasHandle_t handle = cuda_cublas_for_tier(0);
const float alpha = 1.0f;
const float beta = 0.0f;
for (uint32_t head0 = 0; head0 < n_head; head0 += max_group_heads) {
const uint32_t group_heads = min(max_group_heads, n_head - head0);
const uint64_t group_q_count =
(uint64_t)group_heads * n_q * kv_lora_dim;
glm_dense_attn_pack_q_f16_kernel<<<
(group_q_count + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
q_half, (const float *)qk_low->ptr, n_q, n_head, head0,
group_heads, kv_lora_dim);
if (!cuda_ok(cudaGetLastError(),
"GLM dense attention Q pack launch")) return 0;
cublasStatus_t st = cublasGemmStridedBatchedEx(
handle,
CUBLAS_OP_T, CUBLAS_OP_N,
(int)visible_kv, (int)n_q, (int)kv_lora_dim,
&alpha,
kv, CUDA_R_16F, (int)kv_lora_dim, 0,
q_half, CUDA_R_16F, (int)kv_lora_dim,
(long long)n_q * kv_lora_dim,
&beta,
scores, CUDA_R_32F, (int)visible_kv,
(long long)n_q * visible_kv,
(int)group_heads,
CUBLAS_COMPUTE_32F,
CUBLAS_GEMM_DEFAULT);
if (!cublas_ok(st, "GLM dense attention score GEMM")) return 0;
dim3 softmax_grid(n_q, group_heads, 1u);
glm_dense_attn_causal_softmax_f16_kernel<<<
softmax_grid, 256, 0, cuda_decode_stream()>>>(
prob, scores, n_q, visible_kv, q_row0, attn_scale);
if (!cuda_ok(cudaGetLastError(),
"GLM dense attention softmax launch")) return 0;
st = cublasGemmStridedBatchedEx(
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
(int)kv_lora_dim, (int)n_q, (int)visible_kv,
&alpha,
kv, CUDA_R_16F, (int)kv_lora_dim, 0,
prob, CUDA_R_16F, (int)visible_kv,
(long long)n_q * visible_kv,
&beta,
packed_out, CUDA_R_32F, (int)kv_lora_dim,
(long long)n_q * kv_lora_dim,
(int)group_heads,
CUBLAS_COMPUTE_32F,
CUBLAS_GEMM_DEFAULT);
if (!cublas_ok(st, "GLM dense attention value GEMM")) return 0;
glm_dense_attn_unpack_f32_kernel<<<
(group_q_count + 255u) / 256u, 256, 0, cuda_decode_stream()>>>(
(float *)lora_out->ptr, packed_out, n_q, n_head, head0,
group_heads, kv_lora_dim);
if (!cuda_ok(cudaGetLastError(),
"GLM dense attention output unpack launch")) return 0;
}
return 1;
}
/* Scalar-correct MLA attention: one warp per head, grid
* (ceil(n_head/8), n_tokens). The row loop handles either a contiguous
* causal range or an explicit selected-row list and mirrors the Metal online
* softmax so numerics stay comparable. */
template <typename CT, bool selected_rows>
__global__ static void glm_attention_lora_causal_kernel(
float *lora_out,
const float *q,
const float *qk_low,
const CT *kv_lora_cache,
const CT *k_rope_cache,
const uint32_t *selected,
uint32_t cache_cap,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float scale) {
const uint32_t token = blockIdx.y;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t head = blockIdx.x * 8u + warp;
if (token >= n_tokens || head >= n_head || kv_lora_dim != 512u ||
(qk_rope != 0u && qk_rope != 64u)) {
return;
}
const uint32_t visible = selected_rows
? n_selected : min(n_selected, pos0 + token + 1u);
if (visible == 0u) return;
const uint32_t qk_dim = qk_nope + qk_rope;
const float *qh = q + (uint64_t)token * n_head * qk_dim +
(uint64_t)head * qk_dim;
const float4 *low4 = (const float4 *)(qk_low +
(uint64_t)token * n_head * kv_lora_dim +
(uint64_t)head * kv_lora_dim);
float4 low0 = low4[lane];
float4 low1 = low4[lane + 32u];
float4 low2 = low4[lane + 64u];
float4 low3 = low4[lane + 96u];
float4 qrope = make_float4(0.f, 0.f, 0.f, 0.f);
const uint32_t rope_vecs = qk_rope >> 2; /* 16 */
if (lane < rope_vecs) {
qrope = *((const float4 *)(qh + qk_nope + lane * 4u));
}
float corr_dims[2] = {0.0f, 0.0f};
if (qk_rope != 0u && ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig,
beta_fast, freq_base)));
corr_dims[1] = fminf((float)qk_rope - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope, (int)n_ctx_orig,
beta_slow, freq_base)));
}
float M = -FLT_MAX / 2.0f;
float S = 0.0f;
float4 o0 = make_float4(0.f,0.f,0.f,0.f);
float4 o1 = o0, o2 = o0, o3 = o0;
for (uint32_t ri = 0u; ri < visible; ri++) {
const uint32_t row = selected_rows
? selected[(uint64_t)token * n_selected + ri] : ri;
if (row >= cache_cap) continue;
const CT *kvrow = kv_lora_cache + (uint64_t)row * kv_lora_dim;
float partial = 0.0f;
{
const float4 k0 = make_float4(
(float)(kvrow[lane*4u+0u]), (float)(kvrow[lane*4u+1u]),
(float)(kvrow[lane*4u+2u]), (float)(kvrow[lane*4u+3u]));
const float4 k1 = make_float4(
(float)(kvrow[(lane+32u)*4u+0u]), (float)(kvrow[(lane+32u)*4u+1u]),
(float)(kvrow[(lane+32u)*4u+2u]), (float)(kvrow[(lane+32u)*4u+3u]));
const float4 k2 = make_float4(
(float)(kvrow[(lane+64u)*4u+0u]), (float)(kvrow[(lane+64u)*4u+1u]),
(float)(kvrow[(lane+64u)*4u+2u]), (float)(kvrow[(lane+64u)*4u+3u]));
const float4 k3 = make_float4(
(float)(kvrow[(lane+96u)*4u+0u]), (float)(kvrow[(lane+96u)*4u+1u]),
(float)(kvrow[(lane+96u)*4u+2u]), (float)(kvrow[(lane+96u)*4u+3u]));
partial += low0.x*k0.x + low0.y*k0.y + low0.z*k0.z + low0.w*k0.w;
partial += low1.x*k1.x + low1.y*k1.y + low1.z*k1.z + low1.w*k1.w;
partial += low2.x*k2.x + low2.y*k2.y + low2.z*k2.z + low2.w*k2.w;
partial += low3.x*k3.x + low3.y*k3.y + low3.z*k3.z + low3.w*k3.w;
if (lane < rope_vecs) {
const uint64_t rope_base = (uint64_t)row * qk_rope;
const uint32_t r = lane * 4u;
const float2 y0 = glm_cache_rope_pair_f16_dev(
k_rope_cache, rope_base, r, row, qk_rope, freq_base,
freq_scale, ext_factor, attn_factor, corr_dims);
const float2 y1 = glm_cache_rope_pair_f16_dev(
k_rope_cache, rope_base, r + 2u, row, qk_rope,
freq_base, freq_scale, ext_factor, attn_factor,
corr_dims);
partial += qrope.x*y0.x + qrope.y*y0.y +
qrope.z*y1.x + qrope.w*y1.y;
}
for (uint32_t off = 16u; off > 0u; off >>= 1u) {
partial += __shfl_xor_sync(0xffffffffu, partial, off);
}
const float score = partial * scale;
const float new_m = fmaxf(M, score);
const float old_scale = expf(M - new_m);
const float row_scale = expf(score - new_m);
o0.x = o0.x*old_scale + k0.x*row_scale; o0.y = o0.y*old_scale + k0.y*row_scale;
o0.z = o0.z*old_scale + k0.z*row_scale; o0.w = o0.w*old_scale + k0.w*row_scale;
o1.x = o1.x*old_scale + k1.x*row_scale; o1.y = o1.y*old_scale + k1.y*row_scale;
o1.z = o1.z*old_scale + k1.z*row_scale; o1.w = o1.w*old_scale + k1.w*row_scale;
o2.x = o2.x*old_scale + k2.x*row_scale; o2.y = o2.y*old_scale + k2.y*row_scale;
o2.z = o2.z*old_scale + k2.z*row_scale; o2.w = o2.w*old_scale + k2.w*row_scale;
o3.x = o3.x*old_scale + k3.x*row_scale; o3.y = o3.y*old_scale + k3.y*row_scale;
o3.z = o3.z*old_scale + k3.z*row_scale; o3.w = o3.w*old_scale + k3.w*row_scale;
S = S*old_scale + row_scale;
M = new_m;
}
}
const float inv_s = S > 0.0f ? 1.0f / S : 0.0f;
float4 *out4 = (float4 *)(lora_out +
((uint64_t)token * n_head + head) * kv_lora_dim);
o0.x*=inv_s; o0.y*=inv_s; o0.z*=inv_s; o0.w*=inv_s;
o1.x*=inv_s; o1.y*=inv_s; o1.z*=inv_s; o1.w*=inv_s;
o2.x*=inv_s; o2.y*=inv_s; o2.z*=inv_s; o2.w*=inv_s;
o3.x*=inv_s; o3.y*=inv_s; o3.z*=inv_s; o3.w*=inv_s;
out4[lane] = o0;
out4[lane + 32u] = o1;
out4[lane + 64u] = o2;
out4[lane + 96u] = o3;
}
extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
n_tokens == 0 || n_head == 0 || kv_lora_dim != 512u ||
(qk_rope != 0u && qk_rope != 64u)) {
fprintf(stderr, "ds4: glm attn lora causal: unsupported config "
"(n_tok=%u head=%u lora=%u rope=%u f16=%d)\n",
n_tokens, n_head, kv_lora_dim, qk_rope, (int)cache_f16);
return 0;
}
const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope));
dim3 grid((n_head + 7u) / 8u, n_tokens, 1);
if (cache_f16) {
glm_attention_lora_causal_kernel<__half, false><<<grid, 256>>>(
(float *)lora_out->ptr,
(const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
NULL, cache_cap,
n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope,
qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, scale);
} else {
glm_attention_lora_causal_kernel<float, false><<<grid, 256>>>(
(float *)lora_out->ptr,
(const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
NULL, cache_cap,
n_tokens, pos0, n_selected, n_head, kv_lora_dim, qk_nope,
qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, scale);
}
return cuda_ok(cudaGetLastError(), "glm attn lora causal launch");
}
extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float);
const uint64_t qk_dim = (uint64_t)qk_nope + qk_rope;
if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
!selected || n_tokens == 0u || n_selected == 0u || n_head == 0u ||
kv_lora_dim != 512u || (qk_rope != 0u && qk_rope != 64u) ||
cache_cap == 0u ||
selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t) ||
q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) ||
qk_low->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) ||
kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem ||
k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem ||
lora_out->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) {
fprintf(stderr, "ds4: glm attn lora selected: unsupported config "
"(n_tok=%u selected=%u head=%u lora=%u rope=%u f16=%d)\n",
n_tokens, n_selected, n_head, kv_lora_dim, qk_rope,
(int)cache_f16);
return 0;
}
const float scale = 1.0f / sqrtf((float)(qk_nope + qk_rope));
dim3 grid((n_head + 7u) / 8u, n_tokens, 1);
if (cache_f16) {
glm_attention_lora_causal_kernel<__half, true><<<grid, 256>>>(
(float *)lora_out->ptr,
(const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr, cache_cap,
n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope,
qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, scale);
} else {
glm_attention_lora_causal_kernel<float, true><<<grid, 256>>>(
(float *)lora_out->ptr,
(const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr, cache_cap,
n_tokens, 0u, n_selected, n_head, kv_lora_dim, qk_nope,
qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, scale);
}
return cuda_ok(cudaGetLastError(), "glm attn lora selected launch");
}
extern "C" int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens, uint32_t n_selected, uint32_t cache_cap,
bool cache_f16, uint32_t n_head, uint32_t kv_lora_dim,
uint32_t qk_nope, uint32_t qk_rope, uint32_t n_ctx_orig,
float freq_base, float freq_scale, float ext_factor,
float attn_factor, float beta_fast, float beta_slow) {
return ds4_gpu_glm_attention_indexed_batch_lora_tensor(
lora_out, q, qk_low, kv_lora_cache, k_rope_cache, selected,
n_tokens, n_selected, cache_cap, cache_f16, n_head, kv_lora_dim,
qk_nope, qk_rope, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow);
}
extern "C" int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_batch_typed_tensor\n");
return 0;
}
extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(
ds4_gpu_tensor *heads,
ds4_gpu_tensor *partial_lora,
ds4_gpu_tensor *partial_ms,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
bool selected_rows_valid,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
uint32_t block_rows,
uint32_t n_blocks,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor\n");
return 0;
}
__device__ __forceinline__ static float glm_q8_0_dot_row_dev(
const char *row, const float *x, uint32_t n_cols) {
float acc = 0.0f;
const uint32_t nb = n_cols >> 5;
for (uint32_t b = 0; b < nb; b++) {
const char *blk = row + (uint64_t)b * 34u;
const float d = __half2float(*(const __half *)blk);
const int8_t *q = (const int8_t *)(blk + 2);
float s = 0.0f;
#pragma unroll 8
for (uint32_t k = 0; k < 32u; k++) s += (float)q[k] * x[b * 32u + k];
acc += d * s;
}
return acc;
}
template <typename CT>
__device__ __forceinline__ static float2 glm_cache_value_pair_dev(
const CT *p) {
return make_float2((float)p[0], (float)p[1]);
}
template <>
__device__ __forceinline__ float2 glm_cache_value_pair_dev<__half>(
const __half *p) {
return __half22float2(*(const __half2 *)p);
}
template <>
__device__ __forceinline__ float2 glm_cache_value_pair_dev<float>(
const float *p) {
return *(const float2 *)p;
}
/* Exact staged decode attention. The original fused kernel owns one block per
* head, which leaves more than half of an L40S idle. These stages preserve the
* fused kernel's arithmetic order for every score, softmax lane, lora output,
* and value-projection row while exposing independent rows/dimensions as
* separate blocks. */
template <typename CT, bool RANGE_TOK2 = false>
__global__ static void glm_attention_decode_weights_staged_kernel(
float *weights,
float *denom,
const float *q,
const float *qk_low,
const CT *kv_lora_cache,
const CT *k_rope_cache,
const uint32_t *selected,
uint32_t n_selected,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
float scale,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool score_vec2) {
const uint32_t head = blockIdx.x;
const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u;
const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u);
const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u);
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
if (head >= n_head || row_count == 0u) return;
const uint32_t qk_dim = qk_nope + qk_rope;
extern __shared__ float glm_dec_stage_sh[];
float *red = glm_dec_stage_sh;
float *scores = glm_dec_stage_sh + 256u;
const float *qh = q +
((uint64_t)token * n_head + head) * qk_dim;
const float *low = qk_low +
((uint64_t)token * n_head + head) * kv_lora_dim;
float corr_dims[2] = {0.0f, 0.0f};
if (qk_rope != 0u && ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope,
(int)n_ctx_orig, beta_fast, freq_base)));
corr_dims[1] = fminf((float)qk_rope - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope,
(int)n_ctx_orig, beta_slow, freq_base)));
}
float local_max = -FLT_MAX;
for (uint32_t s = tid; s < row_count; s += nth) {
const uint32_t row = RANGE_TOK2 ? s : selected[s];
float score = -FLT_MAX;
if (row < cache_cap) {
float dotv = 0.0f;
const uint64_t lora_base = (uint64_t)row * kv_lora_dim;
if (score_vec2) {
for (uint32_t j = 0; j < kv_lora_dim; j += 2u) {
const float2 x = *(const float2 *)(low + j);
const float2 y = glm_cache_value_pair_dev(
kv_lora_cache + lora_base + j);
dotv += x.x * y.x;
dotv += x.y * y.y;
}
} else {
for (uint32_t j = 0; j < kv_lora_dim; j++) {
dotv += low[j] * (float)kv_lora_cache[lora_base + j];
}
}
const uint64_t rope_base = (uint64_t)row * qk_rope;
for (uint32_t r = 0; r < qk_rope; r += 2u) {
const float2 y = glm_cache_rope_pair_f16_dev<CT>(
k_rope_cache, rope_base, r, row, qk_rope,
freq_base, freq_scale, ext_factor, attn_factor,
corr_dims);
dotv += qh[qk_nope + r] * y.x +
qh[qk_nope + r + 1u] * y.y;
}
score = dotv * scale;
}
scores[s] = score;
local_max = fmaxf(local_max, score);
}
red[tid] = local_max;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]);
__syncthreads();
}
const float max_score = red[0];
__syncthreads();
float local_sum = 0.0f;
for (uint32_t s = tid; s < row_count; s += nth) {
const float w = expf(scores[s] - max_score);
scores[s] = w;
local_sum += w;
}
red[tid] = local_sum;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) red[tid] += red[tid + step];
__syncthreads();
}
const uint64_t head_index = (uint64_t)token * n_head + head;
if (tid == 0u) denom[head_index] = fmaxf(red[0], 1.0e-20f);
float *head_weights = weights + head_index * score_stride;
for (uint32_t s = tid; s < row_count; s += nth) {
head_weights[s] = scores[s];
}
}
template <typename CT, bool RANGE_TOK2 = false>
__global__ static void glm_attention_decode_lora_staged_kernel(
float *lora_sum,
const float *scores,
const float *denom,
const CT *kv_lora_cache,
const uint32_t *selected,
uint32_t n_selected,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_lora_dim) {
const uint32_t head = blockIdx.y;
const uint32_t token = RANGE_TOK2 ? blockIdx.z : 0u;
const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u);
const uint32_t score_stride = n_selected + (RANGE_TOK2 ? 1u : 0u);
const uint32_t pair = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t j = pair * 2u;
if (head >= n_head || j >= kv_lora_dim) return;
const uint64_t head_index = (uint64_t)token * n_head + head;
const float *head_scores = scores + head_index * score_stride;
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t s = 0; s < row_count; s++) {
const uint32_t row = RANGE_TOK2 ? s : selected[s];
if (row < cache_cap) {
const float2 v = glm_cache_value_pair_dev(
kv_lora_cache + (uint64_t)row * kv_lora_dim + j);
const float w = head_scores[s];
acc0 += w * v.x;
acc1 += w * v.y;
}
}
float *out = lora_sum + head_index * kv_lora_dim + j;
out[0] = acc0 / denom[head_index];
out[1] = acc1 / denom[head_index];
}
template <bool TOK2 = false>
__global__ static void glm_attention_decode_value_staged_kernel(
float *heads,
const float *lora_sum,
const char *value_weight,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim,
uint32_t value_row_bytes) {
const uint32_t head = blockIdx.y;
const uint32_t token = TOK2 ? blockIdx.z : 0u;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t nwarps = blockDim.x >> 5;
const uint32_t out_warp = blockIdx.x * nwarps + warp;
const uint32_t total_warps = gridDim.x * nwarps;
if (head >= n_head) return;
const float *low = lora_sum +
((uint64_t)token * n_head + head) * kv_lora_dim;
float *out = heads +
((uint64_t)token * n_head + head) * value_dim;
const uint32_t nblk = kv_lora_dim >> 5;
for (uint32_t d = out_warp; d < value_dim; d += total_warps) {
const char *row = value_weight +
((uint64_t)head * value_dim + d) * value_row_bytes;
float acc = 0.0f;
for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) {
const char *b = row + (uint64_t)blk * 34u;
const float dscale = __half2float(*(const __half *)b);
const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u;
const float *xs = low + blk * 32u + (lane & 1u) * 16u;
float s = 0.0f;
#pragma unroll
for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k];
acc += dscale * s;
}
for (int off = 16; off > 0; off >>= 1) {
acc += __shfl_down_sync(0xffffffffu, acc, off);
}
if (lane == 0u) out[d] = acc;
}
}
/* Single-token indexed MLA decode attention, one block per head.
* Fuses score (qk_low . kv_lora + q_rope . rope(k_rope@row)), softmax over
* the indexer-selected rows, the weighted kv_lora sum, and the per-head
* value projection (q8_0). Dynamic shared: red[256] + scores[n_selected] +
* lora_sum[kv_lora_dim]. */
template <typename CT, bool RANGE_TOK2 = false>
__global__ static void glm_attention_indexed_decode_kernel(
float *heads,
const float *q,
const float *qk_low,
const CT *kv_lora_cache,
const CT *k_rope_cache,
const char *value_weight,
const uint32_t *selected,
uint32_t n_selected,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t value_row_bytes,
bool lora_vec2,
bool score_vec2,
float scale,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint32_t head = blockIdx.x;
const uint32_t token = RANGE_TOK2 ? blockIdx.y : 0u;
const uint32_t row_count = n_selected + (RANGE_TOK2 ? token : 0u);
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
if (head >= n_head || row_count == 0u) return;
const uint32_t qk_dim = qk_nope + qk_rope;
extern __shared__ float glm_dec_sh[];
float *red = glm_dec_sh;
float *scores = glm_dec_sh + 256u;
float *lora_sum = scores + row_count;
const float *qh = q +
((uint64_t)token * n_head + head) * qk_dim;
const float *low = qk_low +
((uint64_t)token * n_head + head) * kv_lora_dim;
float corr_dims[2] = {0.0f, 0.0f};
if (qk_rope != 0u && ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev((int)qk_rope,
(int)n_ctx_orig, beta_fast, freq_base)));
corr_dims[1] = fminf((float)qk_rope - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev((int)qk_rope,
(int)n_ctx_orig, beta_slow, freq_base)));
}
float local_max = -FLT_MAX;
for (uint32_t s = tid; s < row_count; s += nth) {
const uint32_t row = RANGE_TOK2 ? s : selected[s];
float score = -FLT_MAX;
if (row < cache_cap) {
float dotv = 0.0f;
const uint64_t lora_base = (uint64_t)row * kv_lora_dim;
if (score_vec2) {
for (uint32_t j = 0; j < kv_lora_dim; j += 2u) {
const float2 x = *(const float2 *)(low + j);
const float2 y = glm_cache_value_pair_dev(
kv_lora_cache + lora_base + j);
dotv += x.x * y.x;
dotv += x.y * y.y;
}
} else {
for (uint32_t j = 0; j < kv_lora_dim; j++) {
dotv += low[j] * (float)kv_lora_cache[lora_base + j];
}
}
const uint64_t rope_base = (uint64_t)row * qk_rope;
for (uint32_t r = 0; r < qk_rope; r += 2u) {
const float2 y = glm_cache_rope_pair_f16_dev<CT>(
k_rope_cache, rope_base, r, row, qk_rope,
freq_base, freq_scale, ext_factor, attn_factor,
corr_dims);
dotv += qh[qk_nope + r] * y.x + qh[qk_nope + r + 1u] * y.y;
}
score = dotv * scale;
}
scores[s] = score;
local_max = fmaxf(local_max, score);
}
red[tid] = local_max;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]);
__syncthreads();
}
const float max_score = red[0];
__syncthreads();
float local_sum = 0.0f;
for (uint32_t s = tid; s < row_count; s += nth) {
const float w = expf(scores[s] - max_score);
scores[s] = w;
local_sum += w;
}
red[tid] = local_sum;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) red[tid] += red[tid + step];
__syncthreads();
}
const float denom = fmaxf(red[0], 1.0e-20f);
__syncthreads();
if (lora_vec2) {
for (uint32_t j = tid * 2u; j < kv_lora_dim; j += nth * 2u) {
float acc0 = 0.0f;
float acc1 = 0.0f;
for (uint32_t s = 0; s < row_count; s++) {
const uint32_t row = RANGE_TOK2 ? s : selected[s];
if (row < cache_cap) {
const float2 v = glm_cache_value_pair_dev(
kv_lora_cache + (uint64_t)row * kv_lora_dim + j);
const float w = scores[s];
acc0 += w * v.x;
acc1 += w * v.y;
}
}
lora_sum[j] = acc0 / denom;
lora_sum[j + 1u] = acc1 / denom;
}
} else {
for (uint32_t j = tid; j < kv_lora_dim; j += nth) {
float acc = 0.0f;
for (uint32_t s = 0; s < row_count; s++) {
const uint32_t row = RANGE_TOK2 ? s : selected[s];
if (row < cache_cap) {
acc += scores[s] *
(float)kv_lora_cache[(uint64_t)row * kv_lora_dim + j];
}
}
lora_sum[j] = acc / denom;
}
}
__syncthreads();
float *out = heads +
((uint64_t)token * n_head + head) * value_dim;
/* Warp-cooperative value projection: one warp per output dim, two
* lanes per q8_0 block (16 cols each). */
const uint32_t nwarps = nth >> 5;
const uint32_t warp = tid >> 5;
const uint32_t lane = tid & 31u;
const uint32_t nblk = kv_lora_dim >> 5;
for (uint32_t d = warp; d < value_dim; d += nwarps) {
const char *row = value_weight +
((uint64_t)head * value_dim + d) * value_row_bytes;
float acc = 0.0f;
for (uint32_t blk = lane >> 1; blk < nblk; blk += 16u) {
const char *b = row + (uint64_t)blk * 34u;
const float dscale = __half2float(*(const __half *)b);
const int8_t *q = (const int8_t *)(b + 2) + (lane & 1u) * 16u;
const float *xs = lora_sum + blk * 32u + (lane & 1u) * 16u;
float s = 0.0f;
#pragma unroll
for (int k = 0; k < 16; k++) s += (float)q[k] * xs[k];
acc += dscale * s;
}
for (int off = 16; off > 0; off >>= 1) {
acc += __shfl_down_sync(0xffffffffu, acc, off);
}
if (lane == 0u) out[d] = acc;
}
}
extern "C" int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint32_t qk_dim = qk_nope + qk_rope;
if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
!model_map || !selected ||
n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
n_head == 0 || kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u ||
qk_nope == 0 || (qk_rope & 1u) != 0u ||
value_dim == 0) {
return 0;
}
if (value_weight_type != 8u) { /* DS4_TENSOR_Q8_0 */
fprintf(stderr,
"ds4: glm indexed decode attention: unsupported value type %u\n",
value_weight_type);
return 0;
}
const uint64_t value_row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u;
const uint64_t value_weight_bytes =
(uint64_t)n_head * value_dim * value_row_bytes;
if (value_weight_offset > model_size ||
value_weight_bytes > model_size - value_weight_offset) {
return 0;
}
const uint64_t cache_elem = cache_f16 ? 2u : 4u;
if (heads->bytes < (uint64_t)n_head * value_dim * sizeof(float) ||
q->bytes < (uint64_t)n_head * qk_dim * sizeof(float) ||
qk_low->bytes < (uint64_t)n_head * kv_lora_dim * sizeof(float) ||
kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem ||
k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem ||
selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const char *vw = cuda_resolve_weight_ptr(model_map, value_weight_offset,
value_weight_bytes, logical_tier, "glm_v_b_decode");
if (!vw) return 0;
const float scale = 1.0f / sqrtf((float)qk_dim);
const bool score_vec2 = getenv("DS4_GLM_ATTN_NO_SCORE_VEC2") == NULL;
const bool range_tok2 =
g_glm_mtp_verify_mode &&
getenv("DS4_GLM_MTP_NO_ATTN_TOK2") == NULL &&
n_selected < cache_cap &&
heads->bytes >= 2u * (uint64_t)n_head * value_dim * sizeof(float) &&
q->bytes >= 2u * (uint64_t)n_head * qk_dim * sizeof(float) &&
qk_low->bytes >=
2u * (uint64_t)n_head * kv_lora_dim * sizeof(float);
if (range_tok2 && n_selected < 512u) {
const bool lora_vec2 =
getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL;
const uint32_t shmem =
(256u + n_selected + 1u + kv_lora_dim) *
(uint32_t)sizeof(float);
const dim3 grid(n_head, 2u, 1u);
if (cache_f16) {
glm_attention_indexed_decode_kernel<__half, true>
<<<grid, 256, shmem>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
vw, (const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, value_dim,
(uint32_t)value_row_bytes,
lora_vec2, score_vec2, scale, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
} else {
glm_attention_indexed_decode_kernel<float, true>
<<<grid, 256, shmem>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
vw, (const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, value_dim,
(uint32_t)value_row_bytes,
lora_vec2, score_vec2, scale, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
}
return cuda_ok(cudaGetLastError(),
"glm indexed decode attention tok2 range");
}
if (n_selected >= 512u &&
getenv("DS4_GLM_ATTN_NO_STAGED_DECODE") == NULL) {
const uint32_t token_count = range_tok2 ? 2u : 1u;
const uint32_t score_stride = n_selected + (range_tok2 ? 1u : 0u);
const uint64_t head_count = (uint64_t)token_count * n_head;
if (head_count > UINT64_MAX / score_stride ||
head_count > UINT64_MAX / kv_lora_dim) {
return 0;
}
const uint64_t score_count = head_count * score_stride;
const uint64_t lora_count = head_count * kv_lora_dim;
if (score_count > UINT64_MAX - head_count - lora_count ||
score_count + head_count + lora_count >
UINT64_MAX / sizeof(float)) {
return 0;
}
const uint64_t scratch_bytes =
(score_count + head_count + lora_count) * sizeof(float);
float *scratch = (float *)cuda_tmp_alloc_on(
ds4_tensor_device_idx(heads), scratch_bytes,
"glm staged decode attention");
if (!scratch) return 0;
float *softmax_denom = scratch + score_count;
float *lora_sum = softmax_denom + head_count;
const uint32_t weight_shmem =
(256u + score_stride) * (uint32_t)sizeof(float);
const dim3 weight_grid(n_head, token_count, 1u);
if (cache_f16 && range_tok2) {
glm_attention_decode_weights_staged_kernel<__half, true>
<<<weight_grid, 256, weight_shmem>>>(
scratch, softmax_denom, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, scale, n_ctx_orig, freq_base,
freq_scale, ext_factor, attn_factor, beta_fast,
beta_slow, score_vec2);
} else if (cache_f16) {
glm_attention_decode_weights_staged_kernel<__half>
<<<weight_grid, 256, weight_shmem>>>(
scratch, softmax_denom, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, scale, n_ctx_orig, freq_base,
freq_scale, ext_factor, attn_factor, beta_fast,
beta_slow, score_vec2);
} else if (range_tok2) {
glm_attention_decode_weights_staged_kernel<float, true>
<<<weight_grid, 256, weight_shmem>>>(
scratch, softmax_denom, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, scale, n_ctx_orig, freq_base,
freq_scale, ext_factor, attn_factor, beta_fast,
beta_slow, score_vec2);
} else {
glm_attention_decode_weights_staged_kernel<float>
<<<weight_grid, 256, weight_shmem>>>(
scratch, softmax_denom, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, scale, n_ctx_orig, freq_base,
freq_scale, ext_factor, attn_factor, beta_fast,
beta_slow, score_vec2);
}
if (!cuda_ok(cudaGetLastError(),
"glm staged decode weights launch")) {
return 0;
}
dim3 lora_grid((kv_lora_dim / 2u + 63u) / 64u,
n_head, token_count);
if (cache_f16 && range_tok2) {
glm_attention_decode_lora_staged_kernel<__half, true>
<<<lora_grid, 64>>>(
lora_sum, scratch, softmax_denom,
(const __half *)kv_lora_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim);
} else if (cache_f16) {
glm_attention_decode_lora_staged_kernel<__half>
<<<lora_grid, 64>>>(
lora_sum, scratch, softmax_denom,
(const __half *)kv_lora_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim);
} else if (range_tok2) {
glm_attention_decode_lora_staged_kernel<float, true>
<<<lora_grid, 64>>>(
lora_sum, scratch, softmax_denom,
(const float *)kv_lora_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim);
} else {
glm_attention_decode_lora_staged_kernel<float>
<<<lora_grid, 64>>>(
lora_sum, scratch, softmax_denom,
(const float *)kv_lora_cache->ptr,
(const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim);
}
if (!cuda_ok(cudaGetLastError(),
"glm staged decode lora launch")) {
return 0;
}
dim3 value_grid((value_dim + 127u) / 128u,
n_head, token_count);
if (range_tok2) {
glm_attention_decode_value_staged_kernel<true>
<<<value_grid, 128>>>(
(float *)heads->ptr, lora_sum, vw,
n_head, kv_lora_dim, value_dim,
(uint32_t)value_row_bytes);
} else {
glm_attention_decode_value_staged_kernel<false>
<<<value_grid, 128>>>(
(float *)heads->ptr, lora_sum, vw,
n_head, kv_lora_dim, value_dim,
(uint32_t)value_row_bytes);
}
return cuda_ok(cudaGetLastError(),
"glm staged decode value launch");
}
const bool lora_vec2 = getenv("DS4_GLM_ATTN_NO_LORA_VEC2") == NULL;
const uint32_t shmem =
(256u + n_selected + kv_lora_dim) * (uint32_t)sizeof(float);
if (cache_f16) {
glm_attention_indexed_decode_kernel<__half><<<n_head, 256, shmem>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const __half *)kv_lora_cache->ptr,
(const __half *)k_rope_cache->ptr,
vw, (const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes,
lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
} else {
glm_attention_indexed_decode_kernel<float><<<n_head, 256, shmem>>>(
(float *)heads->ptr, (const float *)q->ptr,
(const float *)qk_low->ptr,
(const float *)kv_lora_cache->ptr,
(const float *)k_rope_cache->ptr,
vw, (const uint32_t *)selected->ptr,
n_selected, cache_cap, n_head, kv_lora_dim,
qk_nope, qk_rope, value_dim, (uint32_t)value_row_bytes,
lora_vec2, score_vec2, scale, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
}
return cuda_ok(cudaGetLastError(), "glm indexed decode attention");
}
__global__ static void glm53_indexer_pool_update_kernel(
void *pool_cache,
float *tail,
const float *raw_k,
const float *gate,
const float *norm_weight,
const float *norm_bias,
const uint16_t *ape,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t pool_size,
float eps,
bool cache_f16) {
const uint32_t tid = threadIdx.x;
if (head_dim == 0u || pool_size == 0u || tid >= head_dim ||
n_tokens == 0u) {
return;
}
float *const tail_k = tail;
float *const tail_gate = tail + (uint64_t)pool_size * head_dim;
extern __shared__ float shared[];
float *rows = shared;
float *mean = rows + (uint64_t)pool_size * head_dim;
float *inv = mean + pool_size;
const uint32_t pool = pos0 / pool_size + blockIdx.x;
const uint32_t pool_start = pool * pool_size;
const uint32_t input_end = pos0 + n_tokens;
if (pool_start >= input_end || pool_start + pool_size <= pos0) return;
const bool complete = pool_start + pool_size <= input_end;
for (uint32_t r = 0; r < pool_size; r++) {
const uint32_t pos = pool_start + r;
float k_value = 0.0f;
float gate_value = 0.0f;
if (pos >= pos0 && pos < input_end) {
const uint32_t src_row = pos - pos0;
k_value = raw_k[(uint64_t)src_row * head_dim + tid];
gate_value = gate[(uint64_t)src_row * head_dim + tid];
if (!complete) {
tail_k[(uint64_t)r * head_dim + tid] = k_value;
tail_gate[(uint64_t)r * head_dim + tid] = gate_value;
}
} else {
k_value = tail_k[(uint64_t)r * head_dim + tid];
gate_value = tail_gate[(uint64_t)r * head_dim + tid];
}
rows[(uint64_t)r * head_dim + tid] = k_value;
}
__syncthreads();
if (!complete || pool >= (cache_cap + pool_size - 1u) / pool_size) {
return;
}
if (tid < pool_size) {
const uint32_t r = tid;
float sum = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) {
sum += rows[(uint64_t)r * head_dim + d];
}
const float m = sum / (float)head_dim;
float ss = 0.0f;
for (uint32_t d = 0; d < head_dim; d++) {
const float delta = rows[(uint64_t)r * head_dim + d] - m;
ss = fmaf(delta, delta, ss);
}
mean[r] = m;
inv[r] = rsqrtf(ss / (float)head_dim + eps);
}
__syncthreads();
float max_logit = -INFINITY;
float logits[4];
for (uint32_t r = 0; r < pool_size; r++) {
const uint32_t pos = pool_start + r;
const float gate_value = pos >= pos0
? gate[(uint64_t)(pos - pos0) * head_dim + tid]
: tail_gate[(uint64_t)r * head_dim + tid];
const float ape_value = __uint_as_float(
(uint32_t)ape[(uint64_t)r * head_dim + tid] << 16);
logits[r] = gate_value + ape_value;
max_logit = fmaxf(max_logit, logits[r]);
}
float denom = 0.0f;
for (uint32_t r = 0; r < pool_size; r++) {
logits[r] = expf(logits[r] - max_logit);
denom += logits[r];
}
float pooled = 0.0f;
for (uint32_t r = 0; r < pool_size; r++) {
const float normalized =
(rows[(uint64_t)r * head_dim + tid] - mean[r]) * inv[r] *
norm_weight[tid] + norm_bias[tid];
pooled = fmaf(logits[r] / denom, normalized, pooled);
}
const uint64_t dst_index = (uint64_t)pool * head_dim + tid;
if (cache_f16) {
((__half *)pool_cache)[dst_index] = __float2half(pooled);
} else {
((float *)pool_cache)[dst_index] = pooled;
}
}
extern "C" int ds4_gpu_glm53_indexer_pool_update_tensor(
ds4_gpu_tensor *pool_cache,
ds4_gpu_tensor *tail_k,
ds4_gpu_tensor *tail_gate,
const ds4_gpu_tensor *raw_k,
const ds4_gpu_tensor *gate,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint64_t norm_bias_offset,
uint64_t ape_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t pool_size,
float eps,
bool cache_f16) {
if (!pool_cache || !tail_k || !tail_gate || !raw_k || !gate ||
!model_map || n_tokens == 0u || cache_cap == 0u || head_dim == 0u ||
head_dim > 1024u || pool_size != 4u || !isfinite(eps) || eps <= 0.0f ||
pos0 > cache_cap || n_tokens > cache_cap - pos0) {
return 0;
}
const uint64_t pool_cap =
((uint64_t)cache_cap + pool_size - 1u) / pool_size;
const uint64_t cache_elem_size = cache_f16 ? sizeof(__half) : sizeof(float);
if (!glm53_cuda_tensor_has(pool_cache, pool_cap * head_dim,
cache_elem_size) ||
!glm53_cuda_tensor_has(tail_k, 2u * (uint64_t)pool_size * head_dim,
sizeof(float)) ||
!glm53_cuda_tensor_has(tail_gate, (uint64_t)pool_size * head_dim,
sizeof(float)) ||
(char *)tail_gate->ptr !=
(char *)tail_k->ptr +
(uint64_t)pool_size * head_dim * sizeof(float) ||
!glm53_cuda_tensor_has(raw_k, (uint64_t)n_tokens * head_dim,
sizeof(float)) ||
!glm53_cuda_tensor_has(gate, (uint64_t)n_tokens * head_dim,
sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 indexer pool received invalid buffers\n");
return 0;
}
const int tier = ds4_tensor_device_idx(pool_cache);
const float *norm_weight = glm53_cuda_weight_f32(
model_map, model_size, norm_weight_offset, head_dim, tier,
"indexer pool norm weight");
const float *norm_bias = glm53_cuda_weight_f32(
model_map, model_size, norm_bias_offset, head_dim, tier,
"indexer pool norm bias");
const uint64_t ape_bytes =
(uint64_t)pool_size * head_dim * sizeof(uint16_t);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset) {
fprintf(stderr, "ds4: GLM-5.3 indexer pool APE range is invalid\n");
return 0;
}
const uint16_t *ape = (const uint16_t *)cuda_resolve_weight_ptr(
model_map, ape_offset, ape_bytes, tier, "indexer pool APE");
if (!norm_weight || !norm_bias || !ape) return 0;
const uint64_t shared_bytes =
((uint64_t)pool_size * head_dim + 2u * pool_size) * sizeof(float);
uint32_t done = 0u;
const uint32_t leading = pos0 % pool_size;
if (leading != 0u) {
uint32_t chunk = pool_size - leading;
if (chunk > n_tokens) chunk = n_tokens;
glm53_indexer_pool_update_kernel<<<1u, head_dim, shared_bytes,
cuda_decode_stream()>>>(
pool_cache->ptr, (float *)tail_k->ptr,
(const float *)raw_k->ptr,
(const float *)gate->ptr, norm_weight, norm_bias, ape,
pos0, chunk, cache_cap, head_dim, pool_size, eps, cache_f16);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 leading indexer pool update launch")) {
return 0;
}
done = chunk;
}
const uint32_t remaining = n_tokens - done;
const uint32_t full_tokens = remaining - remaining % pool_size;
if (full_tokens != 0u) {
const uint32_t groups = full_tokens / pool_size;
glm53_indexer_pool_update_kernel<<<groups, head_dim, shared_bytes,
cuda_decode_stream()>>>(
pool_cache->ptr, (float *)tail_k->ptr,
(const float *)raw_k->ptr + (uint64_t)done * head_dim,
(const float *)gate->ptr + (uint64_t)done * head_dim,
norm_weight, norm_bias, ape, pos0 + done, full_tokens, cache_cap,
head_dim, pool_size, eps, cache_f16);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 full indexer pool update launch")) {
return 0;
}
done += full_tokens;
}
if (done < n_tokens) {
const uint32_t tail_tokens = n_tokens - done;
glm53_indexer_pool_update_kernel<<<1u, head_dim, shared_bytes,
cuda_decode_stream()>>>(
pool_cache->ptr, (float *)tail_k->ptr,
(const float *)raw_k->ptr + (uint64_t)done * head_dim,
(const float *)gate->ptr + (uint64_t)done * head_dim,
norm_weight, norm_bias, ape, pos0 + done, tail_tokens,
cache_cap, head_dim, pool_size, eps, cache_f16);
if (!cuda_ok(cudaGetLastError(),
"GLM-5.3 trailing indexer pool update launch")) {
return 0;
}
}
return 1;
}
__global__ static void glm53_expand_pool_selection_kernel(
uint32_t *raw_selected,
const uint32_t *pool_selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t selected_pools,
uint32_t index_topk,
uint32_t pool_size,
uint32_t output_width) {
const uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
const uint64_t total = (uint64_t)n_tokens * output_width;
if (gid >= total || output_width == 0u || pool_size == 0u) return;
const uint32_t token = (uint32_t)(gid / output_width);
const uint32_t slot = (uint32_t)(gid - (uint64_t)token * output_width);
uint32_t value = UINT32_MAX;
if (slot < index_topk) {
const uint32_t pool_slot = slot / pool_size;
if (pool_slot < selected_pools) {
const uint32_t pool = pool_selected[
(uint64_t)token * selected_pools + pool_slot];
value = pool * pool_size + slot % pool_size;
}
} else {
const uint32_t tail_slot = slot - index_topk;
const uint32_t visible = pos0 + token + 1u;
const uint32_t tail_count = visible % pool_size;
if (tail_slot < tail_count) {
value = visible - tail_count + tail_slot;
}
}
raw_selected[gid] = value;
}
extern "C" int ds4_gpu_glm53_expand_pool_selection_tensor(
ds4_gpu_tensor *raw_selected,
const ds4_gpu_tensor *pool_selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t selected_pools,
uint32_t index_topk,
uint32_t pool_size,
uint32_t output_width) {
if (!raw_selected || !pool_selected || n_tokens == 0u ||
selected_pools == 0u || index_topk == 0u || pool_size == 0u ||
selected_pools > index_topk / pool_size ||
output_width < index_topk + pool_size - 1u ||
!glm53_cuda_tensor_has(raw_selected,
(uint64_t)n_tokens * output_width, sizeof(uint32_t)) ||
!glm53_cuda_tensor_has(pool_selected,
(uint64_t)n_tokens * selected_pools, sizeof(uint32_t))) {
return 0;
}
const uint64_t total = (uint64_t)n_tokens * output_width;
glm53_expand_pool_selection_kernel<<<
(unsigned)((total + 255u) / 256u), 256u, 0,
cuda_decode_stream()>>>(
(uint32_t *)raw_selected->ptr,
(const uint32_t *)pool_selected->ptr, n_tokens, pos0,
selected_pools, index_topk, pool_size, output_width);
return cuda_ok(cudaGetLastError(),
"GLM-5.3 pool selection expansion launch");
}
template <typename CT>
__device__ __forceinline__ static CT glm_cache_cast(float value) {
return (CT)value;
}
template <>
__device__ __forceinline__ __half glm_cache_cast<__half>(float value) {
return __float2half(value);
}
template <typename CT>
__global__ static void glm_build_kv_cache_kernel(
CT *key_cache,
CT *value_cache,
const float *kv_raw,
const float *k_nope,
const float *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint32_t token = blockIdx.x;
const uint32_t head = blockIdx.y;
const uint32_t tid = threadIdx.x;
if (token >= n_tokens || head >= n_head) return;
const uint32_t qk_dim = qk_nope + qk_rope;
const uint32_t pos = pos0 + token;
const float *raw = kv_raw + (uint64_t)token * kv_raw_dim;
const float *kn = k_nope +
((uint64_t)token * n_head + head) * qk_nope;
const float *val = value +
((uint64_t)token * n_head + head) * value_dim;
const uint64_t kbase = ((uint64_t)pos * n_head + head) * qk_dim;
const uint64_t vbase = ((uint64_t)pos * n_head + head) * value_dim;
for (uint32_t i = tid; i < qk_nope; i += blockDim.x) {
key_cache[kbase + i] = glm_cache_cast<CT>(kn[i]);
}
float corr_dims[2] = {0.0f, 0.0f};
if (qk_rope != 0u && ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev(
(int)qk_rope, (int)n_ctx_orig, beta_fast, freq_base)));
corr_dims[1] = fminf((float)qk_rope - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev(
(int)qk_rope, (int)n_ctx_orig, beta_slow, freq_base)));
}
const float inv_dims = qk_rope != 0u ? -1.0f / (float)qk_rope : 0.0f;
for (uint32_t r = tid * 2u; r < qk_rope; r += blockDim.x * 2u) {
const float theta = (float)pos *
powf(freq_base, inv_dims * (float)r);
float cos_theta = 0.0f;
float sin_theta = 0.0f;
glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)r,
ext_factor, attn_factor,
&cos_theta, &sin_theta);
const float x0 = raw[kv_lora_dim + r];
const float x1 = raw[kv_lora_dim + r + 1u];
key_cache[kbase + qk_nope + r] = glm_cache_cast<CT>(
x0 * cos_theta - x1 * sin_theta);
key_cache[kbase + qk_nope + r + 1u] = glm_cache_cast<CT>(
x0 * sin_theta + x1 * cos_theta);
}
for (uint32_t i = tid; i < value_dim; i += blockDim.x) {
value_cache[vbase + i] = glm_cache_cast<CT>(val[i]);
}
}
static int glm_build_kv_cache_launch(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
const uint32_t qk_dim = qk_nope + qk_rope;
const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float);
if (!key_cache || !value_cache || !kv_raw || !k_nope || !value ||
n_tokens == 0u || cache_cap == 0u || n_head == 0u ||
kv_raw_dim == 0u || kv_lora_dim == 0u || qk_nope == 0u ||
value_dim == 0u || kv_lora_dim > kv_raw_dim ||
qk_rope > kv_raw_dim - kv_lora_dim || qk_dim < qk_nope ||
(qk_rope & 1u) != 0u || pos0 > cache_cap ||
n_tokens > cache_cap - pos0 ||
(qk_rope != 0u && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0u && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow) ||
key_cache->bytes <
(uint64_t)cache_cap * n_head * qk_dim * cache_elem ||
value_cache->bytes <
(uint64_t)cache_cap * n_head * value_dim * cache_elem ||
kv_raw->bytes <
(uint64_t)n_tokens * kv_raw_dim * sizeof(float) ||
k_nope->bytes <
(uint64_t)n_tokens * n_head * qk_nope * sizeof(float) ||
value->bytes <
(uint64_t)n_tokens * n_head * value_dim * sizeof(float)) {
return 0;
}
const dim3 grid(n_tokens, n_head, 1u);
if (cache_f16) {
glm_build_kv_cache_kernel<__half><<<grid, 256u, 0,
cuda_decode_stream()>>>(
(__half *)key_cache->ptr, (__half *)value_cache->ptr,
(const float *)kv_raw->ptr, (const float *)k_nope->ptr,
(const float *)value->ptr, pos0, n_tokens, n_head,
kv_raw_dim, kv_lora_dim, qk_nope, qk_rope, value_dim,
n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
} else {
glm_build_kv_cache_kernel<float><<<grid, 256u, 0,
cuda_decode_stream()>>>(
(float *)key_cache->ptr, (float *)value_cache->ptr,
(const float *)kv_raw->ptr, (const float *)k_nope->ptr,
(const float *)value->ptr, pos0, n_tokens, n_head,
kv_raw_dim, kv_lora_dim, qk_nope, qk_rope, value_dim,
n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow);
}
return cuda_ok(cudaGetLastError(), "GLM KV cache build launch");
}
extern "C" int ds4_gpu_glm_build_kv_cache_flash_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
if (pos0 != 0u) return 0;
return glm_build_kv_cache_launch(
key_cache, value_cache, kv_raw, k_nope, value, pos0, n_tokens,
cache_cap, n_head, kv_raw_dim, kv_lora_dim, qk_nope, qk_rope,
value_dim, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, cache_f16);
}
extern "C" int ds4_gpu_glm_build_kv_cache_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
return glm_build_kv_cache_launch(
key_cache, value_cache, kv_raw, k_nope, value, pos0, n_tokens,
cache_cap, n_head, kv_raw_dim, kv_lora_dim, qk_nope, qk_rope,
value_dim, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow, cache_f16);
}
__global__ static void glm_fill_selected_range_batch_kernel(
uint32_t *selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t pad_row) {
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
const uint32_t total = n_tokens * n_selected;
if (gid >= total || n_selected == 0u) return;
const uint32_t token = gid / n_selected;
const uint32_t slot = gid - token * n_selected;
const uint32_t visible = pos0 + token + 1u;
selected[gid] = slot < visible ? slot : pad_row;
}
extern "C" int ds4_gpu_glm_fill_selected_range_batch_tensor(
ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t pad_row) {
if (!selected || n_tokens == 0 || n_selected == 0 ||
selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t)) {
return 0;
}
const uint64_t total = (uint64_t)n_tokens * n_selected;
glm_fill_selected_range_batch_kernel<<<(unsigned)((total + 255) / 256), 256>>>(
(uint32_t *)selected->ptr, n_tokens, pos0, n_selected, pad_row);
return cuda_ok(cudaGetLastError(), "glm fill selected batch launch");
}
__global__ static void glm_fill_selected_range_kernel(
uint32_t *selected, uint32_t n_selected) {
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
if (gid < n_selected) selected[gid] = gid;
}
extern "C" int ds4_gpu_glm_fill_selected_range_tensor(
ds4_gpu_tensor *selected,
uint32_t n_selected) {
if (!selected || n_selected == 0 ||
selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) {
return 0;
}
glm_fill_selected_range_kernel<<<(n_selected + 255) / 256, 256>>>(
(uint32_t *)selected->ptr, n_selected);
return cuda_ok(cudaGetLastError(), "glm fill selected launch");
}
static int glm_rope_tail_offset_launch(
ds4_gpu_tensor *x,
uint32_t n_tokens, uint32_t n_head, uint32_t head_dim,
uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0,
uint32_t n_ctx_orig, float freq_base, float freq_scale,
float ext_factor, float attn_factor,
float beta_fast, float beta_slow, const char *what);
extern "C" int ds4_gpu_glm_indexer_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim,
rot_dim, 0, pos0, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow, "glm indexer rope tail");
}
template <typename CT>
__global__ static void glm_indexer_scores_f32_kernel(
float *scores,
const float *q,
const float *weights,
const CT *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t row_group_size,
float scale,
bool causal) {
const uint32_t row = blockIdx.x;
const uint32_t token = blockIdx.y;
const uint32_t tid = threadIdx.x;
if (row >= n_rows || token >= n_tokens || tid >= 128u) return;
const uint32_t visible = (pos0 + token + 1u) / row_group_size;
if (causal && row >= min(n_rows, visible)) {
if (tid == 0u) scores[(uint64_t)token * n_rows + row] = -INFINITY;
return;
}
__shared__ float partial[128];
float total = 0.0f;
const CT *krow = indexer_key_cache + (uint64_t)row * head_dim;
for (uint32_t h = 0; h < n_head; h++) {
const float *qh = q +
((uint64_t)token * n_head + h) * head_dim;
float dot = tid < head_dim ? qh[tid] * (float)krow[tid] : 0.0f;
partial[tid] = dot;
__syncthreads();
for (uint32_t stride = 64u; stride > 0u; stride >>= 1u) {
if (tid < stride) partial[tid] += partial[tid + stride];
__syncthreads();
}
if (tid == 0u) {
total += fmaxf(partial[0], 0.0f) *
weights[(uint64_t)token * n_head + h];
}
__syncthreads();
}
if (tid == 0u) {
scores[(uint64_t)token * n_rows + row] = total * scale;
}
}
/* 16-token x 128-row indexer tile. Q and cached K are staged as fp16,
* matching the model's compact-cache precision; each head's MMA result and
* the weighted head reduction remain fp32. */
template <typename CT>
__global__ static void glm_indexer_scores_wmma128_kernel(
float *scores,
const float *q,
const float *weights,
const CT *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t row_group_size,
float scale,
bool causal) {
#if __CUDA_ARCH__ >= 700
namespace wmma = nvcuda::wmma;
const uint32_t row0 = blockIdx.x * 128u;
const uint32_t token0 = blockIdx.y * 16u;
const uint32_t tid = threadIdx.x;
const uint32_t warp = tid >> 5u;
if (tid >= 256u || head_dim != 128u) return;
if (causal) {
const uint32_t last_token = min(token0 + 16u, n_tokens);
const uint32_t max_visible = last_token > token0
? min((pos0 + last_token) / row_group_size, n_rows) : 0u;
if (row0 >= max_visible) {
for (uint32_t i = tid; i < 16u * 128u; i += 256u) {
const uint32_t token = token0 + (i >> 7u);
const uint32_t row = row0 + (i & 127u);
if (token < n_tokens && row < n_rows) {
scores[(uint64_t)token * n_rows + row] = -INFINITY;
}
}
return;
}
}
__shared__ __half q_sh[16 * 128];
__shared__ __half k_sh[128 * 128];
__shared__ float dot_sh[8 * 16 * 16];
float acc[8] = {0.0f, 0.0f, 0.0f, 0.0f,
0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t i = tid; i < 128u * 128u; i += 256u) {
const uint32_t r = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t row = row0 + r;
const float v = row < n_rows
? (float)indexer_key_cache[(uint64_t)row * head_dim + d]
: 0.0f;
k_sh[d + r * 128u] = __float2half(v);
}
__syncthreads();
for (uint32_t h = 0; h < n_head; h++) {
for (uint32_t i = tid; i < 16u * 128u; i += 256u) {
const uint32_t tr = i >> 7u;
const uint32_t d = i & 127u;
const uint32_t token = token0 + tr;
const float v = token < n_tokens
? q[((uint64_t)token * n_head + h) * head_dim + d]
: 0.0f;
q_sh[i] = __float2half(v);
}
__syncthreads();
wmma::fragment<wmma::matrix_a, 16, 16, 16,
__half, wmma::row_major> q_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16,
__half, wmma::col_major> k_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> dot_frag;
wmma::fill_fragment(dot_frag, 0.0f);
const uint32_t col0 = warp * 16u;
for (uint32_t k0 = 0; k0 < 128u; k0 += 16u) {
wmma::load_matrix_sync(q_frag, q_sh + k0, 128);
wmma::load_matrix_sync(k_frag,
k_sh + col0 * 128u + k0, 128);
wmma::mma_sync(dot_frag, q_frag, k_frag, dot_frag);
}
wmma::store_matrix_sync(dot_sh + warp * 16u * 16u,
dot_frag, 16, wmma::mem_row_major);
__syncthreads();
const uint32_t local0 = tid & 255u;
const uint32_t token = token0 + (local0 >> 4u);
const float w = token < n_tokens
? weights[(uint64_t)token * n_head + h] : 0.0f;
uint32_t slot = 0;
for (uint32_t i = tid; i < 8u * 16u * 16u;
i += 256u, slot++) {
const uint32_t row = row0 + (i >> 8u) * 16u + (i & 15u);
if (token < n_tokens && row < n_rows) {
acc[slot] += fmaxf(dot_sh[i], 0.0f) * w;
}
}
__syncthreads();
}
uint32_t slot = 0;
for (uint32_t i = tid; i < 8u * 16u * 16u;
i += 256u, slot++) {
const uint32_t local = i & 255u;
const uint32_t token = token0 + (local >> 4u);
const uint32_t row = row0 + (i >> 8u) * 16u + (local & 15u);
if (token < n_tokens && row < n_rows) {
float out = acc[slot] * scale;
if (causal &&
row >= (pos0 + token + 1u) / row_group_size) {
out = -INFINITY;
}
scores[(uint64_t)token * n_rows + row] = out;
}
}
#endif
}
static int glm_indexer_scores_launch(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t row_group_size,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16,
bool causal) {
const uint64_t cache_elem = cache_f16 ? sizeof(__half) : sizeof(float);
if (!scores || !q || !weights || !indexer_key_cache || n_rows == 0u ||
n_tokens == 0u || row_group_size == 0u || n_head == 0u ||
head_dim != 128u || !isfinite(scale) || scale <= 0.0f ||
pos0 > UINT32_MAX - n_tokens ||
(causal && (uint64_t)(pos0 + n_tokens) / row_group_size > n_rows) ||
q->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float) ||
weights->bytes < (uint64_t)n_tokens * n_head * sizeof(float) ||
indexer_key_cache->bytes < (uint64_t)n_rows * head_dim * cache_elem ||
scores->bytes < (uint64_t)n_tokens * n_rows * sizeof(float)) {
return 0;
}
if (!g_quality_mode) {
dim3 grid((n_rows + 127u) / 128u,
(n_tokens + 15u) / 16u, 1);
if (cache_f16) {
glm_indexer_scores_wmma128_kernel<__half><<<grid, 256>>>(
(float *)scores->ptr, (const float *)q->ptr,
(const float *)weights->ptr,
(const __half *)indexer_key_cache->ptr,
n_rows, n_tokens, pos0, n_head, head_dim, row_group_size,
scale, causal);
} else {
glm_indexer_scores_wmma128_kernel<float><<<grid, 256>>>(
(float *)scores->ptr, (const float *)q->ptr,
(const float *)weights->ptr,
(const float *)indexer_key_cache->ptr,
n_rows, n_tokens, pos0, n_head, head_dim, row_group_size,
scale, causal);
}
return cuda_ok(cudaGetLastError(), "glm indexer scores wmma launch");
}
dim3 grid(n_rows, n_tokens, 1);
if (cache_f16) {
glm_indexer_scores_f32_kernel<__half><<<grid, 128>>>(
(float *)scores->ptr, (const float *)q->ptr,
(const float *)weights->ptr,
(const __half *)indexer_key_cache->ptr,
n_rows, n_tokens, pos0, n_head, head_dim, row_group_size,
scale, causal);
} else {
glm_indexer_scores_f32_kernel<float><<<grid, 128>>>(
(float *)scores->ptr, (const float *)q->ptr,
(const float *)weights->ptr,
(const float *)indexer_key_cache->ptr,
n_rows, n_tokens, pos0, n_head, head_dim, row_group_size,
scale, causal);
}
return cuda_ok(cudaGetLastError(), "glm indexer scores f32 launch");
}
extern "C" int ds4_gpu_glm_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache,
n_rows, 1u, 0u, 1u, n_head, head_dim,
scale, cache_f16, false);
}
extern "C" int ds4_gpu_glm_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache,
n_rows, n_tokens, pos0, 1u, n_head,
head_dim, scale, cache_f16, true);
}
extern "C" int ds4_gpu_glm53_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t pool_size,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
if (pool_size != 4u) {
fprintf(stderr, "ds4: GLM-5.3 indexer requires pool size 4\n");
return 0;
}
return glm_indexer_scores_launch(scores, q, weights, indexer_key_cache,
n_rows, n_tokens, pos0, pool_size,
n_head, head_dim, scale, cache_f16, true);
}
__global__ static void glm_k_b_project_q8_0_kernel(
float *out,
const char *weight,
const float *kv_norm,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head,
uint64_t row_bytes) {
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t head = blockIdx.x;
const uint32_t token = blockIdx.y;
const uint32_t q = blockIdx.z * 8u + warp;
float sum = 0.0f;
if (head < n_head && token < n_tokens && q < qk_nope) {
const float *kv = kv_norm + (uint64_t)token * kv_lora_dim;
const uint32_t block = q >> 5u;
const uint32_t slot = q & 31u;
for (uint32_t j = lane; j < kv_lora_dim; j += 32u) {
const char *row = weight +
((uint64_t)head * kv_lora_dim + j) * row_bytes;
const char *blk = row + (uint64_t)block * 34u;
const float d = __half2float(*(const __half *)blk);
const int8_t quant = ((const int8_t *)(blk + 2))[slot];
sum = fmaf(d * (float)quant, kv[j], sum);
}
}
sum = warp_sum_f32(sum);
if (lane == 0u && head < n_head && token < n_tokens && q < qk_nope) {
out[((uint64_t)token * n_head + head) * qk_nope + q] = sum;
}
}
extern "C" int ds4_gpu_glm_k_b_project_typed_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head) {
if (!out || !kv_norm || !model_map || n_tokens == 0u ||
kv_lora_dim == 0u || qk_nope == 0u || n_head == 0u ||
weight_type != 8u || (qk_nope & 31u) != 0u ||
kv_norm->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float) ||
out->bytes <
(uint64_t)n_tokens * n_head * qk_nope * sizeof(float)) {
if (weight_type != 8u) {
fprintf(stderr, "ds4: CUDA GLM k_b projection requires Q8_0\n");
}
return 0;
}
const uint64_t row_bytes = ((uint64_t)qk_nope / 32u) * 34u;
const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
if (weight_rows > UINT64_MAX / row_bytes) return 0;
const uint64_t weight_bytes = weight_rows * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
return 0;
}
const int tier = ds4_tensor_device_idx(out);
const char *weight = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, tier, "GLM k_b Q8_0");
if (!weight) return 0;
const dim3 grid(n_head, n_tokens, (qk_nope + 7u) / 8u);
glm_k_b_project_q8_0_kernel<<<grid, 256u, 0,
cuda_decode_stream()>>>(
(float *)out->ptr, weight, (const float *)kv_norm->ptr,
n_tokens, kv_lora_dim, qk_nope, n_head, row_bytes);
return cuda_ok(cudaGetLastError(), "GLM k_b Q8_0 projection launch");
}
__global__ static void glm_kv_lora_rms_norm_kernel(
float *dst,
const float *src,
const float *w,
uint32_t n_tokens,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
float eps) {
const uint32_t row = blockIdx.x;
if (row >= n_tokens) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
const float *x = src + (uint64_t)row * kv_raw_dim;
float *out = dst + (uint64_t)row * kv_lora_dim;
__shared__ float scratch[256];
float ss = 0.0f;
for (uint32_t i = tid; i < kv_lora_dim; i += nth) {
const float v = x[i];
ss += v * v;
}
scratch[tid] = ss;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) scratch[tid] += scratch[tid + step];
__syncthreads();
}
const float inv = rsqrtf(scratch[0] / (float)kv_lora_dim + eps);
for (uint32_t i = tid; i < kv_lora_dim; i += nth) {
out[i] = x[i] * inv * w[i];
}
}
extern "C" int ds4_gpu_glm_kv_lora_rms_norm_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
float eps) {
if (!dst || !src || !model_map || n_tokens == 0 ||
kv_lora_dim == 0 || kv_lora_dim > kv_raw_dim) {
return 0;
}
const uint64_t wb = (uint64_t)kv_lora_dim * sizeof(float);
if (weight_offset > model_size || wb > model_size - weight_offset ||
src->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) ||
dst->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const float *w = (const float *)cuda_resolve_weight_ptr(
model_map, weight_offset, wb, logical_tier, "glm_kv_lora_norm");
if (!w) return 0;
glm_kv_lora_rms_norm_kernel<<<n_tokens, 256>>>(
(float *)dst->ptr, (const float *)src->ptr, w,
n_tokens, kv_raw_dim, kv_lora_dim, eps);
return cuda_ok(cudaGetLastError(), "glm kv lora rms norm launch");
}
__global__ static void glm_qk_lowrank_q8_0_batch_kernel(
float *qk_low,
const char *weight,
const float *q,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim,
uint64_t row_bytes) {
const uint32_t head = blockIdx.x;
const uint32_t token = blockIdx.y;
if (head >= n_head || token >= n_tokens) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
const float *qh = q + (uint64_t)token * n_head * qk_dim +
(uint64_t)head * qk_dim;
float *out = qk_low + (uint64_t)token * n_head * kv_lora_dim +
(uint64_t)head * kv_lora_dim;
for (uint32_t j = tid; j < kv_lora_dim; j += nth) {
const char *row = weight +
((uint64_t)head * kv_lora_dim + j) * row_bytes;
out[j] = glm_q8_0_dot_row_dev(row, qh, qk_nope);
}
}
extern "C" int ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
if (!qk_low || !q || !model_map || n_tokens == 0 || n_head == 0 ||
kv_lora_dim == 0 || qk_nope == 0 || (qk_nope & 31u) != 0u) {
return 0;
}
if (weight_type != 8u) {
fprintf(stderr, "ds4: glm qk_lowrank: unsupported type %u\n",
weight_type);
return 0;
}
const uint64_t row_bytes = ((uint64_t)qk_nope / 32u) * 34u;
const uint64_t wbytes = (uint64_t)n_head * kv_lora_dim * row_bytes;
if (weight_offset > model_size || wbytes > model_size - weight_offset ||
q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) ||
qk_low->bytes <
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const char *w = (const char *)cuda_resolve_weight_ptr(
model_map, weight_offset, wbytes, logical_tier, "glm_k_b_qk");
if (!w) return 0;
if (g_q8_dequant_gemm_enabled && g_cublas_ready && n_tokens >= 128u) {
/* Per-head strided-batched GEMM over a dequantized k_b: the
* per-(token,head) warp kernel was ~65ms/layer at 820 tokens.
* Scratch (executing device): [w_f16][q_f16][out_f32]. */
const uint64_t wh_bytes =
(uint64_t)n_head * kv_lora_dim * qk_nope * sizeof(__half);
const uint64_t xh_off = (wh_bytes + 255u) & ~255ull;
const uint64_t xh_bytes =
(uint64_t)n_tokens * n_head * qk_dim * sizeof(__half);
const uint64_t oo_off = (xh_off + xh_bytes + 255u) & ~255ull;
const uint64_t oo_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
void *tmp = cuda_tmp_alloc_on(logical_tier, oo_off + oo_bytes,
"glm qk_low gemm");
if (tmp) {
__half *wh = (__half *)tmp;
__half *xh = (__half *)((char *)tmp + xh_off);
float *oo = (float *)((char *)tmp + oo_off);
const uint64_t total_blocks =
(uint64_t)n_head * kv_lora_dim * (qk_nope / 32u);
q8_0_dequant_f16_kernel<<<(unsigned)((total_blocks * 2u + 255u) / 256u), 256>>>(
wh, (const unsigned char *)w, total_blocks,
qk_nope / 32u, qk_nope);
const uint64_t xn = (uint64_t)n_tokens * n_head * qk_dim;
f32_to_f16_kernel<<<(xn + 255u) / 256u, 256>>>(
xh, (const float *)q->ptr, xn);
if (cuda_ok(cudaGetLastError(), "glm qk_low gemm staging")) {
const float alpha = 1.0f;
const float beta = 0.0f;
cublasStatus_t st = cublasGemmStridedBatchedEx(
cuda_cublas_for_tier(logical_tier),
CUBLAS_OP_T, CUBLAS_OP_N,
(int)kv_lora_dim, (int)n_tokens, (int)qk_nope,
&alpha,
wh, CUDA_R_16F, (int)qk_nope,
(long long)((uint64_t)kv_lora_dim * qk_nope),
xh, CUDA_R_16F, (int)(n_head * qk_dim),
(long long)qk_dim,
&beta,
oo, CUDA_R_32F, (int)(n_head * kv_lora_dim),
(long long)kv_lora_dim,
(int)n_head,
CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
if (st == CUBLAS_STATUS_SUCCESS &&
cuda_ok(cudaMemcpyAsync(qk_low->ptr, oo, oo_bytes,
cudaMemcpyDeviceToDevice, 0),
"glm qk_low gemm out copy")) {
return 1;
}
fprintf(stderr,
"ds4: glm qk_low gemm failed (status %d); native path\n",
(int)st);
}
}
}
dim3 grid(n_head, n_tokens, 1);
glm_qk_lowrank_q8_0_batch_kernel<<<grid, 128>>>(
(float *)qk_low->ptr, w, (const float *)q->ptr,
n_tokens, n_head, kv_lora_dim, qk_nope, qk_dim, row_bytes);
return cuda_ok(cudaGetLastError(), "glm qk lowrank batch launch");
}
extern "C" int ds4_gpu_glm_qk_lowrank_typed_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low, q, model_map,
model_size,
weight_offset,
weight_type, 1u, n_head,
kv_lora_dim, qk_nope,
qk_dim);
}
/* Fused decode-path QKV norm + compact-KV store, one block per
* (token, part): part 0 rms-norms q into q_out, part 1 rms-norms
* kv_raw[:kv_lora_dim] into the kv_lora ring, part 2 copies the
* UNROTATED rope tail into the k_rope ring (roped at attention read). */
__global__ static void glm_qkv_norm_store_compact_kv_kernel(
float *q_dst,
const float *q_src,
const float *q_w,
uint32_t q_n,
const float *kv_raw,
const float *kv_w,
char *kv_lora_cache,
char *k_rope_cache,
uint32_t pos0,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
int cache_f16,
float eps) {
const uint32_t token = blockIdx.x;
const uint32_t part = blockIdx.y;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
const uint32_t pos = pos0 + token;
if (part == 2u) {
if (pos >= cache_cap) return;
const float *src = kv_raw + (uint64_t)token * kv_raw_dim + kv_lora_dim;
if (cache_f16) {
__half *dst = (__half *)k_rope_cache + (uint64_t)pos * qk_rope;
for (uint32_t i = tid; i < qk_rope; i += nth) {
dst[i] = __float2half(src[i]);
}
} else {
float *dst = (float *)k_rope_cache + (uint64_t)pos * qk_rope;
for (uint32_t i = tid; i < qk_rope; i += nth) {
dst[i] = src[i];
}
}
return;
}
const bool kv_task = part != 0u;
const uint32_t n = kv_task ? kv_lora_dim : q_n;
const float *x = kv_task ? kv_raw + (uint64_t)token * kv_raw_dim
: q_src + (uint64_t)token * q_n;
const float *w = kv_task ? kv_w : q_w;
__shared__ float sh[32];
float sumf = 0.0f;
for (uint32_t i = tid; i < n; i += nth) {
const float v = x[i];
sumf += v * v;
}
for (int off = 16; off > 0; off >>= 1) {
sumf += __shfl_xor_sync(0xffffffffu, sumf, off);
}
if ((tid & 31u) == 0u) sh[tid >> 5] = sumf;
__syncthreads();
if (tid < 32u) {
sumf = (tid < (nth + 31u) / 32u) ? sh[tid] : 0.0f;
for (int off = 16; off > 0; off >>= 1) {
sumf += __shfl_xor_sync(0xffffffffu, sumf, off);
}
if (tid == 0u) sh[0] = sumf;
}
__syncthreads();
const float scale = rsqrtf(sh[0] / (float)n + eps);
if (!kv_task) {
float *y = q_dst + (uint64_t)token * q_n;
for (uint32_t i = tid; i < n; i += nth) {
y[i] = (x[i] * scale) * w[i];
}
return;
}
if (pos >= cache_cap) return;
if (cache_f16) {
__half *dst = (__half *)kv_lora_cache + (uint64_t)pos * kv_lora_dim;
for (uint32_t i = tid; i < kv_lora_dim; i += nth) {
dst[i] = __float2half((x[i] * scale) * w[i]);
}
} else {
float *dst = (float *)kv_lora_cache + (uint64_t)pos * kv_lora_dim;
for (uint32_t i = tid; i < kv_lora_dim; i += nth) {
dst[i] = (x[i] * scale) * w[i];
}
}
}
extern "C" int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_raw,
uint64_t kv_weight_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16,
float eps) {
if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw ||
!model_map || n_tokens == 0 || q_n == 0 || kv_lora_dim == 0 ||
cache_cap == 0 || pos0 > cache_cap || n_tokens > cache_cap - pos0 ||
kv_lora_dim > UINT32_MAX - qk_rope ||
kv_raw_dim < kv_lora_dim + qk_rope ||
q->bytes < (uint64_t)n_tokens * q_n * sizeof(float) ||
q_out->bytes < (uint64_t)n_tokens * q_n * sizeof(float) ||
kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float) ||
kv_lora_cache->bytes <
(uint64_t)cache_cap * kv_lora_dim *
(cache_f16 ? sizeof(__half) : sizeof(float)) ||
k_rope_cache->bytes <
(uint64_t)cache_cap * qk_rope *
(cache_f16 ? sizeof(__half) : sizeof(float))) {
return 0;
}
if (q_weight_offset > model_size ||
(uint64_t)q_n * sizeof(float) > model_size - q_weight_offset ||
kv_weight_offset > model_size ||
(uint64_t)kv_lora_dim * sizeof(float) > model_size - kv_weight_offset) {
return 0;
}
const int logical_tier = cuda_current_tier();
const float *q_w = (const float *)cuda_resolve_weight_ptr(
model_map, q_weight_offset, (uint64_t)q_n * sizeof(float),
logical_tier, "glm_q_norm");
const float *kv_w = (const float *)cuda_resolve_weight_ptr(
model_map, kv_weight_offset,
(uint64_t)kv_lora_dim * sizeof(float),
logical_tier, "glm_kv_norm");
if (!q_w || !kv_w) return 0;
dim3 grid(n_tokens, qk_rope != 0u ? 3u : 2u, 1u);
glm_qkv_norm_store_compact_kv_kernel<<<grid, 256>>>(
(float *)q_out->ptr,
(const float *)q->ptr,
q_w,
q_n,
(const float *)kv_raw->ptr,
kv_w,
(char *)kv_lora_cache->ptr,
(char *)k_rope_cache->ptr,
pos0, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope,
cache_f16 ? 1 : 0, eps);
return cuda_ok(cudaGetLastError(), "glm qkv norm store compact kv");
}
/* In-place interleaved-pair yarn rope on a [n_tokens][n_head][head_dim]
* f32 tensor, rotating rot_dim dims starting at rot_offset. Shared by the
* attention q tail (offset = head_dim - rot_dim) and the DSA indexer
* (offset = 0). Grid (n_head, n_tokens). */
__global__ static void glm_rope_tail_offset_kernel(
float *x,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t rot_offset,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint32_t head = blockIdx.x;
const uint32_t token = blockIdx.y;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
if (rot_dim == 0u || (rot_dim & 1u) != 0u ||
rot_offset > head_dim || rot_dim > head_dim - rot_offset) return;
const uint32_t pos = pos0 + token;
float *row = x + ((uint64_t)token * n_head + head) * head_dim + rot_offset;
float corr_dims[2] = {0.0f, 0.0f};
if (ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim,
(int)n_ctx_orig, beta_fast, freq_base)));
corr_dims[1] = fminf((float)rot_dim - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim,
(int)n_ctx_orig, beta_slow, freq_base)));
}
const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)rot_dim;
for (uint32_t i = tid * 2u; i < rot_dim; i += nth * 2u) {
const float theta =
theta_base * powf(freq_base, inv_ndims * (float)i);
float ct, st;
glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i,
ext_factor, attn_factor, &ct, &st);
const float x0 = row[i];
const float x1 = row[i + 1u];
row[i] = x0 * ct - x1 * st;
row[i + 1u] = x0 * st + x1 * ct;
}
}
static int glm_rope_tail_offset_launch(
ds4_gpu_tensor *x,
uint32_t n_tokens, uint32_t n_head, uint32_t head_dim,
uint32_t rot_dim, uint32_t rot_offset, uint32_t pos0,
uint32_t n_ctx_orig, float freq_base, float freq_scale,
float ext_factor, float attn_factor,
float beta_fast, float beta_slow, const char *what) {
if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 ||
rot_dim == 0 || (rot_dim & 1u) != 0u || rot_offset > head_dim ||
rot_dim > head_dim - rot_offset ||
x->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float)) {
return 0;
}
dim3 grid(n_head, n_tokens, 1);
glm_rope_tail_offset_kernel<<<grid, 32>>>(
(float *)x->ptr, n_head, head_dim, rot_dim, rot_offset,
pos0, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
return cuda_ok(cudaGetLastError(), what);
}
extern "C" int ds4_gpu_glm_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (rot_dim > head_dim) return 0;
return glm_rope_tail_offset_launch(x, n_tokens, n_head, head_dim,
rot_dim, head_dim - rot_dim, pos0, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow, "glm rope tail");
}
extern "C" int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride) {
fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor\n");
return 0;
}
/* Scalar-correct GLM routed MoE (q2_K experts): per (token, slot) block
* quantizes nothing - dots q2_K rows against a q8_K-quantized activation
* staged in shared memory. Grid: (n_tokens, n_expert). Mid buffer holds
* silu(gate)*up per slot; out accumulates expert_weight-scaled down rows.
*/
__global__ static void glm_routed_moe_batch_q2K_gateup_kernel(
float *mid,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t n_tokens,
uint32_t mid_token_stride) {
const uint32_t tok = blockIdx.x;
const uint32_t slot = blockIdx.y;
if (tok >= n_tokens || slot >= n_expert) return;
const int32_t expert = selected[(uint64_t)tok * n_expert + slot];
if (expert < 0) return;
const cuda_block_q8_K *xrow = xq + (uint64_t)tok * xq_blocks;
float *mrow = mid + (uint64_t)tok * mid_token_stride +
(uint64_t)slot * expert_mid_dim;
for (uint32_t r = threadIdx.x; r < expert_mid_dim; r += blockDim.x) {
const cuda_block_q2_K *gr = (const cuda_block_q2_K *)(gate_base +
(uint64_t)expert * gate_expert_bytes + (uint64_t)r * gate_row_bytes);
const cuda_block_q2_K *ur = (const cuda_block_q2_K *)(up_base +
(uint64_t)expert * up_expert_bytes + (uint64_t)r * up_row_bytes);
float g = 0.0f, u = 0.0f;
for (uint32_t b = 0; b < xq_blocks; b++) {
g += dev_dot_q2_K_q8_K_block(gr + b, xrow + b);
u += dev_dot_q2_K_q8_K_block(ur + b, xrow + b);
}
mrow[r] = (g / (1.0f + expf(-g))) * u; /* silu(g)*u */
}
}
__global__ static void glm_routed_moe_batch_q2K_down_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
const float *weights,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.y;
if (tok >= n_tokens) return;
const uint32_t r = blockIdx.x * blockDim.x + threadIdx.x;
if (r >= out_dim) return;
float acc = 0.0f;
for (uint32_t slot = 0; slot < n_expert; slot++) {
const int32_t expert = selected[(uint64_t)tok * n_expert + slot];
if (expert < 0) continue;
const float w = weights[(uint64_t)tok * n_expert + slot];
const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base +
(uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes);
const cuda_block_q8_K *mrow = midq +
((uint64_t)tok * n_expert + slot) * midq_blocks;
float s = 0.0f;
for (uint32_t b = 0; b < midq_blocks; b++) {
s += dev_dot_q2_K_q8_K_block(dr + b, mrow + b);
}
acc += w * s;
}
out[(uint64_t)tok * out_dim + r] = acc;
}
/* Warp-per-row routed MoE (q2_K x q8_K). Each block stages the token's
* q8_K activation row in shared memory; one warp produces one mid row
* (gate dot + up dot + silu*mul fused). Grid:
* (expert_mid_dim/warps, n_expert, n_tokens). */
__global__ static void glm_routed_moe_gateup_warp_kernel(
float *mid,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.z;
const uint32_t slot = blockIdx.y;
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
if (tok >= n_tokens || slot >= n_expert) return;
const int32_t expert = selected[(uint64_t)tok * n_expert + slot];
extern __shared__ unsigned int glm_moe_sh_u32[];
{
const unsigned int *src =
(const unsigned int *)(xq + (uint64_t)tok * xq_blocks);
const uint32_t words = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u;
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) {
glm_moe_sh_u32[i] = src[i];
}
}
__syncthreads();
if (expert < 0) return;
const cuda_block_q8_K *xrow = (const cuda_block_q8_K *)glm_moe_sh_u32;
const uint32_t r = blockIdx.x * warps + warp;
if (r >= expert_mid_dim) return;
const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes +
(uint64_t)r * gate_row_bytes;
const char *ur = up_base + (uint64_t)expert * up_expert_bytes +
(uint64_t)r * up_row_bytes;
float g = 0.0f, u = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
g += dev_dot_q2_K_q8_K_block(
(const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b);
u += dev_dot_q2_K_q8_K_block(
(const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b);
}
for (int off = 16; off > 0; off >>= 1) {
g += __shfl_down_sync(0xffffffffu, g, off);
u += __shfl_down_sync(0xffffffffu, u, off);
}
if (lane == 0u) {
mid[((uint64_t)tok * n_expert + slot) * expert_mid_dim + r] =
(g / (1.0f + expf(-g))) * u;
}
}
/* Exact two-token gate/up with adjacent-token expert reuse. Token 0 owns an
* expert present in both rows; token 1 only launches work for unmatched
* experts. Each token keeps the native lane assignment and warp reduction. */
__global__ static void glm_routed_moe_gateup_tok2_reuse_kernel(
float *mid,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *selected,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert) {
const uint32_t owner = blockIdx.y;
const uint32_t tok = owner / n_expert;
const uint32_t slot = owner - tok * n_expert;
const uint32_t warps = blockDim.x >> 5u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
if (tok >= 2u || slot >= n_expert) return;
const int32_t expert = selected[owner];
if (expert < 0) return;
int32_t mate_slot = -1;
for (uint32_t s = 0; s < n_expert; s++) {
if (selected[(uint64_t)(1u - tok) * n_expert + s] == expert) {
mate_slot = (int32_t)s;
break;
}
}
if (tok == 1u && mate_slot >= 0) return;
const uint32_t np = tok == 0u && mate_slot >= 0 ? 2u : 1u;
const uint32_t pair0 = owner;
const uint32_t pair1 = n_expert + (uint32_t)mate_slot;
extern __shared__ unsigned int glm_moe_tok2_sh_u32[];
const uint32_t words_per_row =
xq_blocks * (uint32_t)sizeof(cuda_block_q8_K) / 4u;
const unsigned int *src0 = (const unsigned int *)(
xq + (uint64_t)tok * xq_blocks);
for (uint32_t i = threadIdx.x; i < words_per_row;
i += blockDim.x) {
glm_moe_tok2_sh_u32[i] = src0[i];
}
if (np == 2u) {
const unsigned int *src1 =
(const unsigned int *)(xq + xq_blocks);
for (uint32_t i = threadIdx.x; i < words_per_row;
i += blockDim.x) {
glm_moe_tok2_sh_u32[words_per_row + i] = src1[i];
}
}
__syncthreads();
const uint32_t r = blockIdx.x * warps + warp;
if (r >= expert_mid_dim) return;
const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes +
(uint64_t)r * gate_row_bytes;
const char *ur = up_base + (uint64_t)expert * up_expert_bytes +
(uint64_t)r * up_row_bytes;
const cuda_block_q8_K *x0 =
(const cuda_block_q8_K *)glm_moe_tok2_sh_u32;
const cuda_block_q8_K *x1 = np == 2u
? x0 + xq_blocks
: NULL;
float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
dev_dot_q2_K_q8_K_block8(
(const cuda_block_q2_K *)(gr + (uint64_t)b * 84u),
x0 + b, np == 2u ? x1 + b : NULL,
NULL, NULL, NULL, NULL, NULL, NULL, np, g);
dev_dot_q2_K_q8_K_block8(
(const cuda_block_q2_K *)(ur + (uint64_t)b * 84u),
x0 + b, np == 2u ? x1 + b : NULL,
NULL, NULL, NULL, NULL, NULL, NULL, np, u);
}
for (uint32_t p = 0; p < np; p++) {
for (int off = 16; off > 0; off >>= 1) {
g[p] += __shfl_down_sync(0xffffffffu, g[p], off);
u[p] += __shfl_down_sync(0xffffffffu, u[p], off);
}
if (lane == 0u) {
const uint32_t pair = p == 0u ? pair0 : pair1;
mid[(uint64_t)pair * expert_mid_dim + r] =
(g[p] / (1.0f + expf(-g[p]))) * u[p];
}
}
}
/* Warp-per-output-row down projection: stages all n_expert quantized mid
* rows for the token in shared memory, each warp accumulates one out row
* across every selected expert. Grid: (out_dim/warps, n_tokens). */
__global__ static void glm_routed_moe_down_warp_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *selected,
const float *weights,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.y;
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
if (tok >= n_tokens) return;
extern __shared__ unsigned int glm_moe_sh_u32[];
{
const unsigned int *src = (const unsigned int *)
(midq + (uint64_t)tok * n_expert * midq_blocks);
const uint32_t words = n_expert * midq_blocks *
(uint32_t)sizeof(cuda_block_q8_K) / 4u;
for (uint32_t i = threadIdx.x; i < words; i += blockDim.x) {
glm_moe_sh_u32[i] = src[i];
}
}
__syncthreads();
const cuda_block_q8_K *msh = (const cuda_block_q8_K *)glm_moe_sh_u32;
const uint32_t r = blockIdx.x * warps + warp;
if (r >= out_dim) return;
const uint32_t units = n_expert * midq_blocks;
float acc = 0.0f;
for (uint32_t idx = lane; idx < units; idx += 32u) {
const uint32_t slot = idx / midq_blocks;
const uint32_t b = idx - slot * midq_blocks;
const int32_t expert = selected[(uint64_t)tok * n_expert + slot];
if (expert < 0) continue;
const float w = weights[(uint64_t)tok * n_expert + slot];
const cuda_block_q2_K *dr = (const cuda_block_q2_K *)(down_base +
(uint64_t)expert * down_expert_bytes + (uint64_t)r * down_row_bytes);
acc += w * dev_dot_q2_K_q8_K_block(dr + b, msh + slot * midq_blocks + b);
}
for (int off = 16; off > 0; off >>= 1) {
acc += __shfl_down_sync(0xffffffffu, acc, off);
}
if (lane == 0u) out[(uint64_t)tok * out_dim + r] = acc;
}
/* Expert-major routed MoE for prefill: build per-expert token lists,
* then walk rows expert-by-expert so weights stream once per layer and
* activations hit L2. pair = tok * n_expert + slot indexes selected/
* weights/mid rows directly. */
__global__ static void glm_moe_expert_map_kernel(
int32_t *counts,
int32_t *lists,
const int32_t *selected,
uint32_t n_pairs,
uint32_t n_total_expert,
uint32_t cap,
uint32_t pair_base) {
const uint32_t p = blockIdx.x * blockDim.x + threadIdx.x;
if (p >= n_pairs) return;
const uint32_t pair = pair_base + p;
const int32_t e = selected[pair];
if (e < 0 || (uint32_t)e >= n_total_expert) return;
const int32_t idx = atomicAdd(&counts[e], 1);
lists[(uint64_t)e * cap + idx] = (int32_t)pair;
}
__global__ static void glm_moe_build_expert_tiles8_kernel(
uint32_t *tile_total,
uint32_t *tile_experts,
uint32_t *tile_starts,
const int32_t *counts,
uint32_t n_total_expert) {
if (blockIdx.x != 0u || threadIdx.x != 0u) return;
uint32_t total = 0;
for (uint32_t e = 0; e < n_total_expert; e++) {
const uint32_t count = counts[e] > 0 ? (uint32_t)counts[e] : 0u;
const uint32_t ntiles = (count + 7u) / 8u;
for (uint32_t t = 0; t < ntiles; t++) {
tile_experts[total] = e;
tile_starts[total] = t * 8u;
total++;
}
}
*tile_total = total;
}
/* Expert-tiled Q2_K gate/up for GLM prefill. One warp keeps the same
* block-to-lane assignment and reduction tree as the token-major W32
* kernel, but evaluates eight pairs against each loaded expert row. */
__global__ static void glm_routed_moe_gateup_expert_tile8_kernel(
float *mid,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *counts,
const int32_t *lists,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t cap) {
const uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t warps = blockDim.x >> 5u;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t expert = tile_experts[tile];
const uint32_t local_start = tile_starts[tile];
const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u;
__shared__ uint32_t pair[8];
__shared__ uint32_t tok[8];
__shared__ uint32_t np;
if (threadIdx.x == 0u) {
uint32_t n = count - local_start;
if (n > 8u) n = 8u;
np = n;
for (uint32_t p = 0; p < n; p++) {
const uint32_t pr =
(uint32_t)lists[(uint64_t)expert * cap + local_start + p];
pair[p] = pr;
tok[p] = pr / n_expert;
}
}
__syncthreads();
const uint32_t r = blockIdx.x * warps + warp;
if (r >= expert_mid_dim) return;
const char *gr = gate_base + (uint64_t)expert * gate_expert_bytes +
(uint64_t)r * gate_row_bytes;
const char *ur = up_base + (uint64_t)expert * up_expert_bytes +
(uint64_t)r * up_row_bytes;
float g[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
float u[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
const cuda_block_q8_K *x0 = np > 0u ? xq + (uint64_t)tok[0] * xq_blocks + b : NULL;
const cuda_block_q8_K *x1 = np > 1u ? xq + (uint64_t)tok[1] * xq_blocks + b : NULL;
const cuda_block_q8_K *x2 = np > 2u ? xq + (uint64_t)tok[2] * xq_blocks + b : NULL;
const cuda_block_q8_K *x3 = np > 3u ? xq + (uint64_t)tok[3] * xq_blocks + b : NULL;
const cuda_block_q8_K *x4 = np > 4u ? xq + (uint64_t)tok[4] * xq_blocks + b : NULL;
const cuda_block_q8_K *x5 = np > 5u ? xq + (uint64_t)tok[5] * xq_blocks + b : NULL;
const cuda_block_q8_K *x6 = np > 6u ? xq + (uint64_t)tok[6] * xq_blocks + b : NULL;
const cuda_block_q8_K *x7 = np > 7u ? xq + (uint64_t)tok[7] * xq_blocks + b : NULL;
dev_dot_q2_K_q8_K_block8(
(const cuda_block_q2_K *)(gr + (uint64_t)b * 84u),
x0, x1, x2, x3, x4, x5, x6, x7, np, g);
dev_dot_q2_K_q8_K_block8(
(const cuda_block_q2_K *)(ur + (uint64_t)b * 84u),
x0, x1, x2, x3, x4, x5, x6, x7, np, u);
}
for (uint32_t p = 0; p < np; p++) {
for (int off = 16; off > 0; off >>= 1) {
g[p] += __shfl_down_sync(0xffffffffu, g[p], off);
u[p] += __shfl_down_sync(0xffffffffu, u[p], off);
}
if (lane == 0u) {
mid[(uint64_t)pair[p] * expert_mid_dim + r] =
(g[p] / (1.0f + expf(-g[p]))) * u[p];
}
}
}
__global__ static void glm_routed_moe_gateup_expert_kernel(
float *mid,
const char *gate_base,
const char *up_base,
const cuda_block_q8_K *xq,
const int32_t *counts,
const int32_t *lists,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint32_t xq_blocks,
uint32_t expert_mid_dim,
uint32_t n_expert,
uint32_t cap) {
const uint32_t e = blockIdx.y;
const int32_t nt = counts[e];
if (nt == 0) return;
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t r = blockIdx.x * warps + warp;
if (r >= expert_mid_dim) return;
const char *gr = gate_base + (uint64_t)e * gate_expert_bytes +
(uint64_t)r * gate_row_bytes;
const char *ur = up_base + (uint64_t)e * up_expert_bytes +
(uint64_t)r * up_row_bytes;
const int32_t *lst = lists + (uint64_t)e * cap;
for (int32_t i = 0; i < nt; i++) {
const uint32_t pair = (uint32_t)lst[i];
const cuda_block_q8_K *xrow = xq + (uint64_t)(pair / n_expert) * xq_blocks;
float g = 0.0f, u = 0.0f;
for (uint32_t b = lane; b < xq_blocks; b += 32u) {
g += dev_dot_q2_K_q8_K_block(
(const cuda_block_q2_K *)(gr + (uint64_t)b * 84u), xrow + b);
u += dev_dot_q2_K_q8_K_block(
(const cuda_block_q2_K *)(ur + (uint64_t)b * 84u), xrow + b);
}
for (int off = 16; off > 0; off >>= 1) {
g += __shfl_down_sync(0xffffffffu, g, off);
u += __shfl_down_sync(0xffffffffu, u, off);
}
if (lane == 0u) {
mid[(uint64_t)pair * expert_mid_dim + r] =
(g / (1.0f + expf(-g))) * u;
}
}
}
__global__ static void glm_routed_moe_down_expert_kernel(
float *out,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *counts,
const int32_t *lists,
const float *weights,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t cap) {
const uint32_t e = blockIdx.y;
const int32_t nt = counts[e];
if (nt == 0) return;
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t r = blockIdx.x * warps + warp;
if (r >= out_dim) return;
const char *dr = down_base + (uint64_t)e * down_expert_bytes +
(uint64_t)r * down_row_bytes;
const int32_t *lst = lists + (uint64_t)e * cap;
for (int32_t i = 0; i < nt; i++) {
const uint32_t pair = (uint32_t)lst[i];
const cuda_block_q8_K *mrow = midq + (uint64_t)pair * midq_blocks;
float s = 0.0f;
for (uint32_t b = lane; b < midq_blocks; b += 32u) {
s += dev_dot_q2_K_q8_K_block(
(const cuda_block_q2_K *)(dr + (uint64_t)b * 84u), mrow + b);
}
for (int off = 16; off > 0; off >>= 1) {
s += __shfl_down_sync(0xffffffffu, s, off);
}
if (lane == 0u) {
atomicAdd(&out[(uint64_t)(pair / n_expert) * out_dim + r],
weights[pair] * s);
}
}
}
/* Expert-tiled down projection with an exact token-major reduction. The
* first kernel reuses each Q2_K row across eight routed pairs, but materializes
* each block dot. The second kernel applies the router weight and consumes the
* dots with the same lane assignment and warp tree as the native kernel. */
__global__ static void glm_routed_moe_down_expert_tile8_terms_kernel(
float *terms,
const char *down_base,
const cuda_block_q8_K *midq,
const int32_t *counts,
const int32_t *lists,
const uint32_t *tile_total,
const uint32_t *tile_experts,
const uint32_t *tile_starts,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t cap,
uint32_t pair_base) {
const uint32_t tile = blockIdx.y;
if (tile >= *tile_total) return;
const uint32_t lane = threadIdx.x & 7u;
const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u);
const uint32_t expert = tile_experts[tile];
const uint32_t local_start = tile_starts[tile];
const uint32_t count = counts[expert] > 0 ? (uint32_t)counts[expert] : 0u;
__shared__ uint32_t pair[8];
__shared__ uint32_t np;
__shared__ cuda_block_q8_K mq[8][8];
if (threadIdx.x == 0u) {
uint32_t n = count - local_start;
if (n > 8u) n = 8u;
np = n;
for (uint32_t p = 0; p < n; p++) {
pair[p] = (uint32_t)lists[
(uint64_t)expert * cap + local_start + p];
}
}
__syncthreads();
for (uint32_t i = threadIdx.x; i < np * midq_blocks; i += blockDim.x) {
const uint32_t p = i / midq_blocks;
const uint32_t b = i - p * midq_blocks;
mq[p][b] = midq[(uint64_t)pair[p] * midq_blocks + b];
}
__syncthreads();
if (row >= out_dim || lane >= midq_blocks) return;
const cuda_block_q2_K *wr = (const cuda_block_q2_K *)(down_base +
(uint64_t)expert * down_expert_bytes +
(uint64_t)row * down_row_bytes);
for (uint32_t p = 0; p < np; p++) {
const uint32_t pr = pair[p];
const float dot = dev_dot_q2_K_q8_K_block(wr + lane, &mq[p][lane]);
terms[((uint64_t)(pr - pair_base) * out_dim + row) *
midq_blocks + lane] =
dot;
}
}
__global__ static void glm_routed_moe_down_terms_reduce_kernel(
float *out,
const float *terms,
const int32_t *selected,
const float *weights,
uint32_t midq_blocks,
uint32_t out_dim,
uint32_t n_expert,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.y;
const uint32_t warp = threadIdx.x >> 5u;
const uint32_t lane = threadIdx.x & 31u;
const uint32_t row = blockIdx.x * 8u + warp;
if (tok >= n_tokens || row >= out_dim) return;
const uint32_t units = n_expert * midq_blocks;
float acc = 0.0f;
for (uint32_t idx = lane; idx < units; idx += 32u) {
const uint32_t slot = idx / midq_blocks;
const uint32_t b = idx - slot * midq_blocks;
const uint32_t pr = tok * n_expert + slot;
if (selected[pr] >= 0) {
acc += weights[pr] *
terms[((uint64_t)pr * out_dim + row) * midq_blocks + b];
}
}
for (int off = 16; off > 0; off >>= 1) {
acc += __shfl_down_sync(0xffffffffu, acc, off);
}
if (lane == 0u) out[(uint64_t)tok * out_dim + row] = acc;
}
static int glm_routed_moe_finish_batch(
ds4_gpu_tensor *out,
float *out_work,
uint64_t out_bytes,
const char *what) {
if (!cuda_ok(cudaGetLastError(), what)) return 0;
if (out_work == (float *)out->ptr) return 1;
return cuda_ok(cudaMemcpyAsync(out->ptr, out_work, out_bytes,
cudaMemcpyDeviceToDevice, 0),
"glm routed moe local output copy");
}
extern "C" int ds4_gpu_glm_routed_moe_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride) {
(void)layer_index; (void)n_total_expert;
if (!out || !mid || !x || !selected || !weights || !model_map ||
n_tokens == 0 || n_expert == 0 ||
(expert_in_dim & 255u) != 0u || (expert_mid_dim & 255u) != 0u) {
return 0;
}
if (gate_type != 10u || up_type != 10u || down_type != 10u) {
fprintf(stderr, "ds4: glm routed moe: unsupported types %u/%u/%u\n",
gate_type, up_type, down_type);
return 0;
}
if (mid_token_stride != n_expert * expert_mid_dim) {
fprintf(stderr,
"ds4: glm routed moe: mid stride %u != %u (packed rows expected)\n",
mid_token_stride, n_expert * expert_mid_dim);
return 0;
}
const int logical_tier = cuda_current_tier();
const char *gw = (const char *)cuda_resolve_weight_ptr(model_map,
gate_offset, (uint64_t)256 * gate_expert_bytes, logical_tier,
"glm_gate_exps");
const char *uw = (const char *)cuda_resolve_weight_ptr(model_map,
up_offset, (uint64_t)256 * up_expert_bytes, logical_tier,
"glm_up_exps");
const char *dw = (const char *)cuda_resolve_weight_ptr(model_map,
down_offset, (uint64_t)256 * down_expert_bytes, logical_tier,
"glm_down_exps");
if (!gw || !uw || !dw) return 0;
/* Stage 1: quantize x rows to q8_K (existing kernel). */
const uint32_t xq_blocks = expert_in_dim / 256u;
const uint32_t midq_blocks = expert_mid_dim / 256u;
static ds4_gpu_tensor *xq_scratch[DS4_MAX_GPUS] = {0};
static ds4_gpu_tensor *midq_scratch[DS4_MAX_GPUS] = {0};
int dev = logical_tier;
const int scratch_tier = getenv("DS4_GLM_MOE_SCRATCH_TIER0") ? 0 : dev;
const uint64_t xq_bytes = (uint64_t)n_tokens * xq_blocks *
sizeof(cuda_block_q8_K);
const uint64_t midq_bytes = (uint64_t)n_tokens * n_expert * midq_blocks *
sizeof(cuda_block_q8_K);
if (!xq_scratch[dev] || xq_scratch[dev]->bytes < xq_bytes) {
if (xq_scratch[dev]) ds4_gpu_tensor_free(xq_scratch[dev]);
xq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, xq_bytes);
}
if (!midq_scratch[dev] || midq_scratch[dev]->bytes < midq_bytes) {
if (midq_scratch[dev]) ds4_gpu_tensor_free(midq_scratch[dev]);
midq_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(scratch_tier, midq_bytes);
}
if (!xq_scratch[dev] || !midq_scratch[dev]) return 0;
static ds4_gpu_tensor *mid_local[DS4_MAX_GPUS] = {0};
static ds4_gpu_tensor *out_local[DS4_MAX_GPUS] = {0};
const uint64_t mid_work_bytes =
(uint64_t)n_tokens * mid_token_stride * sizeof(float);
const uint64_t out_work_bytes =
(uint64_t)n_tokens * out_dim * sizeof(float);
float *mid_work = (float *)mid->ptr;
float *out_work = (float *)out->ptr;
const bool use_local_batch_io =
n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_LOCAL_BATCH_IO");
if (use_local_batch_io && ds4_tensor_device_idx(mid) != dev) {
if (!mid_local[dev] || mid_local[dev]->bytes < mid_work_bytes) {
if (mid_local[dev]) ds4_gpu_tensor_free(mid_local[dev]);
mid_local[dev] =
ds4_gpu_tensor_alloc_ptr_on(dev, mid_work_bytes);
}
if (!mid_local[dev]) return 0;
mid_work = (float *)mid_local[dev]->ptr;
}
if (use_local_batch_io && ds4_tensor_device_idx(out) != dev) {
if (!out_local[dev] || out_local[dev]->bytes < out_work_bytes) {
if (out_local[dev]) ds4_gpu_tensor_free(out_local[dev]);
out_local[dev] =
ds4_gpu_tensor_alloc_ptr_on(dev, out_work_bytes);
}
if (!out_local[dev]) return 0;
out_work = (float *)out_local[dev]->ptr;
}
{
dim3 gq(xq_blocks, n_tokens, 1);
q8_K_quantize_kernel<<<gq, 256>>>(
(cuda_block_q8_K *)xq_scratch[dev]->ptr,
(const float *)x->ptr, expert_in_dim, n_tokens);
}
static ds4_gpu_tensor *map_scratch[DS4_MAX_GPUS] = {0};
static ds4_gpu_tensor *down_terms_scratch[DS4_MAX_GPUS] = {0};
const bool use_expert_tile8 =
n_tokens >= 128u && !getenv("DS4_GLM_MOE_NO_EXPERT_TILE8");
const bool use_expert_major =
n_tokens >= 16u && getenv("DS4_GLM_MOE_EXPERT_MAJOR");
if (use_expert_tile8 || use_expert_major) {
const uint32_t cap = n_tokens;
const uint32_t n_pairs = n_tokens * n_expert;
const uint64_t counts_bytes = 256u * sizeof(int32_t);
const uint64_t lists_off = (counts_bytes + 255u) & ~255ull;
const uint64_t lists_bytes =
(uint64_t)256u * cap * sizeof(int32_t);
const uint32_t tile_capacity =
(n_pairs + 7u) / 8u + 256u;
const uint64_t tile_total_off =
(lists_off + lists_bytes + 255u) & ~255ull;
const uint64_t tile_experts_off =
(tile_total_off + sizeof(uint32_t) + 255u) & ~255ull;
const uint64_t tile_starts_off =
tile_experts_off + (uint64_t)tile_capacity * sizeof(uint32_t);
const uint64_t map_bytes = use_expert_tile8
? tile_starts_off + (uint64_t)tile_capacity * sizeof(uint32_t)
: lists_off + lists_bytes;
if (!map_scratch[dev] || map_scratch[dev]->bytes < map_bytes) {
if (map_scratch[dev]) ds4_gpu_tensor_free(map_scratch[dev]);
map_scratch[dev] = ds4_gpu_tensor_alloc_ptr_on(dev, map_bytes);
}
if (map_scratch[dev]) {
int32_t *counts = (int32_t *)map_scratch[dev]->ptr;
int32_t *lists = (int32_t *)((char *)map_scratch[dev]->ptr + lists_off);
cudaMemsetAsync(counts, 0, counts_bytes);
glm_moe_expert_map_kernel<<<(n_pairs + 255u) / 256u, 256>>>(
counts, lists, (const int32_t *)selected->ptr,
n_pairs, 256u, cap, 0u);
if (use_expert_tile8) {
uint32_t *tile_total = (uint32_t *)(
(char *)map_scratch[dev]->ptr + tile_total_off);
uint32_t *tile_experts = (uint32_t *)(
(char *)map_scratch[dev]->ptr + tile_experts_off);
uint32_t *tile_starts = (uint32_t *)(
(char *)map_scratch[dev]->ptr + tile_starts_off);
glm_moe_build_expert_tiles8_kernel<<<1, 1>>>(
tile_total, tile_experts, tile_starts,
counts, 256u);
dim3 ge1((expert_mid_dim + 7u) / 8u,
tile_capacity, 1);
glm_routed_moe_gateup_expert_tile8_kernel<<<ge1, 256>>>(
mid_work, gw, uw,
(const cuda_block_q8_K *)xq_scratch[dev]->ptr,
counts, lists,
tile_total, tile_experts, tile_starts,
gate_expert_bytes, gate_row_bytes,
up_expert_bytes, up_row_bytes,
xq_blocks, expert_mid_dim, n_expert, cap);
q8_K_quantize_kernel<<<
dim3(midq_blocks, n_tokens * n_expert, 1), 256>>>(
(cuda_block_q8_K *)midq_scratch[dev]->ptr,
mid_work,
expert_mid_dim, n_tokens * n_expert);
if (getenv("DS4_GLM_MOE_NO_DOWN_TILE8_EXACT") == NULL) {
const uint32_t max_chunk_tokens = 512u;
const uint32_t scratch_tokens =
n_tokens < max_chunk_tokens ? n_tokens : max_chunk_tokens;
const uint64_t term_count =
(uint64_t)scratch_tokens * n_expert *
out_dim * midq_blocks;
const uint64_t term_bytes = term_count * sizeof(float);
if (!down_terms_scratch[dev] ||
down_terms_scratch[dev]->bytes < term_bytes) {
if (down_terms_scratch[dev]) {
ds4_gpu_tensor_free(down_terms_scratch[dev]);
}
down_terms_scratch[dev] =
ds4_gpu_tensor_alloc_ptr_on(dev, term_bytes);
}
if (down_terms_scratch[dev]) {
for (uint32_t token0 = 0; token0 < n_tokens;
token0 += max_chunk_tokens) {
uint32_t chunk_tokens = n_tokens - token0;
if (chunk_tokens > max_chunk_tokens) {
chunk_tokens = max_chunk_tokens;
}
const uint32_t pair_base = token0 * n_expert;
const uint32_t chunk_pairs =
chunk_tokens * n_expert;
uint32_t chunk_tile_capacity = tile_capacity;
if (n_tokens > max_chunk_tokens) {
chunk_tile_capacity =
(chunk_pairs + 7u) / 8u + 256u;
cudaMemsetAsync(counts, 0, counts_bytes);
glm_moe_expert_map_kernel<<<
(chunk_pairs + 255u) / 256u, 256>>>(
counts, lists,
(const int32_t *)selected->ptr,
chunk_pairs, 256u, cap, pair_base);
glm_moe_build_expert_tiles8_kernel<<<1, 1>>>(
tile_total, tile_experts, tile_starts,
counts, 256u);
}
dim3 gd1((out_dim + 31u) / 32u,
chunk_tile_capacity, 1);
glm_routed_moe_down_expert_tile8_terms_kernel<<<
gd1, 256>>>(
(float *)down_terms_scratch[dev]->ptr,
dw,
(const cuda_block_q8_K *)midq_scratch[dev]->ptr,
counts, lists,
tile_total, tile_experts, tile_starts,
down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, cap,
pair_base);
dim3 gd2((out_dim + 7u) / 8u,
chunk_tokens, 1);
glm_routed_moe_down_terms_reduce_kernel<<<
gd2, 256>>>(
out_work + (uint64_t)token0 * out_dim,
(const float *)down_terms_scratch[dev]->ptr,
(const int32_t *)selected->ptr + pair_base,
(const float *)weights->ptr + pair_base,
midq_blocks, out_dim, n_expert,
chunk_tokens);
}
return glm_routed_moe_finish_batch(
out, out_work, out_work_bytes,
"glm routed moe exact down tile8");
}
}
const uint32_t warps = 8u;
dim3 ge2((out_dim + warps - 1u) / warps,
n_tokens, 1);
const uint32_t sh2 = n_expert * midq_blocks *
(uint32_t)sizeof(cuda_block_q8_K);
glm_routed_moe_down_warp_kernel<<<
ge2, warps * 32u, sh2>>>(
out_work, dw,
(const cuda_block_q8_K *)midq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, n_tokens);
return glm_routed_moe_finish_batch(
out, out_work, out_work_bytes,
"glm routed moe expert tile8");
}
dim3 ge1((expert_mid_dim + 7u) / 8u, 256u, 1);
glm_routed_moe_gateup_expert_kernel<<<ge1, 256>>>(
mid_work, gw, uw,
(const cuda_block_q8_K *)xq_scratch[dev]->ptr,
counts, lists,
gate_expert_bytes, gate_row_bytes,
up_expert_bytes, up_row_bytes,
xq_blocks, expert_mid_dim, n_expert, cap);
q8_K_quantize_kernel<<<dim3(midq_blocks, n_tokens * n_expert, 1), 256>>>(
(cuda_block_q8_K *)midq_scratch[dev]->ptr,
mid_work, expert_mid_dim, n_tokens * n_expert);
cudaMemsetAsync(out_work, 0,
(uint64_t)n_tokens * out_dim * sizeof(float));
dim3 ge2((out_dim + 7u) / 8u, 256u, 1);
glm_routed_moe_down_expert_kernel<<<ge2, 256>>>(
out_work, dw,
(const cuda_block_q8_K *)midq_scratch[dev]->ptr,
counts, lists, (const float *)weights->ptr,
down_expert_bytes, down_row_bytes,
midq_blocks, out_dim, n_expert, cap);
return glm_routed_moe_finish_batch(
out, out_work, out_work_bytes,
"glm routed moe expert-major");
}
}
if (n_tokens == 2u &&
g_glm_mtp_verify_mode &&
getenv("DS4_GLM_MTP_NO_MOE_TOK2") == NULL) {
const uint32_t warps = 8u;
dim3 g1((expert_mid_dim + warps - 1u) / warps,
2u * n_expert, 1u);
const uint32_t sh1 = 2u * xq_blocks *
(uint32_t)sizeof(cuda_block_q8_K);
glm_routed_moe_gateup_tok2_reuse_kernel<<<
g1, warps * 32u, sh1>>>(
mid_work, gw, uw,
(const cuda_block_q8_K *)xq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
gate_expert_bytes, gate_row_bytes,
up_expert_bytes, up_row_bytes,
xq_blocks, expert_mid_dim, n_expert);
} else if (getenv("DS4_GLM_MOE_SCALAR")) {
dim3 g1(n_tokens, n_expert, 1);
glm_routed_moe_batch_q2K_gateup_kernel<<<g1, 128>>>(
mid_work, gw, uw,
(const cuda_block_q8_K *)xq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes,
xq_blocks, expert_mid_dim, n_expert, n_tokens, mid_token_stride);
} else {
const uint32_t warps = 8u;
dim3 g1((expert_mid_dim + warps - 1u) / warps, n_expert, n_tokens);
const uint32_t sh1 = xq_blocks * (uint32_t)sizeof(cuda_block_q8_K);
glm_routed_moe_gateup_warp_kernel<<<g1, warps * 32u, sh1>>>(
mid_work, gw, uw,
(const cuda_block_q8_K *)xq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes,
xq_blocks, expert_mid_dim, n_expert, n_tokens);
}
{
dim3 gq(midq_blocks, n_tokens * n_expert, 1);
q8_K_quantize_kernel<<<gq, 256>>>(
(cuda_block_q8_K *)midq_scratch[dev]->ptr,
mid_work, expert_mid_dim, n_tokens * n_expert);
}
if (getenv("DS4_GLM_MOE_SCALAR")) {
dim3 g2((out_dim + 127u) / 128u, n_tokens, 1);
glm_routed_moe_batch_q2K_down_kernel<<<g2, 128>>>(
out_work, dw,
(const cuda_block_q8_K *)midq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
down_expert_bytes, down_row_bytes, midq_blocks, out_dim,
n_expert, n_tokens);
} else {
const uint32_t warps = 8u;
dim3 g2((out_dim + warps - 1u) / warps, n_tokens, 1);
const uint32_t sh2 = n_expert * midq_blocks *
(uint32_t)sizeof(cuda_block_q8_K);
glm_routed_moe_down_warp_kernel<<<g2, warps * 32u, sh2>>>(
out_work, dw,
(const cuda_block_q8_K *)midq_scratch[dev]->ptr,
(const int32_t *)selected->ptr,
(const float *)weights->ptr,
down_expert_bytes, down_row_bytes, midq_blocks, out_dim,
n_expert, n_tokens);
}
return glm_routed_moe_finish_batch(
out, out_work, out_work_bytes,
"glm routed moe batch launch");
}
extern "C" int ds4_gpu_glm_routed_moe_one_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t layer_index,
const ds4_gpu_tensor *x,
bool force_resident) {
(void)force_resident;
return ds4_gpu_glm_routed_moe_batch_tensor(out, mid,
model_map, model_size,
gate_offset, up_offset, down_offset,
gate_type, up_type, down_type,
gate_expert_bytes, gate_row_bytes,
up_expert_bytes, up_row_bytes,
down_expert_bytes, down_row_bytes,
expert_in_dim, expert_mid_dim, out_dim,
selected, weights, n_total_expert, n_expert, layer_index,
x, 1, n_expert * expert_mid_dim);
}
/* Parallel router select: 256 threads cover up to two experts each, then top-k
* via k rounds of shared-memory argmax over probs+bias (value desc, index
* asc tie-break -- matches the CPU topk_desc). One block per token. */
__global__ static void glm_router_select_parallel_kernel(
int32_t *selected,
float *weights_out,
float *probs_out,
const float *bias,
const float *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.x;
const uint32_t tid = threadIdx.x;
if (tok >= n_tokens) return;
const float *lg = logits + (uint64_t)tok * n_expert;
float *probs = probs_out + (uint64_t)tok * n_expert;
int32_t *sel = selected + (uint64_t)tok * n_expert_used;
float *w = weights_out + (uint64_t)tok * n_expert_used;
__shared__ float sh_v[256];
__shared__ int sh_i[256];
__shared__ float sh_sum;
const uint32_t i0 = tid;
const uint32_t i1 = tid + blockDim.x;
float v0 = -1e30f;
float v1 = -1e30f;
if (i0 < n_expert) {
const float p = 1.0f / (1.0f + expf(-lg[i0]));
probs[i0] = p;
v0 = p + bias[i0];
}
if (i1 < n_expert) {
const float p = 1.0f / (1.0f + expf(-lg[i1]));
probs[i1] = p;
v1 = p + bias[i1];
}
if (tid == 0u) sh_sum = 0.0f;
__syncthreads();
for (uint32_t k2 = 0; k2 < n_expert_used; k2++) {
const bool take1 = v1 > v0 || (v1 == v0 && i1 < i0);
sh_v[tid] = take1 ? v1 : v0;
sh_i[tid] = (int)(take1 ? i1 : i0);
__syncthreads();
for (uint32_t step = 128u; step > 0u; step >>= 1u) {
if (tid < step) {
const float ov = sh_v[tid + step];
const int oi = sh_i[tid + step];
if (ov > sh_v[tid] || (ov == sh_v[tid] && oi < sh_i[tid])) {
sh_v[tid] = ov;
sh_i[tid] = oi;
}
}
__syncthreads();
}
if (tid == 0u) {
const int best = sh_i[0];
sel[k2] = best;
const float p = probs[best];
w[k2] = p;
sh_sum += p;
}
__syncthreads();
const int best = sh_i[0];
if ((int)i0 == best) v0 = -1e30f;
if ((int)i1 == best) v1 = -1e30f;
__syncthreads();
}
if (tid == 0u) {
float sum = sh_sum;
if (sum < 6.103515625e-5f) sum = 6.103515625e-5f;
for (uint32_t k2 = 0; k2 < n_expert_used; k2++) {
w[k2] = w[k2] / sum * expert_weight_scale;
}
}
}
__global__ static void glm_router_select_batch_kernel(
int32_t *selected,
float *weights_out,
float *probs_out,
const float *bias,
const float *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
const uint32_t tok = blockIdx.x;
if (tok >= n_tokens || threadIdx.x != 0u) return;
const float *lg = logits + (uint64_t)tok * n_expert;
float *probs = probs_out + (uint64_t)tok * n_expert;
int32_t *sel = selected + (uint64_t)tok * n_expert_used;
float *w = weights_out + (uint64_t)tok * n_expert_used;
for (uint32_t i = 0; i < n_expert; i++) {
const float p = 1.0f / (1.0f + expf(-lg[i]));
probs[i] = p;
}
/* top-k over probs+bias, ties by smaller index (matches CPU topk_desc) */
bool taken[384];
for (uint32_t i = 0; i < n_expert; i++) taken[i] = false;
float sum = 0.0f;
for (uint32_t k2 = 0; k2 < n_expert_used; k2++) {
int best = -1; float bv = -1e30f;
for (uint32_t i = 0; i < n_expert; i++) {
if (taken[i]) continue;
const float v = probs[i] + bias[i];
if (v > bv) { bv = v; best = (int)i; }
}
taken[best] = true;
sel[k2] = best;
w[k2] = probs[best];
sum += probs[best];
}
if (sum < 6.103515625e-5f) sum = 6.103515625e-5f;
for (uint32_t k2 = 0; k2 < n_expert_used; k2++) {
w[k2] = w[k2] / sum * expert_weight_scale;
}
}
extern "C" int ds4_gpu_glm_router_select_batch_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
if (!selected || !weights || !probs || !logits || !model_map ||
n_expert == 0 || n_expert > 384u || n_expert_used == 0 ||
n_tokens == 0) {
return 0;
}
const uint64_t bb = (uint64_t)n_expert * sizeof(float);
if (bias_offset > model_size || bb > model_size - bias_offset ||
logits->bytes < (uint64_t)n_tokens * n_expert * sizeof(float) ||
selected->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(int32_t) ||
weights->bytes < (uint64_t)n_tokens * n_expert_used * sizeof(float) ||
probs->bytes < (uint64_t)n_tokens * n_expert * sizeof(float)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const float *bias = (const float *)cuda_resolve_weight_ptr(
model_map, bias_offset, bb, logical_tier, "glm_exp_probs_b");
if (!bias) return 0;
if (!getenv("DS4_GLM_ROUTER_SCALAR")) {
glm_router_select_parallel_kernel<<<n_tokens, 256, 0,
cuda_decode_stream()>>>(
(int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
(const float *)logits->ptr,
n_expert, n_expert_used, expert_weight_scale, n_tokens);
} else glm_router_select_batch_kernel<<<n_tokens, 32, 0,
cuda_decode_stream()>>>(
(int32_t *)selected->ptr,
(float *)weights->ptr,
(float *)probs->ptr,
bias,
(const float *)logits->ptr,
n_expert, n_expert_used, expert_weight_scale, n_tokens);
return cuda_ok(cudaGetLastError(), "glm router select batch launch");
}
extern "C" int ds4_gpu_glm_router_select_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale) {
return ds4_gpu_glm_router_select_batch_tensor(selected, weights, probs,
model_map, model_size,
bias_offset, logits,
n_expert, n_expert_used,
expert_weight_scale, 1u);
}
__global__ static void glm_store_compact_kv_kernel(
char *kv_lora_cache,
char *k_rope_cache,
const float *kv_norm,
const float *kv_raw,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
uint32_t cache_f16) {
const uint32_t token = blockIdx.x;
const uint32_t part = blockIdx.y;
if (token >= n_tokens || part > 1u) return;
const uint32_t pos = pos0 + token;
if (pos >= cache_cap) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
if (part == 0u) {
const float *src = kv_norm + (uint64_t)token * kv_lora_dim;
if (cache_f16) {
__half *dst = (__half *)(kv_lora_cache +
(uint64_t)pos * kv_lora_dim * sizeof(__half));
for (uint32_t i = tid; i < kv_lora_dim; i += nth)
dst[i] = __float2half(src[i]);
} else {
float *dst = (float *)(kv_lora_cache +
(uint64_t)pos * kv_lora_dim * sizeof(float));
for (uint32_t i = tid; i < kv_lora_dim; i += nth)
dst[i] = src[i];
}
} else {
const float *src = kv_raw +
(uint64_t)token * kv_raw_dim + kv_lora_dim;
if (cache_f16) {
__half *dst = (__half *)(k_rope_cache +
(uint64_t)pos * qk_rope * sizeof(__half));
for (uint32_t i = tid; i < qk_rope; i += nth)
dst[i] = __float2half(src[i]);
} else {
float *dst = (float *)(k_rope_cache +
(uint64_t)pos * qk_rope * sizeof(float));
for (uint32_t i = tid; i < qk_rope; i += nth)
dst[i] = src[i];
}
}
}
extern "C" int ds4_gpu_glm_store_compact_kv_tensor(
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_norm,
const ds4_gpu_tensor *kv_raw,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16) {
if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw ||
n_tokens == 0 || cache_cap == 0 || pos0 > cache_cap ||
n_tokens > cache_cap - pos0 || kv_lora_dim == 0 ||
kv_lora_dim > UINT32_MAX - qk_rope ||
kv_lora_dim + qk_rope > kv_raw_dim) {
fprintf(stderr,
"ds4: glm compact KV store received invalid dimensions "
"(pos=%u tokens=%u cap=%u raw=%u lora=%u rope=%u)\n",
pos0, n_tokens, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope);
return 0;
}
const uint64_t es = cache_f16 ? sizeof(__half) : sizeof(float);
const uint64_t norm_need =
(uint64_t)n_tokens * kv_lora_dim * sizeof(float);
const uint64_t raw_need =
(uint64_t)n_tokens * kv_raw_dim * sizeof(float);
const uint64_t lora_need = (uint64_t)cache_cap * kv_lora_dim * es;
const uint64_t rope_need = (uint64_t)cache_cap * qk_rope * es;
if (kv_norm->bytes < norm_need || kv_raw->bytes < raw_need ||
kv_lora_cache->bytes < lora_need || k_rope_cache->bytes < rope_need) {
fprintf(stderr,
"ds4: glm compact KV store received undersized buffers "
"(norm=%llu/%llu raw=%llu/%llu lora=%llu/%llu rope=%llu/%llu)\n",
(unsigned long long)kv_norm->bytes,
(unsigned long long)norm_need,
(unsigned long long)kv_raw->bytes,
(unsigned long long)raw_need,
(unsigned long long)kv_lora_cache->bytes,
(unsigned long long)lora_need,
(unsigned long long)k_rope_cache->bytes,
(unsigned long long)rope_need);
return 0;
}
dim3 grid(n_tokens, qk_rope != 0u ? 2u : 1u, 1u);
glm_store_compact_kv_kernel<<<grid, 128>>>(
(char *)kv_lora_cache->ptr,
(char *)k_rope_cache->ptr,
(const float *)kv_norm->ptr,
(const float *)kv_raw->ptr,
pos0, n_tokens, cache_cap, kv_raw_dim, kv_lora_dim, qk_rope,
cache_f16 ? 1u : 0u);
return cuda_ok(cudaGetLastError(), "glm store compact kv launch");
}
__global__ static void glm_store_indexer_k_kernel(
char *cache,
const float *raw_k,
const float *w,
const float *b,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t n_ctx_orig,
float eps,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
uint32_t cache_f16) {
const uint32_t token = blockIdx.x;
if (token >= n_tokens) return;
const uint32_t pos = pos0 + token;
if (pos >= cache_cap) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
const float *src = raw_k + (uint64_t)token * head_dim;
__shared__ float scratch[256];
float sum = 0.0f;
for (uint32_t i = tid; i < head_dim; i += nth) sum += src[i];
scratch[tid] = sum;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) scratch[tid] += scratch[tid + step];
__syncthreads();
}
const float mean = scratch[0] / (float)head_dim;
__syncthreads();
float ss = 0.0f;
for (uint32_t i = tid; i < head_dim; i += nth) {
const float d = src[i] - mean;
ss += d * d;
}
scratch[tid] = ss;
__syncthreads();
for (uint32_t step = nth >> 1; step > 0; step >>= 1) {
if (tid < step) scratch[tid] += scratch[tid + step];
__syncthreads();
}
const float inv = rsqrtf(scratch[0] / (float)head_dim + eps);
float corr_dims[2] = {0.0f, 0.0f};
if (ext_factor != 0.0f) {
corr_dims[0] = fmaxf(0.0f,
floorf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig,
beta_fast, freq_base)));
corr_dims[1] = fminf((float)rot_dim - 1.0f,
ceilf(glm_rope_yarn_corr_factor_dev((int)rot_dim, (int)n_ctx_orig,
beta_slow, freq_base)));
}
const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)rot_dim;
for (uint32_t i = tid; i < head_dim; i += nth) {
float v0, v1; bool pair = false;
if (i < rot_dim) {
if ((i & 1u) != 0u) continue;
const float theta = theta_base * powf(freq_base, inv_ndims * (float)i);
float ct, st;
glm_rope_yarn_dev(theta, freq_scale, corr_dims, (int)i,
ext_factor, attn_factor, &ct, &st);
const float x0 = (src[i] - mean) * inv * w[i] + b[i];
const float x1 = (src[i + 1u] - mean) * inv * w[i + 1u] + b[i + 1u];
v0 = x0 * ct - x1 * st;
v1 = x0 * st + x1 * ct;
pair = true;
} else {
v0 = (src[i] - mean) * inv * w[i] + b[i];
}
if (cache_f16) {
__half *dst = (__half *)(cache + (uint64_t)pos * head_dim * sizeof(__half));
dst[i] = __float2half(v0);
if (pair) dst[i + 1u] = __float2half(v1);
} else {
float *dst = (float *)(cache + (uint64_t)pos * head_dim * sizeof(float));
dst[i] = v0;
if (pair) dst[i + 1u] = v1;
}
}
}
extern "C" int ds4_gpu_glm_store_indexer_k_tensor(
ds4_gpu_tensor *indexer_key_cache,
const ds4_gpu_tensor *raw_k,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t bias_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t n_ctx_orig,
float eps,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
if (!indexer_key_cache || !raw_k || !model_map || n_tokens == 0 ||
head_dim == 0 || head_dim > 256u || (rot_dim & 1u) != 0u) {
return 0;
}
const uint64_t wb = (uint64_t)head_dim * sizeof(float);
if (weight_offset > model_size || wb > model_size - weight_offset ||
bias_offset > model_size || wb > model_size - bias_offset ||
raw_k->bytes < (uint64_t)n_tokens * head_dim * sizeof(float) ||
indexer_key_cache->bytes <
(uint64_t)cache_cap * head_dim *
(cache_f16 ? sizeof(__half) : sizeof(float))) {
return 0;
}
const int logical_tier = cuda_current_tier();
const float *w = (const float *)cuda_resolve_weight_ptr(
model_map, weight_offset, wb, logical_tier, "glm_indexer_k_norm");
const float *b = (const float *)cuda_resolve_weight_ptr(
model_map, bias_offset, wb, logical_tier, "glm_indexer_k_norm_b");
if (!w || !b) return 0;
glm_store_indexer_k_kernel<<<n_tokens, 128>>>(
(char *)indexer_key_cache->ptr,
(const float *)raw_k->ptr,
w, b, pos0, n_tokens, cache_cap, head_dim, rot_dim, n_ctx_orig,
eps, freq_base, freq_scale, ext_factor, attn_factor,
beta_fast, beta_slow, cache_f16 ? 1u : 0u);
return cuda_ok(cudaGetLastError(), "glm store indexer k launch");
}
extern "C" int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
const ds4_gpu_stream_expert_table *table,
const ds4_gpu_tensor *selected,
uint32_t n_selected) {
if (!g_ssd_streaming_mode) return 1;
if (!table || !selected || n_selected == 0 ||
selected->bytes < (uint64_t)n_selected * sizeof(int32_t)) {
return 0;
}
std::vector<int32_t> ids;
try {
ids.resize(n_selected);
} catch (...) {
return 0;
}
if (!cuda_ok(cudaMemcpy(ids.data(), selected->ptr,
(size_t)n_selected * sizeof(int32_t),
cudaMemcpyDeviceToHost),
"GLM streaming selected-id read")) {
return 0;
}
return cuda_stream_selected_cache_begin_load(table, ids.data(), n_selected);
}
__global__ static void glm_value_project_q8_0_batch_heads_kernel(
float *heads,
const char *weight,
const float *lora,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim,
uint64_t row_bytes) {
const uint32_t head = blockIdx.x;
const uint32_t token = blockIdx.y;
if (head >= n_head || token >= n_tokens) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
extern __shared__ float xsh[];
const float *src = lora + (uint64_t)token * n_head * kv_lora_dim +
(uint64_t)head * kv_lora_dim;
float *out = heads + (uint64_t)token * n_head * value_dim +
(uint64_t)head * value_dim;
for (uint32_t j = tid; j < kv_lora_dim; j += nth) xsh[j] = src[j];
__syncthreads();
for (uint32_t d = tid; d < value_dim; d += nth) {
const char *row = weight + ((uint64_t)head * value_dim + d) * row_bytes;
out[d] = glm_q8_0_dot_row_dev(row, xsh, kv_lora_dim);
}
}
/* Reuse each Q8 row across a token tile while retaining the scalar kernel's
* block/k accumulation order independently for every output token. */
template <uint32_t token_tile>
__global__ static void glm_value_project_q8_0_batch_heads_tiled_kernel(
float *heads,
const char *weight,
const float *lora,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim,
uint64_t row_bytes) {
const uint32_t head = blockIdx.x;
const uint32_t token0 = blockIdx.y * token_tile;
if (head >= n_head || token0 >= n_tokens) return;
const uint32_t tid = threadIdx.x;
const uint32_t nth = blockDim.x;
extern __shared__ float xsh[];
#pragma unroll
for (uint32_t t = 0; t < token_tile; t++) {
const uint32_t token = token0 + t;
if (token >= n_tokens) break;
const float *src = lora + (uint64_t)token * n_head * kv_lora_dim +
(uint64_t)head * kv_lora_dim;
for (uint32_t j = tid; j < kv_lora_dim; j += nth) {
xsh[(uint64_t)t * kv_lora_dim + j] = src[j];
}
}
__syncthreads();
for (uint32_t od = tid; od < value_dim; od += nth) {
const char *row = weight +
((uint64_t)head * value_dim + od) * row_bytes;
float acc[token_tile] = { 0.0f };
const uint32_t nb = kv_lora_dim >> 5;
for (uint32_t b = 0; b < nb; b++) {
const char *blk = row + (uint64_t)b * 34u;
const float d = __half2float(*(const __half *)blk);
const int8_t *q = (const int8_t *)(blk + 2);
float s[token_tile] = { 0.0f };
#pragma unroll 8
for (uint32_t k = 0; k < 32u; k++) {
const float w = (float)q[k];
#pragma unroll
for (uint32_t t = 0; t < token_tile; t++) {
if (token0 + t < n_tokens) {
s[t] += w * xsh[(uint64_t)t * kv_lora_dim +
b * 32u + k];
}
}
}
#pragma unroll
for (uint32_t t = 0; t < token_tile; t++) {
if (token0 + t < n_tokens) acc[t] += d * s[t];
}
}
#pragma unroll
for (uint32_t t = 0; t < token_tile; t++) {
const uint32_t token = token0 + t;
if (token < n_tokens) {
heads[(uint64_t)token * n_head * value_dim +
(uint64_t)head * value_dim + od] = acc[t];
}
}
}
}
extern "C" int ds4_gpu_glm_value_project_typed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *lora,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim) {
if (!heads || !lora || !model_map || n_tokens == 0 || n_head == 0 ||
kv_lora_dim == 0 || (kv_lora_dim & 31u) != 0u || value_dim == 0) {
return 0;
}
if (weight_type != 8u) {
fprintf(stderr, "ds4: glm value project: unsupported type %u\n",
weight_type);
return 0;
}
const uint64_t row_bytes = ((uint64_t)kv_lora_dim / 32u) * 34u;
const uint64_t wbytes = (uint64_t)n_head * value_dim * row_bytes;
if (weight_offset > model_size || wbytes > model_size - weight_offset ||
lora->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) ||
heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float)) {
return 0;
}
const int logical_tier = cuda_current_tier();
const char *w = (const char *)cuda_resolve_weight_ptr(
model_map, weight_offset, wbytes, logical_tier, "glm_v_b");
if (!w) return 0;
if (n_tokens >= 16u && getenv("DS4_GLM_VALUE_NO_TILE16") == NULL) {
dim3 grid(n_head, (n_tokens + 15u) / 16u, 1);
const size_t shmem = 16ull * kv_lora_dim * sizeof(float);
glm_value_project_q8_0_batch_heads_tiled_kernel<16><<<grid, 128, shmem>>>(
(float *)heads->ptr, w, (const float *)lora->ptr,
n_tokens, n_head, kv_lora_dim, value_dim, row_bytes);
return cuda_ok(cudaGetLastError(), "glm value project tile16 launch");
}
dim3 grid(n_head, n_tokens, 1);
const size_t shmem = (size_t)kv_lora_dim * sizeof(float);
glm_value_project_q8_0_batch_heads_kernel<<<grid, 128, shmem>>>(
(float *)heads->ptr, w, (const float *)lora->ptr,
n_tokens, n_head, kv_lora_dim, value_dim, row_bytes);
return cuda_ok(cudaGetLastError(), "glm value project launch");
}
/* Decode-time (n_tok small) quant matvec. The Metal "mpp/model-view"
* variant is a bandwidth-tuned matvec; on CUDA the generic quant matmul
* already dispatches per type, so delegate. Revisit in the perf pass. */
extern "C" int ds4_gpu_matmul_quant_tensor(ds4_gpu_tensor *out,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t weight_type, uint64_t in_dim, uint64_t out_dim,
const ds4_gpu_tensor *x, uint64_t n_tok);
extern "C" int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_quant_tensor(out, model_map, model_size,
weight_offset, weight_type,
in_dim, out_dim, x, n_tok);
}
extern "C" int ds4_gpu_matmul_quant_rows_scalar_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
fprintf(stderr, "ds4: CUDA stub called: ds4_gpu_matmul_quant_rows_scalar_tensor\n");
return 0;
}
static int cuda_matmul_mmq_dense_quant(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
uint64_t block_elems = 0u;
uint64_t block_bytes = 0u;
const char *label = NULL;
switch (weight_type) {
case 10u: block_elems = 256u; block_bytes = 84u; label = "Q2_K"; break;
case 12u: block_elems = 256u; block_bytes = 144u; label = "Q4_K"; break;
case 16u: block_elems = 256u; block_bytes = 66u; label = "IQ2_XXS"; break;
case 39u: block_elems = 32u; block_bytes = 17u; label = "MXFP4"; break;
default: return 0;
}
if (!out || !x || !model_map || in_dim == 0u || out_dim == 0u ||
n_tok == 0u || in_dim % block_elems != 0u ||
in_dim > INT_MAX || out_dim > INT_MAX || n_tok > INT_MAX ||
out_dim > UINT64_MAX / ((in_dim / block_elems) * block_bytes)) {
return 0;
}
const uint64_t row_bytes = (in_dim / block_elems) * block_bytes;
const uint64_t weight_bytes = out_dim * row_bytes;
if (in_dim > UINT64_MAX / n_tok || out_dim > UINT64_MAX / n_tok) {
return 0;
}
const uint64_t input_elements = n_tok * in_dim;
const uint64_t output_elements = n_tok * out_dim;
if (weight_offset > model_size ||
weight_bytes > model_size - weight_offset ||
input_elements > UINT64_MAX / sizeof(float) ||
output_elements > UINT64_MAX / sizeof(float) ||
x->bytes < input_elements * sizeof(float) ||
out->bytes < output_elements * sizeof(float)) {
return 0;
}
const int tier = ds4_tensor_device_idx(out);
const void *weights = cuda_resolve_weight_ptr(
model_map, weight_offset, weight_bytes, tier, label);
if (!weights) return 0;
int rc = -1;
switch (weight_type) {
case 10u:
rc = ds4_mmq_q2_K_dense(weights, (const float *)x->ptr,
(float *)out->ptr, (int)out_dim, (int)n_tok, (int)in_dim,
cuda_decode_stream());
break;
case 12u:
rc = ds4_mmq_q4_K_dense(weights, (const float *)x->ptr,
(float *)out->ptr, (int)out_dim, (int)n_tok, (int)in_dim,
cuda_decode_stream());
break;
case 16u:
rc = ds4_mmq_iq2_xxs_dense(weights, (const float *)x->ptr,
(float *)out->ptr, (int)out_dim, (int)n_tok, (int)in_dim,
cuda_decode_stream());
break;
case 39u:
rc = ds4_mmq_mxfp4_dense(weights, (const float *)x->ptr,
(float *)out->ptr, (int)out_dim, (int)n_tok, (int)in_dim,
cuda_decode_stream());
break;
}
if (rc != 0) {
fprintf(stderr, "ds4: CUDA dense %s MMQ failed (%d)\n", label, rc);
return 0;
}
return 1;
}
extern "C" int ds4_gpu_matmul_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
switch (weight_type) {
case 8u: /* Q8_0 */
return ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size,
weight_offset, in_dim, out_dim,
x, n_tok);
case 1u: /* F16 */
return ds4_gpu_matmul_f16_tensor(out, model_map, model_size,
weight_offset, in_dim, out_dim,
x, n_tok);
case 10u: /* Q2_K */
case 12u: /* Q4_K */
case 16u: /* IQ2_XXS */
case 39u: /* MXFP4 */
return cuda_matmul_mmq_dense_quant(
out, model_map, model_size, weight_offset, weight_type,
in_dim, out_dim, x, n_tok);
default:
fprintf(stderr, "ds4: matmul_quant: unsupported type %u\n",
weight_type);
return 0;
}
}
extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) {
/* GLM graph memory guard: on this backend the model weights are
* distributed across all devices by the multi-tier placement, so the
* relevant budget is the aggregate VRAM. */
int n = 0;
if (cudaGetDeviceCount(&n) != cudaSuccess || n <= 0) return 0;
size_t free_b = 0, total_b = 0;
if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) return 0;
return (uint64_t)total_b * (uint64_t)n;
}
extern "C" int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) {
(void)selected;
(void)n_selected;
return 1;
}
extern "C" void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) {
(void)enabled; /* CUDA stages the current batch's selected experts. */
}
extern "C" void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) {
g_glm_mtp_verify_mode = enabled;
}
extern "C" int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model_size, const uint64_t *offsets, const uint64_t *sizes, uint32_t count, uint64_t max_tensor_bytes) {
(void)max_tensor_bytes;
if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) {
return 0;
}
for (uint32_t i = 0; i < count; i++) {
if (offsets[i] > model_size || sizes[i] == 0 ||
sizes[i] > model_size - offsets[i]) {
return 0;
}
}
if (!cuda_register_model_map(model_map, model_size, false)) return 0;
if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL) {
for (uint32_t i = 0; i < count; i++) {
(void)cuda_model_prefetch_range(model_map, model_size,
offsets[i], sizes[i]);
}
}
return 1;
}
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid,
const void *model_map, uint64_t model_size,
uint64_t gate_offset, uint64_t up_offset,
uint64_t in_dim, uint64_t out_dim,
const ds4_gpu_tensor *x, uint64_t n_tok, float clamp);
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(gate, up, mid,
model_map, model_size, gate_offset, up_offset,
in_dim, out_dim, x, 1, clamp);
}
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp) {
(void)gate; (void)up; (void)mid; (void)model_map; (void)model_size;
(void)gate_offset; (void)up_offset; (void)in_dim; (void)out_dim;
(void)x; (void)n_tok; (void)clamp;
return 0;
}
/* Fused single-token shared-expert gate+up+swiglu: one warp per output
* row computes both q8_0 dots against a shared-staged f32 x and writes
* silu(gate)*up directly. Falls back to the split path for n_tok > 1. */
__global__ static void glm_shared_gate_up_swiglu_one_kernel(
float *mid,
const char *gw,
const char *uw,
const float *x,
uint32_t in_dim,
uint32_t out_dim,
float clamp) {
extern __shared__ float glm_sgu_sh[];
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
glm_sgu_sh[i] = x[i];
}
__syncthreads();
const uint32_t r = blockIdx.x * warps + warp;
if (r >= out_dim) return;
const uint32_t nblk = in_dim >> 5;
const uint64_t row_bytes = (uint64_t)nblk * 34u;
const char *grow = gw + (uint64_t)r * row_bytes;
const char *urow = uw + (uint64_t)r * row_bytes;
float g = 0.0f, u = 0.0f;
for (uint32_t blk = lane; blk < nblk; blk += 32u) {
const char *gb = grow + (uint64_t)blk * 34u;
const char *ub = urow + (uint64_t)blk * 34u;
const float gd = __half2float(*(const __half *)gb);
const float ud = __half2float(*(const __half *)ub);
const int8_t *gq = (const int8_t *)(gb + 2);
const int8_t *uq = (const int8_t *)(ub + 2);
const float *xs = glm_sgu_sh + blk * 32u;
float gs = 0.0f, us = 0.0f;
#pragma unroll 8
for (int k = 0; k < 32; k++) {
gs += (float)gq[k] * xs[k];
us += (float)uq[k] * xs[k];
}
g += gd * gs;
u += ud * us;
}
for (int off = 16; off > 0; off >>= 1) {
g += __shfl_down_sync(0xffffffffu, g, off);
u += __shfl_down_sync(0xffffffffu, u, off);
}
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (g > clamp) g = clamp;
if (u > clamp) u = clamp;
if (u < -clamp) u = -clamp;
}
mid[r] = (g / (1.0f + expf(-g))) * u;
}
}
/* Two-token verifier variant of the decode kernel above. Each warp loads a
* gate/up weight row once, while each token keeps the decode kernel's block
* order and warp reduction tree independently. */
__global__ static void glm_shared_gate_up_swiglu_tok2_exact_kernel(
float *mid,
const char *gw,
const char *uw,
const float *x,
uint32_t in_dim,
uint32_t out_dim,
float clamp) {
extern __shared__ float glm_sgu2_sh[];
float *x0 = glm_sgu2_sh;
float *x1 = glm_sgu2_sh + in_dim;
const uint32_t warps = blockDim.x >> 5;
const uint32_t warp = threadIdx.x >> 5;
const uint32_t lane = threadIdx.x & 31u;
for (uint32_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
x0[i] = x[i];
x1[i] = x[in_dim + i];
}
__syncthreads();
const uint32_t r = blockIdx.x * warps + warp;
if (r >= out_dim) return;
const uint32_t nblk = in_dim >> 5;
const uint64_t row_bytes = (uint64_t)nblk * 34u;
const char *grow = gw + (uint64_t)r * row_bytes;
const char *urow = uw + (uint64_t)r * row_bytes;
float g0 = 0.0f, u0 = 0.0f;
float g1 = 0.0f, u1 = 0.0f;
for (uint32_t blk = lane; blk < nblk; blk += 32u) {
const char *gb = grow + (uint64_t)blk * 34u;
const char *ub = urow + (uint64_t)blk * 34u;
const float gd = __half2float(*(const __half *)gb);
const float ud = __half2float(*(const __half *)ub);
const int8_t *gq = (const int8_t *)(gb + 2);
const int8_t *uq = (const int8_t *)(ub + 2);
const float *xs0 = x0 + blk * 32u;
const float *xs1 = x1 + blk * 32u;
float gs0 = 0.0f, us0 = 0.0f;
float gs1 = 0.0f, us1 = 0.0f;
#pragma unroll 8
for (int k = 0; k < 32; k++) {
const float gk = (float)gq[k];
const float uk = (float)uq[k];
gs0 += gk * xs0[k];
us0 += uk * xs0[k];
gs1 += gk * xs1[k];
us1 += uk * xs1[k];
}
g0 += gd * gs0;
u0 += ud * us0;
g1 += gd * gs1;
u1 += ud * us1;
}
for (int off = 16; off > 0; off >>= 1) {
g0 += __shfl_down_sync(0xffffffffu, g0, off);
u0 += __shfl_down_sync(0xffffffffu, u0, off);
g1 += __shfl_down_sync(0xffffffffu, g1, off);
u1 += __shfl_down_sync(0xffffffffu, u1, off);
}
if (lane == 0u) {
if (clamp > 1.0e-6f) {
if (g0 > clamp) g0 = clamp;
if (u0 > clamp) u0 = clamp;
if (u0 < -clamp) u0 = -clamp;
if (g1 > clamp) g1 = clamp;
if (u1 > clamp) u1 = clamp;
if (u1 < -clamp) u1 = -clamp;
}
mid[r] = (g0 / (1.0f + expf(-g0))) * u0;
mid[out_dim + r] = (g1 / (1.0f + expf(-g1))) * u1;
}
}
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp) {
if (!gate || !up || !mid || !x || n_tok == 0) return 0;
if (n_tok == 2 && (in_dim & 31u) == 0u &&
g_glm_mtp_verify_mode &&
getenv("DS4_GLM_MTP_NO_SHARED_TOK2") == NULL &&
mid->bytes >= 2u * out_dim * sizeof(float) &&
x->bytes >= 2u * in_dim * sizeof(float)) {
const uint64_t row_bytes = (in_dim / 32u) * 34u;
const uint64_t wb = out_dim * row_bytes;
if (gate_offset <= model_size && wb <= model_size - gate_offset &&
up_offset <= model_size && wb <= model_size - up_offset) {
const int logical_tier = cuda_current_tier();
const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset,
wb, logical_tier, "glm_shared_gate");
const char *uw = cuda_resolve_weight_ptr(model_map, up_offset,
wb, logical_tier, "glm_shared_up");
if (gw && uw) {
const uint32_t warps = 8u;
const uint32_t sh = 2u * (uint32_t)in_dim * sizeof(float);
glm_shared_gate_up_swiglu_tok2_exact_kernel
<<<(unsigned)((out_dim + warps - 1u) / warps),
warps * 32u, sh>>>(
(float *)mid->ptr, gw, uw, (const float *)x->ptr,
(uint32_t)in_dim, (uint32_t)out_dim, clamp);
return cuda_ok(cudaGetLastError(),
"glm shared swiglu tok2 exact");
}
}
}
if (n_tok == 1 &&
cuda_matmul_q8_0_aligned_pair_try(
gate, up, model_map, model_size,
gate_offset, up_offset, in_dim, out_dim, out_dim, x) > 0) {
return ds4_gpu_swiglu_tensor(mid, gate, up,
(uint32_t)out_dim, clamp, 1.0f);
}
if (n_tok == 1 && (in_dim & 31u) == 0u &&
!getenv("DS4_GLM_SHARED_SPLIT") &&
mid->bytes >= out_dim * sizeof(float) &&
x->bytes >= in_dim * sizeof(float)) {
const uint64_t row_bytes = (in_dim / 32u) * 34u;
const uint64_t wb = out_dim * row_bytes;
if (gate_offset <= model_size && wb <= model_size - gate_offset &&
up_offset <= model_size && wb <= model_size - up_offset) {
const int logical_tier = cuda_current_tier();
const char *gw = cuda_resolve_weight_ptr(model_map, gate_offset,
wb, logical_tier, "glm_shared_gate");
const char *uw = cuda_resolve_weight_ptr(model_map, up_offset,
wb, logical_tier, "glm_shared_up");
if (gw && uw) {
const uint32_t warps = 8u;
const uint32_t sh = (uint32_t)in_dim * sizeof(float);
glm_shared_gate_up_swiglu_one_kernel
<<<(unsigned)((out_dim + warps - 1u) / warps),
warps * 32u, sh>>>(
(float *)mid->ptr, gw, uw, (const float *)x->ptr,
(uint32_t)in_dim, (uint32_t)out_dim, clamp);
return cuda_ok(cudaGetLastError(), "glm shared swiglu one");
}
}
}
if (!ds4_gpu_matmul_q8_0_tensor(gate, model_map, model_size, gate_offset,
in_dim, out_dim, x, n_tok) ||
!ds4_gpu_matmul_q8_0_tensor(up, model_map, model_size, up_offset,
in_dim, out_dim, x, n_tok)) {
return 0;
}
return ds4_gpu_swiglu_tensor(mid, gate, up,
(uint32_t)(out_dim * n_tok), clamp, 1.0f);
}
extern "C" int ds4_gpu_shared_mid_swiglu_q8_0_tensor(
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
static ds4_gpu_tensor *gu_scratch[DS4_MAX_GPUS][2] = {{0}};
const int dev = cuda_current_tier();
const uint64_t need = out_dim * sizeof(float);
for (int i = 0; i < 2; i++) {
if (!gu_scratch[dev][i] || gu_scratch[dev][i]->bytes < need) {
if (gu_scratch[dev][i]) ds4_gpu_tensor_free(gu_scratch[dev][i]);
gu_scratch[dev][i] = ds4_gpu_tensor_alloc(need);
}
if (!gu_scratch[dev][i]) return 0;
}
return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
gu_scratch[dev][0], gu_scratch[dev][1], mid,
model_map, model_size, gate_offset, up_offset,
in_dim, out_dim, x, 1, clamp);
}
extern "C" int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) {
if (event_value) *event_value = 1;
return cuda_ok(cudaDeviceSynchronize(), "selected readback signal");
}
extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected) {
return cuda_stream_selected_cache_begin_load(table, selected_ids,
n_selected);
}
extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
if (!gate_expert_bytes || !down_expert_bytes ||
gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2u) return 0;
const uint64_t bytes = 2u * gate_expert_bytes + down_expert_bytes;
if (!g_stream_expert_bytes) return g_stream_expert_budget;
const uint64_t budget = g_stream_expert_budget > UINT64_MAX / g_stream_expert_bytes ?
UINT64_MAX : (uint64_t)g_stream_expert_budget * g_stream_expert_bytes;
return (uint32_t)std::min(budget / bytes, (uint64_t)INT32_MAX);
}
extern "C" int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t count) {
if (!dst || !src) return 0;
if (count == 0) return 1;
if (count > UINT64_MAX / sizeof(float) ||
count > UINT64_MAX / sizeof(__half)) {
return 0;
}
const uint64_t src_bytes = count * sizeof(float);
const uint64_t dst_bytes = count * sizeof(__half);
if (src_offset > src->bytes || src_bytes > src->bytes - src_offset ||
dst_offset > dst->bytes || dst_bytes > dst->bytes - dst_offset ||
ds4_tensor_device_idx(dst) != ds4_tensor_device_idx(src)) {
return 0;
}
const int tier = ds4_tensor_device_idx(dst);
if (ds4_gpu_set_current_device(tier) != 0) return 0;
const uint64_t blocks = (count + 255u) / 256u;
if (blocks > UINT32_MAX) return 0;
f32_to_f16_kernel<<<(unsigned)blocks, 256>>>(
(__half *)((char *)dst->ptr + dst_offset),
(const float *)((const char *)src->ptr + src_offset),
count);
return cuda_ok(cudaGetLastError(), "tensor f32-to-f16 copy launch");
}
extern "C" int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *tensor,
uint64_t offset,
void *data,
uint64_t bytes,
uint64_t event_value,
const char *label) {
(void)event_value;
if (!tensor || !data || offset > tensor->bytes ||
bytes > tensor->bytes - offset) {
return 0;
}
if (!cuda_ok(cudaDeviceSynchronize(),
label ? label : "selected readback wait")) {
return 0;
}
return cuda_ok(cudaMemcpy(data, (const char *)tensor->ptr + offset,
(size_t)bytes, cudaMemcpyDeviceToHost),
"selected tensor read");
}
static struct {
ds4_gpu_tp_exchange_fn row;
ds4_gpu_tp_batch_exchange_fn batch;
ds4_gpu_tp_big_exchange_fn big;
void *ud, *staging;
uint64_t vec_bytes, staging_bytes, row_seq, batch_seq;
bool failed;
} g_cuda_tp;
extern "C" void ds4_gpu_tp_shutdown(void) {
if (g_cuda_tp.row) (void)cudaDeviceSynchronize();
if (g_cuda_tp.staging) (void)cudaFreeHost(g_cuda_tp.staging);
g_cuda_tp = {};
}
extern "C" int ds4_gpu_tp_init(uint32_t rank, ds4_gpu_tensor *slab,
uint64_t gpu_flags_off, uint64_t out_off, uint64_t vec_bytes,
ds4_gpu_tp_exchange_fn fn, void *ud) {
if (g_cuda_tp.row || rank > 1 || !slab || !fn || !vec_bytes ||
out_off > slab->bytes || vec_bytes > slab->bytes - out_off ||
gpu_flags_off >= slab->bytes) return 0;
g_cuda_tp.row = fn;
g_cuda_tp.ud = ud;
g_cuda_tp.vec_bytes = vec_bytes;
return 1;
}
extern "C" void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) {
g_cuda_tp.batch = fn;
}
extern "C" void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) {
g_cuda_tp.big = fn;
}
extern "C" void ds4_gpu_tp_set_session_batch_mode(int enabled) { (void)enabled; }
extern "C" int ds4_gpu_tp_decode_split_flush_safe(void) { return 1; }
extern "C" int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) {
if (!g_cuda_tp.row || g_cuda_tp.failed) return 0;
/* Do not leave a GPU polling kernel waiting for a remote process. */
const int ok = cuda_ok(cudaStreamSynchronize(cuda_decode_stream()), "TP row arrival") &&
g_cuda_tp.row(g_cuda_tp.ud, layer, gate, ++g_cuda_tp.row_seq);
if (!ok) g_cuda_tp.failed = true;
return ok;
}
extern "C" int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) {
if (!g_cuda_tp.row || !g_cuda_tp.batch || g_cuda_tp.failed || !rows || rows > 8) return 0;
const int ok = cuda_ok(cudaStreamSynchronize(cuda_decode_stream()), "TP batch arrival") &&
g_cuda_tp.batch(g_cuda_tp.ud, layer, rows, ++g_cuda_tp.batch_seq);
if (!ok) g_cuda_tp.failed = true;
return ok;
}
extern "C" int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t, ds4_gpu_tensor *in_t, uint64_t bytes) {
if (!g_cuda_tp.row || !g_cuda_tp.big || g_cuda_tp.failed || !out_t || !in_t ||
!rows || bytes != (uint64_t)rows * g_cuda_tp.vec_bytes ||
bytes > SIZE_MAX / 2 || bytes > out_t->bytes || bytes > in_t->bytes) return 0;
if (bytes > g_cuda_tp.staging_bytes) {
if (g_cuda_tp.staging) (void)cudaFreeHost(g_cuda_tp.staging);
g_cuda_tp.staging = NULL;
g_cuda_tp.staging_bytes = 0;
if (!cuda_ok(cudaHostAlloc(&g_cuda_tp.staging, (size_t)bytes * 2,
cudaHostAllocDefault), "TP bulk staging")) {
g_cuda_tp.failed = true;
return 0;
}
g_cuda_tp.staging_bytes = bytes;
}
void *peer = (char *)g_cuda_tp.staging + g_cuda_tp.staging_bytes;
const int ok = cuda_ok(cudaStreamSynchronize(cuda_decode_stream()), "TP bulk arrival") &&
ds4_gpu_tensor_read(out_t, 0, g_cuda_tp.staging, bytes) &&
g_cuda_tp.big(g_cuda_tp.ud, layer, ++g_cuda_tp.batch_seq,
g_cuda_tp.staging, peer, bytes) &&
ds4_gpu_tensor_write(in_t, 0, peer, bytes);
if (!ok) g_cuda_tp.failed = true;
return ok;
}
extern "C" void ds4_gpu_tp_set_attn_head_split(int enabled) {
(void)enabled; /* Mac network-TP head split: no-op on CUDA */
}
extern "C" int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) {
(void)event_value;
return cuda_ok(cudaDeviceSynchronize(),
label ? label : "selected readback wait");
}
/* Compatibility surface shared with the canonical Metal/ROCm graph. CUDA
* either delegates to its equivalent primitive or reports an unavailable
* optional fast path so the graph can use its established fallback. */
extern "C" int ds4_gpu_commit_and_wait_selected_readback(
uint64_t event_value, const char *label) {
(void)event_value;
return cuda_ok(cudaDeviceSynchronize(),
label ? label : "selected readback wait");
}
extern "C" int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) {
const int ok = ds4_gpu_set_model_fd(fd);
if (ok) g_model_fd_host_base = model_map;
return ok;
}
extern "C" int ds4_gpu_pro_q4_expert_table_auto_available(void) {
return 0;
}
extern "C" int ds4_gpu_preload_q4_expert_tables(
const void *model_map, uint64_t model_size,
uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset,
uint64_t gate_expert_bytes, uint64_t down_expert_bytes,
uint32_t n_total_expert) {
(void)model_map; (void)model_size;
(void)gate_offset; (void)up_offset; (void)down_offset;
(void)gate_expert_bytes; (void)down_expert_bytes;
(void)n_total_expert;
return 1;
}
extern "C" void ds4_gpu_set_glm_model(bool enabled) {
(void)enabled;
}
extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) {
g_ssd_streaming_mode = enabled ? 1 : 0;
cuda_stream_selected_cache_invalidate();
if (!g_ssd_streaming_mode) cuda_stream_selected_cache_release();
}
extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) {
if (experts != g_stream_expert_budget) cuda_stream_selected_cache_release();
g_stream_expert_budget = experts;
}
extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) {
if (bytes != g_stream_expert_bytes) cuda_stream_selected_cache_release();
g_stream_expert_bytes = bytes;
}
extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) {
return g_stream_expert_budget;
}
extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) {
return (uint32_t)g_stream_expert_by_gate.size();
}
extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) {
}
extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) {
cuda_stream_selected_cache_release();
}
extern "C" int ds4_gpu_stream_expert_cache_seed_selected(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected) {
return !n_selected || cuda_stream_selected_cache_begin_load(table, selected_ids, n_selected);
}
extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_tokens,
uint32_t n_selected) {
if (n_tokens == 0 || n_selected == 0 ||
(uint64_t)n_tokens * n_selected > UINT32_MAX) {
return 0;
}
return cuda_stream_selected_cache_begin_load(
table, selected_ids, n_tokens * n_selected);
}
extern "C" int ds4_gpu_stream_expert_cache_seed_experts(
const ds4_gpu_stream_expert_table *table,
const int32_t *expert_ids,
const uint32_t *expert_priorities,
uint32_t n_experts) {
(void)expert_priorities;
return !n_experts || cuda_stream_selected_cache_begin_load(table, expert_ids, n_experts);
}
extern "C" int ds4_gpu_argmax_tensor(
ds4_gpu_tensor *out_idx,
const ds4_gpu_tensor *logits,
uint32_t n_vocab) {
return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1u, 1u);
}
extern "C" int ds4_gpu_embed_token_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd) {
return ds4_gpu_embed_token_quant_tensor(out, model_map, model_size,
weight_offset, 8u, n_vocab,
token, n_embd);
}
extern "C" int ds4_gpu_embed_tokens_q8_0_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
return ds4_gpu_embed_tokens_quant_tensor(out, tokens, model_map,
model_size, weight_offset, 8u,
n_vocab, n_tokens, n_embd);
}
extern "C" int ds4_gpu_glm_k_b_project_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head) {
return ds4_gpu_glm_k_b_project_typed_tensor(
out, kv_norm, model_map, model_size, weight_offset, 8u,
n_tokens, kv_lora_dim, qk_nope, n_head);
}
extern "C" int ds4_gpu_matmul_q8_0_kslice_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off) {
if (!x || x_elem_off > x->bytes / sizeof(float) ||
k_cnt > x->bytes / sizeof(float) - x_elem_off) {
return 0;
}
ds4_gpu_tensor x_slice = *x;
x_slice.ptr = (char *)x->ptr + x_elem_off * sizeof(float);
x_slice.bytes = k_cnt * sizeof(float);
x_slice.owner = 0;
return ds4_gpu_matmul_q8_0_kslice_rows_tensor(
out, model_map, model_size, weight_offset,
full_in_dim, out_dim, k_off, k_cnt, &x_slice, 1u);
}
extern "C" int ds4_gpu_matmul_quant_kslice_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off) {
if (weight_type != 8u) return 0;
return ds4_gpu_matmul_q8_0_kslice_tensor(
out, model_map, model_size, weight_offset,
full_in_dim, k_off, k_cnt, out_dim, x, x_elem_off);
}
extern "C" int ds4_gpu_matmul_q8_0_f16_out_tensor(
ds4_gpu_tensor *out_h,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
(void)out_h; (void)model_map; (void)model_size; (void)weight_offset;
(void)in_dim; (void)out_dim; (void)x; (void)n_tok;
return 0;
}
extern "C" int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *q_half,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x,
uint32_t n_tok, uint32_t n_head, uint32_t head_dim,
uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse,
float freq_base, float freq_scale, float ext_factor,
float attn_factor, float beta_fast, float beta_slow, float eps) {
(void)out; (void)q_half; (void)model_map; (void)model_size;
(void)weight_offset; (void)in_dim; (void)out_dim; (void)x;
(void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0;
(void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale;
(void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow;
(void)eps;
return 0;
}
extern "C" int ds4_gpu_attention_prefill_raw_heads_range_tensor(
ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size,
uint64_t sinks_offset, const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv, uint32_t q_row0, uint32_t n_q,
uint32_t n_kv, uint32_t window, uint32_t n_head,
uint32_t head_dim) {
(void)heads; (void)model_map; (void)model_size; (void)sinks_offset;
(void)q; (void)raw_kv; (void)q_row0; (void)n_q; (void)n_kv;
(void)window; (void)n_head; (void)head_dim;
return 0;
}
extern "C" int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(
ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size,
uint64_t sinks_offset, const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv, const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16, uint32_t q_row0, uint32_t n_q,
uint32_t n_tokens, uint32_t n_comp, uint32_t window,
uint32_t ratio, uint32_t n_head, uint32_t head_dim) {
(void)heads; (void)model_map; (void)model_size; (void)sinks_offset;
(void)q; (void)raw_kv; (void)comp_kv; (void)comp_kv_f16;
(void)q_row0; (void)n_q; (void)n_tokens; (void)n_comp; (void)window;
(void)ratio; (void)n_head; (void)head_dim;
return 0;
}
extern "C" int ds4_gpu_attention_output_q8_batch_f16_tensor(
ds4_gpu_tensor *out_h, ds4_gpu_tensor *low,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
uint64_t group_dim, uint64_t rank, uint32_t n_groups,
uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) {
(void)out_h; (void)low; (void)model_map; (void)model_size;
(void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank;
(void)n_groups; (void)out_dim; (void)heads; (void)n_tokens;
return 0;
}
extern "C" int ds4_gpu_attention_output_q4_K_batch_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp, ds4_gpu_tensor *low_tmp,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset, uint32_t out_b_type,
uint64_t group_dim, uint64_t rank, uint32_t n_groups,
uint64_t out_dim, const ds4_gpu_tensor *heads, uint32_t n_tokens) {
(void)out; (void)low; (void)group_tmp; (void)low_tmp;
(void)model_map; (void)model_size; (void)out_a_offset;
(void)out_b_offset; (void)out_b_type; (void)group_dim; (void)rank;
(void)n_groups; (void)out_dim; (void)heads; (void)n_tokens;
return 0;
}
extern "C" int ds4_gpu_attention_output_low_q4_K_slice_tensor(
ds4_gpu_tensor *low, const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t group_dim, uint64_t rank,
uint32_t group0, uint32_t group_cnt,
const ds4_gpu_tensor *heads) {
(void)low; (void)model_map; (void)model_size; (void)out_a_offset;
(void)group_dim; (void)rank; (void)group0; (void)group_cnt;
(void)heads;
return 0;
}
extern "C" int ds4_gpu_hc_expand_split_half_tensor(
ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out_h,
const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split,
uint32_t n_embd, uint32_t n_hc) {
(void)out_hc; (void)block_out_h; (void)residual_hc; (void)split;
(void)n_embd; (void)n_hc;
return 0;
}
extern "C" int ds4_gpu_hc_expand_add_split_half_add_tensor(
ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add_h,
const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split,
uint32_t n_embd, uint32_t n_hc) {
(void)out_hc; (void)block_out; (void)block_add_h;
(void)residual_hc; (void)split; (void)n_embd; (void)n_hc;
return 0;
}
extern "C" void ds4_gpu_tp_suspend_expert_sharding(int suspend) {
(void)suspend;
}
extern "C" void ds4_gpu_tp_keepalive_pause(int paused) {
(void)paused;
}
extern "C" int ds4_gpu_tp_failed(void) { return g_cuda_tp.failed; }
extern "C" void ds4_gpu_model_residency_skip(int skip) {
(void)skip;
}
extern "C" uint64_t ds4_gpu_tp_big_gate_kick(
uint32_t layer, uint32_t rows, const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t, uint64_t bytes) {
(void)layer; (void)rows; (void)out_t; (void)in_t; (void)bytes;
return 0;
}
extern "C" int ds4_gpu_tp_big_gate_wait(uint64_t seq) {
(void)seq;
return 0;
}
#pragma GCC diagnostic pop
#define DS4_GLM53_VISION_STREAM cuda_decode_stream()
#include "ds4_glm53_vision_gpu.cuh"
#include "ds4_deepseek4_vision_gpu.cuh"
#include "ds4_deepseek41_cuda.cuh"
#include "ds4_qwen4_cuda.cuh"
Reference in new issue View Git Blame Copy Permalink