diff --git a/Makefile b/Makefile index fdfb4986f..708d4a163 100644 --- a/Makefile +++ b/Makefile @@ -215,7 +215,9 @@ ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kv ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_SRCS) $(CC) $(OBJCFLAGS) -c -o $@ ds4_metal.m -ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc +ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc ds4_cuda_gqa.inc \ + ds4_cuda_glm_kv.inc ds4_cuda_glm_indexer.inc ds4_cuda_glm_attn.inc \ + ds4_cuda_glm_moe.inc ds4_cuda_glm_stubs.inc $(NVCC) $(NVCCFLAGS) -c -o $@ ds4_cuda.cu ds4_rocm.o: ds4_rocm.cu ds4_gpu.h ds4_iq2_tables_cuda.inc $(ROCM_SRCS) diff --git a/ds4.c b/ds4.c index edf566489..23bfab68f 100644 --- a/ds4.c +++ b/ds4.c @@ -111,14 +111,14 @@ static bool ds4_backend_supports_streaming_auto_cache(ds4_backend backend) { */ enum { - DS4_MAX_LAYER = 79, + DS4_MAX_LAYER = 81, DS4_MAX_EMBD = 7168, DS4_MAX_VOCAB = 154880, DS4_MAX_HEAD = 128, - DS4_MAX_HEAD_KV = 1, + DS4_MAX_HEAD_KV = 8, DS4_MAX_HEAD_DIM = 576, DS4_MAX_VALUE_DIM = 512, - DS4_MAX_ROT = 64, + DS4_MAX_ROT = 128, DS4_MAX_OUT_GROUP = 16, DS4_MAX_LORA_Q = 2048, DS4_MAX_LORA_O = 1024, @@ -138,12 +138,14 @@ enum { typedef enum { DS4_MODEL_FAMILY_DEEPSEEK4 = 0, DS4_MODEL_FAMILY_GLM_DSA = 1, + DS4_MODEL_FAMILY_HY_V3 = 2, } ds4_model_family; typedef enum { DS4_VARIANT_FLASH = 0, DS4_VARIANT_PRO = 1, DS4_VARIANT_GLM52 = 2, + DS4_VARIANT_HY3 = 3, } ds4_variant; typedef struct { @@ -310,6 +312,55 @@ static const ds4_shape DS4_SHAPE_GLM52 = { .rope_orig_ctx = 1048576, }; +/* Tencent Hy3 (hy_v3): plain GQA attention (no MLA, no indexer), 80 + * transformer layers plus 1 nextn MTP layer bound like GLM's. n_head_dim + * doubles as the GQA head_dim for q/k/v; n_rot = head_dim (full-head + * rotate-half RoPE, no yarn). expert_weight_scale carries the config's + * router_scaling_factor; sigmoid routing with expert bias mirrors GLM. */ +static const ds4_shape DS4_SHAPE_HY3 = { + .name = "Hy3", + .family = DS4_MODEL_FAMILY_HY_V3, + .variant = DS4_VARIANT_HY3, + .n_layer = 81, + .n_embd = 4096, + .n_vocab = 120832, + .n_head = 64, + .n_head_kv = 8, + .n_head_dim = 128, + .n_value_dim = 128, + .n_rot = 128, + .n_out_group = 0, + .n_lora_q = 0, + .n_lora_o = 0, + .n_expert = 192, + .n_expert_used = 8, + .n_expert_shared = 1, + .n_ff_exp = 1536, + .n_ff_dense = 13312, + .n_hash_layer = 0, + .n_swa = 0, + .n_indexer_head = 0, + .n_indexer_head_dim = 0, + .n_indexer_top_k = 0, + .n_hc = 0, + .n_hc_sinkhorn_iter = 0, + .n_nextn_predict = 1, + .n_leading_dense = 1, + .n_kv_lora = 0, + .n_key_mla = 0, + .n_value_mla = 0, + .rms_eps = 1.0e-5f, + .hc_eps = 0.0f, + .expert_weight_scale = 2.826f, + .swiglu_clamp_exp = 0.0f, + .rope_freq_base = 11158840.0f, + .rope_scale_factor = 1.0f, + .rope_yarn_beta_fast = 0.0f, + .rope_yarn_beta_slow = 0.0f, + .compress_rope_freq_base = 0.0f, + .rope_orig_ctx = 262144, +}; + static ds4_shape g_ds4_shape = { .name = "DeepSeek V4 Flash", .family = DS4_MODEL_FAMILY_DEEPSEEK4, @@ -3406,6 +3457,12 @@ typedef struct { ds4_tensor *hc_attn_scale; ds4_tensor *hc_attn_base; ds4_tensor *attn_norm; + /* GQA families (hy_v3): direct q/k/v projections + per-head qk norms. */ + ds4_tensor *attn_q; + ds4_tensor *attn_k; + ds4_tensor *attn_v; + ds4_tensor *attn_q_norm; + ds4_tensor *attn_k_norm; ds4_tensor *attn_q_a; ds4_tensor *attn_q_a_norm; ds4_tensor *attn_q_b; @@ -3942,7 +3999,7 @@ static void tensor_expect_routed_expert( } static bool weights_have_output_head(const ds4_weights *w) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { return w && w->output_norm && w->output; } return w && @@ -3954,7 +4011,7 @@ static bool weights_have_output_head(const ds4_weights *w) { } static bool weights_have_partial_output_head(const ds4_weights *w) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { return w && (w->output_norm || w->output); } return w && @@ -4174,6 +4231,77 @@ static void weights_validate_glm_dsa_layout( } } +static void weights_validate_hy_v3_layout( + const ds4_weights *w, + uint32_t layer_start, + uint32_t layer_end, + bool require_token_embd, + bool require_output) { + const uint64_t q_dim = (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM; + const uint64_t kv_dim = (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM; + + if (!w) ds4_die("internal error: missing weights while validating hy_v3 layout"); + if (layer_start >= DS4_N_LAYER) ds4_die("invalid first layer in hy_v3 weight layout validation"); + if (layer_end == UINT32_MAX) layer_end = DS4_N_LAYER - 1u; + if (layer_end >= DS4_N_LAYER || layer_end < layer_start) { + ds4_die("invalid layer range in hy_v3 weight layout validation"); + } + + if (require_token_embd && !w->token_embd) ds4_die("required token embedding tensor is missing"); + if (w->token_embd) { + tensor_expect_layout(w->token_embd, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); + } + + const bool have_output = weights_have_output_head(w); + if (require_output && !have_output) ds4_die("required output head tensors are missing"); + if (weights_have_partial_output_head(w) && !have_output) ds4_die("partial output head in GGUF"); + if (have_output) { + tensor_expect_layout(w->output_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(w->output, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_VOCAB, 0); + } + + for (uint32_t il = layer_start; il <= layer_end; il++) { + const ds4_layer_weights *l = &w->layer[il]; + if (!l->attn_norm) continue; /* unbound MTP tail layer without nextn block */ + + tensor_expect_layout(l->attn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->attn_q, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, q_dim, 0); + tensor_expect_layout(l->attn_k, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, kv_dim, 0); + tensor_expect_layout(l->attn_v, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, kv_dim, 0); + tensor_expect_layout(l->attn_q_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); + tensor_expect_layout(l->attn_k_norm, DS4_TENSOR_F32, 1, DS4_N_HEAD_DIM, 0, 0); + tensor_expect_layout(l->attn_output, DS4_TENSOR_Q8_0, 2, q_dim, DS4_N_EMBD, 0); + tensor_expect_layout(l->ffn_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + + if (il < DS4_N_LEADING_DENSE) { + tensor_expect_layout(l->ffn_gate, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); + tensor_expect_layout(l->ffn_up, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_DENSE, 0); + tensor_expect_layout(l->ffn_down, DS4_TENSOR_Q8_0, 2, DS4_N_FF_DENSE, DS4_N_EMBD, 0); + } else { + tensor_expect_layout(l->ffn_gate_inp, DS4_TENSOR_F32, 2, DS4_N_EMBD, DS4_N_EXPERT, 0); + tensor_expect_layout(l->ffn_exp_probs_b, DS4_TENSOR_F32, 1, DS4_N_EXPERT, 0, 0); + tensor_expect_routed_expert(l->ffn_gate_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_up_exps, 3, DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EXPERT); + tensor_expect_routed_expert(l->ffn_down_exps, 3, DS4_N_FF_EXP, DS4_N_EMBD, DS4_N_EXPERT); + if (l->ffn_gate_exps->type != l->ffn_up_exps->type) { + fprintf(stderr, "ds4: hy_v3 routed gate/up experts use different quant types in layer %u\n", il); + exit(1); + } + tensor_expect_layout(l->ffn_gate_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_layout(l->ffn_up_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_EMBD, DS4_N_FF_EXP, 0); + tensor_expect_layout(l->ffn_down_shexp, DS4_TENSOR_Q8_0, 2, DS4_N_FF_EXP, DS4_N_EMBD, 0); + } + + if (DS4_N_NEXTN_PREDICT != 0 && + il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { + tensor_expect_layout(l->nextn_eh_proj, DS4_TENSOR_Q8_0, 2, 2u * DS4_N_EMBD, DS4_N_EMBD, 0); + tensor_expect_layout(l->nextn_enorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->nextn_hnorm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + tensor_expect_layout(l->nextn_shared_head_norm, DS4_TENSOR_F32, 1, DS4_N_EMBD, 0, 0); + } + } +} + static void weights_validate_layout( const ds4_weights *w, uint32_t layer_start, @@ -4188,6 +4316,14 @@ static void weights_validate_layout( require_output); return; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + weights_validate_hy_v3_layout(w, + layer_start, + layer_end, + require_token_embd, + require_output); + return; + } const uint64_t hc_dim = (uint64_t)DS4_N_EMBD * DS4_N_HC; const uint64_t hc_mix_dim = 2u * DS4_N_HC + (uint64_t)DS4_N_HC * DS4_N_HC; @@ -4712,18 +4848,73 @@ static void config_validate_glm_dsa_model(const ds4_model *m) { config_expect_bool("expert_weights_norm", expert_weight_norm, true); } +static void config_validate_hy_v3_model(const ds4_model *m) { + g_ds4_shape = DS4_SHAPE_HY3; + memset(g_ds4_compress_ratios, 0, sizeof(g_ds4_compress_ratios)); + + const uint32_t n_layer = required_u32(m, "hy-v3.block_count"); + const uint64_t n_ctx = required_u64_compat(m, "hy-v3.context_length"); + const uint32_t n_embd = required_u32(m, "hy-v3.embedding_length"); + /* llama.cpp-converted hy-v3 ggufs carry no .vocab_size key (the + * vocab is implied by the tokenizer arrays); fall back to the model's. */ + uint32_t n_vocab = 120832; + (void)model_get_u32(m, "hy-v3.vocab_size", &n_vocab); + const uint32_t n_ff_dense = required_u32(m, "hy-v3.feed_forward_length"); + const uint32_t n_head = required_u32(m, "hy-v3.attention.head_count"); + const uint32_t n_head_kv = required_u32(m, "hy-v3.attention.head_count_kv"); + const uint32_t n_head_dim = required_u32(m, "hy-v3.attention.key_length"); + const uint32_t n_rot = required_u32(m, "hy-v3.rope.dimension_count"); + const uint32_t n_expert = required_u32(m, "hy-v3.expert_count"); + const uint32_t n_expert_used = required_u32(m, "hy-v3.expert_used_count"); + const uint32_t n_ff_exp = required_u32(m, "hy-v3.expert_feed_forward_length"); + const uint32_t n_expert_shared = required_u32(m, "hy-v3.expert_shared_count"); + const uint32_t n_leading_dense = required_u32(m, "hy-v3.leading_dense_block_count"); + const uint32_t n_nextn = required_u32(m, "hy-v3.nextn_predict_layers"); + + config_expect_u32("block_count", n_layer, DS4_N_LAYER); + config_expect_u64("context_length", n_ctx, DS4_ROPE_ORIG_CTX); + config_expect_u32("embedding_length", n_embd, DS4_N_EMBD); + config_expect_u32("vocab_size", n_vocab, DS4_N_VOCAB); + config_expect_u32("feed_forward_length", n_ff_dense, DS4_N_FF_DENSE); + config_expect_u32("attention.head_count", n_head, DS4_N_HEAD); + config_expect_u32("attention.head_count_kv", n_head_kv, DS4_N_HEAD_KV); + config_expect_u32("attention.key_length", n_head_dim, DS4_N_HEAD_DIM); + config_expect_u32("rope.dimension_count", n_rot, DS4_N_ROT); + config_expect_u32("expert_count", n_expert, DS4_N_EXPERT); + config_expect_u32("expert_used_count", n_expert_used, DS4_N_EXPERT_USED); + config_expect_u32("expert_feed_forward_length", n_ff_exp, DS4_N_FF_EXP); + config_expect_u32("expert_shared_count", n_expert_shared, DS4_N_EXPERT_SHARED); + config_expect_u32("leading_dense_block_count", n_leading_dense, DS4_N_LEADING_DENSE); + config_expect_u32("nextn_predict_layers", n_nextn, DS4_N_NEXTN_PREDICT); + + const float rope_freq_base = required_f32(m, "hy-v3.rope.freq_base"); + config_expect_f32("rope.freq_base", rope_freq_base, DS4_ROPE_FREQ_BASE); + const float rms_eps = required_f32(m, "hy-v3.attention.layer_norm_rms_epsilon"); + config_expect_f32("attention.layer_norm_rms_epsilon", rms_eps, DS4_RMS_EPS); + const float expert_weight_scale = required_f32(m, "hy-v3.expert_weights_scale"); + config_expect_f32("expert_weights_scale", expert_weight_scale, DS4_EXPERT_WEIGHT_SCALE); + /* Sigmoid routing with post-top-k weight normalization (route_norm). */ + const bool expert_weight_norm = required_bool(m, "hy-v3.expert_weights_norm"); + config_expect_bool("expert_weights_norm", expert_weight_norm, true); +} + static void config_validate_model(const ds4_model *m) { ds4_str arch = {0}; - if (model_get_string(m, "general.architecture", &arch) && - ds4_streq(arch, "glm-dsa")) { - config_validate_glm_dsa_model(m); - return; + if (model_get_string(m, "general.architecture", &arch)) { + if (ds4_streq(arch, "glm-dsa")) { + config_validate_glm_dsa_model(m); + return; + } + if (ds4_streq(arch, "hy-v3")) { + config_validate_hy_v3_model(m); + return; + } } config_validate_deepseek4_model(m); } static void weights_bind_output(ds4_weights *w, const ds4_model *m, bool required) { - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { if (required) { w->output_norm = required_tensor(m, "output_norm.weight"); w->output = required_tensor(m, "output.weight"); @@ -4763,7 +4954,63 @@ static void weights_bind_glm_dsa_layer(ds4_layer_weights *l, const ds4_model *m, l->ffn_down = required_tensorf(m, "blk.%u.ffn_down.weight", il); } else { l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); - l->ffn_exp_probs_b = required_tensorf(m, "blk.%u.exp_probs_b.bias", il); + l->ffn_exp_probs_b = required_tensorf(m, "blk.%u.exp_probs_b", il); + l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); + l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); + l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); + l->ffn_gate_shexp = required_tensorf(m, "blk.%u.ffn_gate_shexp.weight", il); + l->ffn_up_shexp = required_tensorf(m, "blk.%u.ffn_up_shexp.weight", il); + l->ffn_down_shexp = required_tensorf(m, "blk.%u.ffn_down_shexp.weight", il); + } + + if (DS4_N_NEXTN_PREDICT != 0 && + il + DS4_N_NEXTN_PREDICT >= DS4_N_LAYER) { + l->nextn_eh_proj = required_tensorf(m, "blk.%u.nextn.eh_proj.weight", il); + l->nextn_enorm = required_tensorf(m, "blk.%u.nextn.enorm.weight", il); + l->nextn_hnorm = required_tensorf(m, "blk.%u.nextn.hnorm.weight", il); + l->nextn_shared_head_norm = + required_tensorf(m, "blk.%u.nextn.shared_head_norm.weight", il); + } +} + +/* hy_v3 GGUF tensor-name contract (defined here; our converter emits it, + * llama.cpp-style names): + * token_embd.weight, output_norm.weight, output.weight + * blk.N.attn_norm.weight + * blk.N.attn_q.weight [4096 -> 64*128] + * blk.N.attn_k.weight [4096 -> 8*128] + * blk.N.attn_v.weight [4096 -> 8*128] + * blk.N.attn_q_norm.weight [128] per-head RMSNorm weight + * blk.N.attn_k_norm.weight [128] + * blk.N.attn_output.weight [64*128 -> 4096] + * blk.N.ffn_norm.weight + * dense layer (il < leading_dense=1): + * blk.N.ffn_gate.weight / ffn_up.weight / ffn_down.weight [13312] + * MoE layers: + * blk.N.ffn_gate_inp.weight router [4096 -> 192] + * blk.N.exp_probs_b expert bias [192] (no suffix) + * blk.N.ffn_{gate,up,down}_exps.weight [192 experts, ff 1536] + * blk.N.ffn_{gate,up,down}_shexp.weight shared expert + * MTP layer (blk.80): the regular names above plus + * blk.N.nextn.eh_proj.weight / nextn.enorm.weight / nextn.hnorm.weight / + * nextn.shared_head_norm.weight (same glue contract as GLM). */ +static void weights_bind_hy_v3_layer(ds4_layer_weights *l, const ds4_model *m, uint32_t il) { + l->attn_norm = required_tensorf(m, "blk.%u.attn_norm.weight", il); + l->attn_q = required_tensorf(m, "blk.%u.attn_q.weight", il); + l->attn_k = required_tensorf(m, "blk.%u.attn_k.weight", il); + l->attn_v = required_tensorf(m, "blk.%u.attn_v.weight", il); + l->attn_q_norm = required_tensorf(m, "blk.%u.attn_q_norm.weight", il); + l->attn_k_norm = required_tensorf(m, "blk.%u.attn_k_norm.weight", il); + l->attn_output = required_tensorf(m, "blk.%u.attn_output.weight", il); + l->ffn_norm = required_tensorf(m, "blk.%u.ffn_norm.weight", il); + + if (il < DS4_N_LEADING_DENSE) { + l->ffn_gate = required_tensorf(m, "blk.%u.ffn_gate.weight", il); + l->ffn_up = required_tensorf(m, "blk.%u.ffn_up.weight", il); + l->ffn_down = required_tensorf(m, "blk.%u.ffn_down.weight", il); + } else { + l->ffn_gate_inp = required_tensorf(m, "blk.%u.ffn_gate_inp.weight", il); + l->ffn_exp_probs_b = required_tensorf(m, "blk.%u.exp_probs_b", il); l->ffn_gate_exps = required_tensorf(m, "blk.%u.ffn_gate_exps.weight", il); l->ffn_up_exps = required_tensorf(m, "blk.%u.ffn_up_exps.weight", il); l->ffn_down_exps = required_tensorf(m, "blk.%u.ffn_down_exps.weight", il); @@ -4787,6 +5034,10 @@ static void weights_bind_layer(ds4_layer_weights *l, const ds4_model *m, uint32_ weights_bind_glm_dsa_layer(l, m, il); return; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + weights_bind_hy_v3_layer(l, m, il); + return; + } const uint32_t compress_ratio = ds4_layer_compress_ratio(il); @@ -4846,7 +5097,7 @@ static void weights_bind( memset(w, 0, sizeof(*w)); uint32_t executable_layers = DS4_N_LAYER; - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4 && DS4_N_LAYER > DS4_N_NEXTN_PREDICT) { executable_layers = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; } @@ -4874,6 +5125,23 @@ static void weights_bind( weights_bind_layer(&w->layer[il], m, il); } + /* Bind the GLM MTP draft layer(s) that sit beyond the executable range so + * weights.layer[N_LAYER-1] is populated for --mtp (it holds the nextn glue + * plus a full GLM layer). These never run in the decode loop; they only + * feed glm_graph_eval_mtp_draft. Absent in GLM ggufs without an MTP block, + * so bind only when the block is actually present. */ + if (!load_slice && + DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4 && + executable_layers < DS4_N_LAYER) { + for (uint32_t il = executable_layers; il < DS4_N_LAYER; il++) { + char probe[64]; + snprintf(probe, sizeof(probe), "blk.%u.nextn.eh_proj.weight", il); + if (model_find_tensor(m, probe) != NULL) { + weights_bind_layer(&w->layer[il], m, il); + } + } + } + weights_validate_layout(w, start, end, require_token_embd, require_output); } @@ -25319,6 +25587,35 @@ static void vocab_load(ds4_vocab *vocab, const ds4_model *model) { return; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + /* Hy3 ggufs carry the special ids as metadata; the token strings + * differ from the DeepSeek literals, so resolve by id. */ + if (!model_get_token_id(model, "tokenizer.ggml.bos_token_id", &vocab->bos_id) || + !model_get_token_id(model, "tokenizer.ggml.eos_token_id", &vocab->eos_id)) { + fprintf(stderr, "ds4: hy-v3 gguf is missing bos/eos token ids\n"); + exit(1); + } + vocab->system_id = -1; + vocab->user_id = vocab_lookup_optional(vocab, "<｜hy_User:opensource｜>"); + vocab->assistant_id = vocab_lookup_optional(vocab, "<｜hy_Assistant:opensource｜>"); + vocab->observation_id = -1; + /* ponytail: reuse the (GLM-only) sop_id slot to hold Hy3's + * reasoning-mode marker; no dedicated field needed for one token. */ + vocab->sop_id = vocab_lookup_optional(vocab, "<｜reasoning_mode:opensource｜>"); + vocab->think_start_id = vocab_lookup_optional(vocab, ""); + vocab->think_end_id = vocab_lookup_optional(vocab, ""); + vocab->tool_call_start_id = -1; + vocab->tool_call_end_id = -1; + vocab->tool_response_start_id = -1; + vocab->tool_response_end_id = -1; + vocab->arg_key_start_id = -1; + vocab->arg_key_end_id = -1; + vocab->arg_value_start_id = -1; + vocab->arg_value_end_id = -1; + vocab->dsml_id = -1; + return; + } + vocab->bos_id = vocab_lookup(vocab, "<｜begin▁of▁sentence｜>"); vocab->eos_id = vocab_lookup(vocab, "<｜end▁of▁sentence｜>"); vocab->system_id = -1; @@ -25386,7 +25683,8 @@ static void encode_chat_prompt( token_vec *out) { const bool need_think_start = ds4_think_mode_enabled(think_mode) || - DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3; if (vocab->bos_id < 0 || vocab->user_id < 0 || vocab->assistant_id < 0 || @@ -25408,7 +25706,8 @@ static void encode_chat_prompt( token_vec_push(out, vocab->assistant_id); if (ds4_think_mode_enabled(think_mode)) { token_vec_push(out, vocab->think_start_id); - } else if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { + } else if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { token_vec_push(out, vocab->think_start_id); token_vec_push(out, vocab->think_end_id); } else { @@ -25570,6 +25869,23 @@ void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role return; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + if (!strcmp(role, "system") || !strcmp(role, "developer")) { + bpe_tokenize_text(vocab, content, tokens); + } else if (!strcmp(role, "assistant")) { + token_vec_push(tokens, vocab->assistant_id); + /* ponytail: no_think framing; v1 does not reconstruct prior + * reasoning content into the transcript. */ + token_vec_push(tokens, vocab->think_start_id); + token_vec_push(tokens, vocab->think_end_id); + bpe_tokenize_text(vocab, content, tokens); + } else { + token_vec_push(tokens, vocab->user_id); + bpe_tokenize_text(vocab, content, tokens); + } + return; + } + if (!strcmp(role, "system") || !strcmp(role, "developer")) { bpe_tokenize_text(vocab, content, tokens); } else if (!strcmp(role, "assistant")) { @@ -25591,7 +25907,8 @@ void ds4_chat_append_message(ds4_engine *e, ds4_tokens *tokens, const char *role void ds4_chat_append_assistant_prefix(ds4_engine *e, ds4_tokens *tokens, ds4_think_mode think_mode) { token_vec_push(tokens, e->vocab.assistant_id); - if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA && + if ((DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA || + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) && !ds4_think_mode_enabled(think_mode)) { token_vec_push(tokens, e->vocab.think_start_id); token_vec_push(tokens, e->vocab.think_end_id); @@ -30697,9 +31014,13 @@ static bool glm_graph_forward_indexed_tokens( !force_scalar_attn && glm_graph_indexed_prefill_batch_qk_low(); const bool use_batch_attn_kernel = !force_scalar_attn && glm_graph_indexed_prefill_batch_attn_kernel(); + /* The split value-projection path assumes the f16 compact cache layout + * (Metal); on CUDA the compact cache is f32 and the split path fails + * the batch attention-lora encode at layer 0. */ const bool use_split_value_proj = use_batch_attn_kernel && - g->batch_attn_lora; + g->batch_attn_lora && + glm_graph_compact_cache_is_f16() != 0; const bool use_batch_q_rank_proj = true; const bool use_batch_q_proj = true; const bool use_batch_indexer_k_proj = true; @@ -33313,6 +33634,380 @@ static int generate_glm_metal_argmax( return 0; } +/* ---- Hy3 (hy_v3) CUDA inference ------------------------------------- + * Straight-line per-token forward, no graph machinery: Hy3 is plain GQA + * (no MLA latents, no indexer, no compressed KV), so the whole decode is + * existing Q8_0 matmul / rmsnorm primitives, the ds4_cuda_gqa.inc kernels, + * and the GLM streaming MoE machinery called with Hy3 dimensions. + * ponytail: token-major prefill (one forward per prompt token); batch + * prefill is a later workstream, same as GLM's was. */ + +typedef struct hy3_gpu_state { + uint32_t ctx; + uint32_t n_exec_layer; + ds4_gpu_tensor *tok; /* i32 [1] */ + ds4_gpu_tensor *cur; /* f32 [embd] residual stream */ + ds4_gpu_tensor *attn_norm; /* f32 [embd] */ + ds4_gpu_tensor *q; /* f32 [64*128] */ + ds4_gpu_tensor *k; /* f32 [8*128] */ + ds4_gpu_tensor *v; /* f32 [8*128] */ + ds4_gpu_tensor *heads; /* f32 [64*128] */ + ds4_gpu_tensor *attn_out; /* f32 [embd] */ + ds4_gpu_tensor *after_attn; /* f32 [embd] */ + ds4_gpu_tensor *ffn_norm; /* f32 [embd] */ + ds4_gpu_tensor *dense_gate; /* f32 [ff_dense] */ + ds4_gpu_tensor *dense_up; /* f32 [ff_dense] */ + ds4_gpu_tensor *dense_mid; /* f32 [ff_dense] */ + ds4_gpu_tensor *shared_mid; /* f32 [ff_exp] */ + ds4_gpu_tensor *shared_out; /* f32 [embd] */ + ds4_gpu_tensor *router_logits; /* f32 [n_expert] */ + ds4_gpu_tensor *router_probs; /* f32 [n_expert] */ + ds4_gpu_tensor *router_selected; /* i32 [n_expert_used] */ + ds4_gpu_tensor *router_weights; /* f32 [n_expert_used] */ + ds4_gpu_tensor *routed_gate; /* f32 [used*ff_exp] */ + ds4_gpu_tensor *routed_up; /* f32 [used*ff_exp] */ + ds4_gpu_tensor *routed_mid; /* f32 [used*ff_exp] */ + ds4_gpu_tensor *routed_down; /* f32 [used*embd] */ + ds4_gpu_tensor *routed_out; /* f32 [embd] */ + ds4_gpu_tensor *out_norm; /* f32 [embd] */ + ds4_gpu_tensor *logits; /* f32 [vocab] */ + /* per executable layer */ + ds4_gpu_tensor **kcache; /* f32 [n_head_kv][ctx][head_dim] */ + ds4_gpu_tensor **vcache; + ds4_gpu_tensor **qn; /* f32 [head_dim] qk-norm weights */ + ds4_gpu_tensor **kn; + bool ssd_streaming; +} hy3_gpu_state; + +static void hy3_state_free(hy3_gpu_state *s) { + if (!s) return; + ds4_gpu_tensor_free(s->tok); ds4_gpu_tensor_free(s->cur); + ds4_gpu_tensor_free(s->attn_norm); + ds4_gpu_tensor_free(s->q); ds4_gpu_tensor_free(s->k); + ds4_gpu_tensor_free(s->v); ds4_gpu_tensor_free(s->heads); + ds4_gpu_tensor_free(s->attn_out); ds4_gpu_tensor_free(s->after_attn); + ds4_gpu_tensor_free(s->ffn_norm); + ds4_gpu_tensor_free(s->dense_gate); ds4_gpu_tensor_free(s->dense_up); + ds4_gpu_tensor_free(s->dense_mid); + ds4_gpu_tensor_free(s->shared_mid); ds4_gpu_tensor_free(s->shared_out); + ds4_gpu_tensor_free(s->router_logits); ds4_gpu_tensor_free(s->router_probs); + ds4_gpu_tensor_free(s->router_selected); ds4_gpu_tensor_free(s->router_weights); + ds4_gpu_tensor_free(s->routed_gate); ds4_gpu_tensor_free(s->routed_up); + ds4_gpu_tensor_free(s->routed_mid); ds4_gpu_tensor_free(s->routed_down); + ds4_gpu_tensor_free(s->routed_out); + ds4_gpu_tensor_free(s->out_norm); ds4_gpu_tensor_free(s->logits); + for (uint32_t il = 0; s->kcache && il < s->n_exec_layer; il++) { + ds4_gpu_tensor_free(s->kcache[il]); + ds4_gpu_tensor_free(s->vcache[il]); + ds4_gpu_tensor_free(s->qn[il]); + ds4_gpu_tensor_free(s->kn[il]); + } + free(s->kcache); free(s->vcache); free(s->qn); free(s->kn); + memset(s, 0, sizeof(*s)); +} + +static bool hy3_state_alloc( + hy3_gpu_state *s, + const ds4_model *model, + const ds4_weights *weights, + uint32_t ctx, + bool ssd_streaming) { + memset(s, 0, sizeof(*s)); + s->ctx = ctx; + s->n_exec_layer = DS4_N_LAYER - DS4_N_NEXTN_PREDICT; + s->ssd_streaming = ssd_streaming; + const uint64_t embd = DS4_N_EMBD; + const uint64_t hd = DS4_N_HEAD_DIM; + s->tok = ds4_gpu_tensor_alloc(sizeof(int32_t)); + s->cur = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->attn_norm = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->q = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD * hd * sizeof(float)); + s->k = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_KV * hd * sizeof(float)); + s->v = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD_KV * hd * sizeof(float)); + s->heads = ds4_gpu_tensor_alloc((uint64_t)DS4_N_HEAD * hd * sizeof(float)); + s->attn_out = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->after_attn = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->ffn_norm = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->dense_gate = ds4_gpu_tensor_alloc((uint64_t)DS4_N_FF_DENSE * sizeof(float)); + s->dense_up = ds4_gpu_tensor_alloc((uint64_t)DS4_N_FF_DENSE * sizeof(float)); + s->dense_mid = ds4_gpu_tensor_alloc((uint64_t)DS4_N_FF_DENSE * sizeof(float)); + s->shared_mid = ds4_gpu_tensor_alloc((uint64_t)DS4_N_FF_EXP * sizeof(float)); + s->shared_out = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->router_logits = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT * sizeof(float)); + s->router_probs = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT * sizeof(float)); + s->router_selected = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * sizeof(int32_t)); + s->router_weights = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * sizeof(float)); + s->routed_gate = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); + s->routed_up = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); + s->routed_mid = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * DS4_N_FF_EXP * sizeof(float)); + s->routed_down = ds4_gpu_tensor_alloc((uint64_t)DS4_N_EXPERT_USED * embd * sizeof(float)); + s->routed_out = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->out_norm = ds4_gpu_tensor_alloc(embd * sizeof(float)); + s->logits = ds4_gpu_tensor_alloc((uint64_t)DS4_N_VOCAB * sizeof(float)); + if (!s->tok || !s->cur || !s->attn_norm || !s->q || !s->k || !s->v || + !s->heads || !s->attn_out || !s->after_attn || !s->ffn_norm || + !s->dense_gate || !s->dense_up || !s->dense_mid || + !s->shared_mid || !s->shared_out || + !s->router_logits || !s->router_probs || !s->router_selected || + !s->router_weights || !s->routed_gate || !s->routed_up || + !s->routed_mid || !s->routed_down || !s->routed_out || + !s->out_norm || !s->logits) { + hy3_state_free(s); + return false; + } + s->kcache = xcalloc(s->n_exec_layer, sizeof(s->kcache[0])); + s->vcache = xcalloc(s->n_exec_layer, sizeof(s->vcache[0])); + s->qn = xcalloc(s->n_exec_layer, sizeof(s->qn[0])); + s->kn = xcalloc(s->n_exec_layer, sizeof(s->kn[0])); + const uint64_t cache_bytes = + (uint64_t)DS4_N_HEAD_KV * ctx * hd * sizeof(float); + for (uint32_t il = 0; il < s->n_exec_layer; il++) { + const ds4_layer_weights *l = &weights->layer[il]; + s->kcache[il] = ds4_gpu_tensor_alloc(cache_bytes); + s->vcache[il] = ds4_gpu_tensor_alloc(cache_bytes); + s->qn[il] = ds4_gpu_tensor_alloc(hd * sizeof(float)); + s->kn[il] = ds4_gpu_tensor_alloc(hd * sizeof(float)); + if (!s->kcache[il] || !s->vcache[il] || !s->qn[il] || !s->kn[il] || + ds4_gpu_tensor_write(s->qn[il], 0, + tensor_data(model, l->attn_q_norm), + hd * sizeof(float)) == 0 || + ds4_gpu_tensor_write(s->kn[il], 0, + tensor_data(model, l->attn_k_norm), + hd * sizeof(float)) == 0) { + hy3_state_free(s); + return false; + } + } + return true; +} + +/* One full forward for one token at absolute position pos. Reads the + * embedding, runs every executable layer, and (optionally) produces host + * logits. The residual stream lives in s->cur throughout. */ +static bool hy3_forward_token( + hy3_gpu_state *s, + const ds4_model *model, + const ds4_weights *weights, + int token, + uint32_t pos, + float *logits_out) { + const int32_t tok32 = (int32_t)token; + if (pos >= s->ctx) return false; + if (ds4_gpu_tensor_write(s->tok, 0, &tok32, sizeof(tok32)) == 0) return false; + bool ok = ds4_gpu_embed_tokens_q8_0_tensor(s->cur, + s->tok, + model->map, + model->size, + weights->token_embd->abs_offset, + DS4_N_VOCAB, + 1, + DS4_N_EMBD) != 0; + for (uint32_t il = 0; ok && il < s->n_exec_layer; il++) { + const ds4_layer_weights *l = &weights->layer[il]; + /* attention */ + ok = ds4_gpu_rms_norm_weight_rows_tensor(s->attn_norm, s->cur, + model->map, model->size, l->attn_norm->abs_offset, + DS4_N_EMBD, 1, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->q, model->map, model->size, + l->attn_q->abs_offset, DS4_N_EMBD, + (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM, s->attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->k, model->map, model->size, + l->attn_k->abs_offset, DS4_N_EMBD, + (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM, s->attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->v, model->map, model->size, + l->attn_v->abs_offset, DS4_N_EMBD, + (uint64_t)DS4_N_HEAD_KV * DS4_N_HEAD_DIM, s->attn_norm, 1) != 0; + if (ok) ok = ds4_gpu_gqa_head_rms_norm_weight(s->q, s->qn[il], + DS4_N_HEAD, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_gqa_head_rms_norm_weight(s->k, s->kn[il], + DS4_N_HEAD_KV, DS4_N_HEAD_DIM, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_gqa_rope(s->q, 1, DS4_N_HEAD, DS4_N_HEAD_DIM, + pos, DS4_ROPE_FREQ_BASE) != 0; + if (ok) ok = ds4_gpu_gqa_rope(s->k, 1, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, + pos, DS4_ROPE_FREQ_BASE) != 0; + if (ok) ok = ds4_gpu_gqa_kv_cache_append(s->kcache[il], s->k, 1, + DS4_N_HEAD_KV, DS4_N_HEAD_DIM, s->ctx, pos) != 0; + if (ok) ok = ds4_gpu_gqa_kv_cache_append(s->vcache[il], s->v, 1, + DS4_N_HEAD_KV, DS4_N_HEAD_DIM, s->ctx, pos) != 0; + if (ok) ok = ds4_gpu_gqa_attention(s->heads, s->q, + s->kcache[il], s->vcache[il], 1, + DS4_N_HEAD, DS4_N_HEAD_KV, DS4_N_HEAD_DIM, s->ctx, pos) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->attn_out, model->map, + model->size, l->attn_output->abs_offset, + (uint64_t)DS4_N_HEAD * DS4_N_HEAD_DIM, DS4_N_EMBD, + s->heads, 1) != 0; + if (ok) ok = ds4_gpu_add_tensor(s->after_attn, s->cur, s->attn_out, + DS4_N_EMBD) != 0; + /* ffn */ + if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(s->ffn_norm, + s->after_attn, model->map, model->size, + l->ffn_norm->abs_offset, DS4_N_EMBD, 1, DS4_RMS_EPS) != 0; + if (!ok) break; + if (il < DS4_N_LEADING_DENSE) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + s->dense_gate, s->dense_up, s->dense_mid, + model->map, model->size, + l->ffn_gate->abs_offset, l->ffn_up->abs_offset, + DS4_N_EMBD, DS4_N_FF_DENSE, s->ffn_norm, 0.0f) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->routed_out, + model->map, model->size, l->ffn_down->abs_offset, + DS4_N_FF_DENSE, DS4_N_EMBD, s->dense_mid, 1) != 0; + if (ok) ok = ds4_gpu_add_tensor(s->cur, s->after_attn, + s->routed_out, DS4_N_EMBD) != 0; + continue; + } + /* router: sigmoid + bias select, route-norm weights * scale */ + ok = ds4_gpu_matmul_f32_tensor(s->router_logits, model->map, + model->size, l->ffn_gate_inp->abs_offset, + DS4_N_EMBD, DS4_N_EXPERT, s->ffn_norm, 1) != 0; + if (ok) ok = ds4_gpu_glm_router_select_tensor(s->router_selected, + s->router_weights, s->router_probs, model->map, model->size, + l->ffn_exp_probs_b->abs_offset, s->router_logits, + DS4_N_EXPERT, DS4_N_EXPERT_USED, DS4_EXPERT_WEIGHT_SCALE) != 0; + uint64_t gate_in = 0, gate_out = 0, gate_row_bytes = 0; + uint64_t up_in = 0, up_out = 0, up_row_bytes = 0; + uint64_t down_in = 0, down_out = 0, down_row_bytes = 0; + (void)tensor_expert_bytes(model, l->ffn_gate_exps, 0, &gate_in, &gate_out, &gate_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_up_exps, 0, &up_in, &up_out, &up_row_bytes); + (void)tensor_expert_bytes(model, l->ffn_down_exps, 0, &down_in, &down_out, &down_row_bytes); + if (ok && s->ssd_streaming) { + const ds4_gpu_stream_expert_table table = { + .model_map = model->map, + .model_size = model->size, + .layer = il, + .n_total_expert = DS4_N_EXPERT, + .gate_offset = l->ffn_gate_exps->abs_offset, + .up_offset = l->ffn_up_exps->abs_offset, + .down_offset = l->ffn_down_exps->abs_offset, + .gate_expert_bytes = gate_out * gate_row_bytes, + .down_expert_bytes = down_out * down_row_bytes, + }; + ok = ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor( + &table, s->router_selected, DS4_N_EXPERT_USED) != 0; + } + /* shared expert */ + if (ok) ok = ds4_gpu_shared_mid_swiglu_q8_0_tensor(s->shared_mid, + model->map, model->size, + l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, DS4_N_FF_EXP, s->ffn_norm, 0.0f) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->shared_out, model->map, + model->size, l->ffn_down_shexp->abs_offset, + DS4_N_FF_EXP, DS4_N_EMBD, s->shared_mid, 1) != 0; + /* routed experts (streaming selected cache picked up inside) */ + if (ok) ok = ds4_gpu_routed_moe_one_tensor(s->routed_out, + s->routed_gate, s->routed_up, s->routed_mid, s->routed_down, + model->map, model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_down_exps->type, + gate_out * gate_row_bytes, gate_row_bytes, + down_out * down_row_bytes, down_row_bytes, + DS4_N_EMBD, DS4_N_FF_EXP, DS4_N_EMBD, + s->router_selected, s->router_weights, + DS4_N_EXPERT, DS4_N_EXPERT_USED, + 0.0f, s->ffn_norm, il, false) != 0; + if (ok) ok = ds4_gpu_add3_tensor(s->cur, s->after_attn, + s->routed_out, s->shared_out, DS4_N_EMBD) != 0; + } + if (ok && logits_out) { + ok = ds4_gpu_rms_norm_weight_rows_tensor(s->out_norm, s->cur, + model->map, model->size, weights->output_norm->abs_offset, + DS4_N_EMBD, 1, DS4_RMS_EPS) != 0; + if (ok) ok = ds4_gpu_matmul_q8_0_tensor(s->logits, model->map, + model->size, weights->output->abs_offset, + DS4_N_EMBD, DS4_N_VOCAB, s->out_norm, 1) != 0; + if (ok) ok = ds4_gpu_tensor_read(s->logits, 0, logits_out, + (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + } + return ok; +} + +static int generate_hy3_cuda_argmax( + const ds4_model * model, + const ds4_vocab * vocab, + const ds4_weights * weights, + const token_vec * prompt, + int n_predict, + int ctx_size, + bool ssd_streaming, + ds4_token_emit_fn emit, + ds4_generation_done_fn done, + void * emit_ud, + ds4_session_progress_fn progress, + void * progress_ud) { + fprintf(stderr, "ds4: using Hy3 CUDA argmax generation path\n"); + if (!prompt || prompt->len <= 0 || prompt->len >= ctx_size) { + fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); + return 1; + } + if (n_predict <= 0) { + if (done) done(emit_ud); + return 0; + } + hy3_gpu_state s; + if (!hy3_state_alloc(&s, model, weights, (uint32_t)ctx_size, ssd_streaming)) { + fprintf(stderr, "ds4: failed to allocate Hy3 CUDA state\n"); + return 1; + } + float *logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(logits[0])); + bool ok = true; + const double t_prefill0 = now_sec(); + for (int i = 0; ok && i < prompt->len; i++) { + const bool last = i + 1 == prompt->len; + ok = hy3_forward_token(&s, model, weights, prompt->v[i], + (uint32_t)i, last ? logits : NULL); + if (progress) progress(progress_ud, "prefill_chunk", i + 1, prompt->len); + } + const double t_prefill1 = now_sec(); + if (!ok) { + fprintf(stderr, "ds4: Hy3 prefill failed\n"); + free(logits); + hy3_state_free(&s); + return 1; + } + int n_generated = 0; + uint32_t pos = (uint32_t)prompt->len; + const bool token_timing = getenv("DS4_TOKEN_TIMING") != NULL; + const double t_decode0 = now_sec(); + for (int i = 0; i < n_predict && pos < s.ctx; i++) { + const int token = sample_argmax(logits, DS4_N_VOCAB); + if (vocab_token_is_generation_stop(vocab, token)) break; + if (emit) emit(emit_ud, token); + n_generated++; + if (i == n_predict - 1 || pos + 1u >= s.ctx) { + pos++; + break; + } + const double t_eval0 = token_timing ? now_sec() : 0.0; + ok = hy3_forward_token(&s, model, weights, token, pos, logits); + if (!ok) { + fprintf(stderr, "ds4: Hy3 decode failed at position %u\n", pos); + free(logits); + hy3_state_free(&s); + return 1; + } + if (token_timing) { + fprintf(stderr, "ds4: Hy3 decode eval %d took %.3f ms\n", + i + 1, (now_sec() - t_eval0) * 1000.0); + } + pos++; + } + const double t_decode1 = now_sec(); + if (done) done(emit_ud); + const double prefill_s = t_prefill1 - t_prefill0; + const double decode_s = t_decode1 - t_decode0; + ds4_log(stderr, + DS4_LOG_TIMING, + "ds4: Hy3 prefill: %.2f t/s, generation: %.2f t/s\n", + prefill_s > 0.0 ? (double)prompt->len / prefill_s : 0.0, + decode_s > 0.0 ? (double)n_generated / decode_s : 0.0); + free(logits); + hy3_state_free(&s); + return 0; +} + /* Metal generation entry point. The model runs as one local whole-graph * pipeline: graph prefill followed by graph decode steps. Streaming PRO may * use decode-style prefill for short prompts. */ @@ -33343,6 +34038,22 @@ static int generate_metal_graph_raw_swa( fprintf(stderr, "ds4: prompt is empty or exceeds context size\n"); return 1; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + /* ponytail: no --power/--prefill-chunk/steering support in v1; the + * straight-line path ignores them rather than growing option glue. */ + return generate_hy3_cuda_argmax(model, + vocab, + weights, + prompt, + n_predict, + ctx_size, + ssd_streaming, + emit, + done, + emit_ud, + progress, + progress_ud); + } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (power_percent > 0 && power_percent < 100) { fprintf(stderr, "ds4: --power is not supported by the GLM Metal path yet\n"); @@ -33709,6 +34420,8 @@ struct ds4_session { ds4_glm_gpu_graph glm_graph; bool glm_graph_ready; uint32_t glm_dense_cache_len; + hy3_gpu_state *hy3_state; + bool hy3_ready; #endif ds4_kv_cache cpu_cache; ds4_cpu_decode_scratch cpu_scratch; @@ -34208,6 +34921,10 @@ static bool ds4_session_is_glm(const ds4_session *s) { return s && s->engine && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; } +static bool ds4_session_is_hy3(const ds4_session *s) { + return s && s->engine && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3; +} + #ifndef DS4_NO_GPU static void ds4_session_glm_reset_dense_cache(ds4_session *s) { if (s) s->glm_dense_cache_len = 0; @@ -38673,6 +39390,27 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { load_layer_start, load_layer_end, load_output); + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + if (!opt->inspect_only) { + if (e->backend != DS4_BACKEND_CUDA || + !ds4_backend_uses_graph(e->backend)) { + fprintf(stderr, + "ds4: Hy3 (hy_v3) inference is CUDA-only in this " + "fork; use --cuda or --inspect\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } + if (opt->mtp_path && opt->mtp_path[0]) { + fprintf(stderr, "ds4: --mtp is not wired for Hy3 yet\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } + } + /* fall through: the generic tail below handles vocab_load, + * streaming budgets, and inspect_only. */ + } if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { if (opt->inspect_only) { *out = e; @@ -38717,10 +39455,11 @@ int ds4_engine_open(ds4_engine **out, const ds4_engine_options *opt) { *out = e; return 0; } - if (e->backend != DS4_BACKEND_METAL || !ds4_backend_uses_graph(e->backend)) { + if ((e->backend != DS4_BACKEND_METAL && e->backend != DS4_BACKEND_CUDA) || + !ds4_backend_uses_graph(e->backend)) { fprintf(stderr, "ds4: GLM 5.2 tensor layout is recognized, but GLM inference is " - "currently Metal-only; use --inspect, --cpu --first-token-test, " + "currently Metal/CUDA-only; use --inspect, --cpu --first-token-test, " "--metal --metal-graph-test, or limited --metal generation\n"); ds4_engine_close(e); *out = NULL; @@ -39219,6 +39958,10 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { fprintf(stderr, "ds4: GLM sessions currently require --metal\n"); return 1; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + fprintf(stderr, "ds4: Hy3 sessions require --cuda\n"); + return 1; + } if (e->distributed.role == DS4_DISTRIBUTED_COORDINATOR) { fprintf(stderr, "ds4: distributed coordinator sessions require the graph backend\n"); return 1; @@ -39301,6 +40044,20 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { *out = s; return 0; } + if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_HY_V3) { + s->hy3_state = xcalloc(1, sizeof(*s->hy3_state)); + if (!hy3_state_alloc(s->hy3_state, &e->model, &e->weights, + (uint32_t)ctx_size, e->ssd_streaming)) { + free(s->hy3_state); + free(s); + return 1; + } + s->hy3_ready = true; + s->prefill_cap = (uint32_t)ctx_size; + s->logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->logits[0])); + *out = s; + return 0; + } s->prefill_cap = metal_graph_prefill_cap_for_prompt(ctx_size, e->prefill_chunk); const uint32_t raw_cap = metal_graph_raw_cap_for_context(ctx_size, s->prefill_cap); @@ -39367,7 +40124,12 @@ void ds4_session_free(ds4_session *s) { } #ifndef DS4_NO_GPU else { - if (ds4_session_is_glm(s)) { + if (ds4_session_is_hy3(s)) { + if (s->hy3_state) { + hy3_state_free(s->hy3_state); + free(s->hy3_state); + } + } else if (ds4_session_is_glm(s)) { glm_graph_free(&s->glm_graph); } else { metal_graph_free(&s->graph); @@ -40193,6 +40955,44 @@ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t ds4_engine *e = s->engine; const char *backend_name = ds4_backend_name(e->backend); + if (ds4_session_is_hy3(s)) { + if (!s->hy3_ready) { + snprintf(err, errlen, "%s Hy3 state is not initialized", backend_name); + return 1; + } + /* Reuse the running KV cache when the new prompt extends the last one; + * the per-layer caches are positional so we only forward the suffix. */ + int start = 0; + if (s->checkpoint_valid && + prompt->len >= s->checkpoint.len && + ds4_tokens_starts_with(prompt, &s->checkpoint)) { + start = s->checkpoint.len; + } else { + s->checkpoint.len = 0; + s->checkpoint_valid = false; + } + for (int i = start; i < prompt->len; i++) { + if (ds4_session_cancelled(s)) { + snprintf(err, errlen, "interrupted"); + s->checkpoint_valid = (start != 0); + return DS4_SESSION_SYNC_INTERRUPTED; + } + const bool last = (i + 1 == prompt->len); + if (!hy3_forward_token(s->hy3_state, &e->model, &e->weights, + prompt->v[i], (uint32_t)i, + last ? s->logits : NULL)) { + snprintf(err, errlen, "Hy3 prefill failed"); + s->checkpoint_valid = false; + return 1; + } + token_vec_push(&s->checkpoint, prompt->v[i]); + if (s->progress) s->progress(s->progress_ud, "prefill_chunk", i + 1, prompt->len); + } + s->checkpoint_valid = true; + s->mtp_draft_valid = false; + return 0; + } + if (ds4_session_is_glm(s)) { if (!s->glm_graph_ready) { snprintf(err, errlen, "%s GLM graph is not initialized", backend_name); @@ -41056,6 +41856,30 @@ static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp, return 1; #else ds4_engine *e = s->engine; + if (ds4_session_is_hy3(s)) { + (void)probe_mtp; + if (!s->hy3_ready) { + if (errlen) snprintf(err, errlen, "%s Hy3 state is not initialized", + ds4_backend_name(e->backend)); + return 1; + } + if ((uint32_t)s->checkpoint.len >= (uint32_t)s->ctx_size) { + if (errlen) snprintf(err, errlen, "Hy3 context reached (%d)", s->ctx_size); + return 1; + } + const uint32_t pos = (uint32_t)s->checkpoint.len; + if (!hy3_forward_token(s->hy3_state, &e->model, &e->weights, + token, pos, s->logits)) { + if (errlen) snprintf(err, errlen, "%s Hy3 decode failed", + ds4_backend_name(e->backend)); + s->checkpoint_valid = false; + return 1; + } + token_vec_push(&s->checkpoint, token); + s->checkpoint_valid = true; + s->mtp_draft_valid = false; + return 0; + } if (ds4_session_is_glm(s)) { if (!s->glm_graph_ready) { if (errlen) snprintf(err, errlen, "%s GLM graph is not initialized", @@ -41179,6 +42003,14 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token, accepted[0] = first_token; return 1; } + if (ds4_session_is_hy3(s)) { + (void)max_tokens; + (void)eos_token; + if (!accepted || accepted_cap <= 0) return 0; + if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1; + accepted[0] = first_token; + return 1; + } #ifdef DS4_NO_GPU (void)s; (void)first_token; (void)max_tokens; (void)eos_token; (void)accepted; (void)accepted_cap; diff --git a/ds4_cli.c b/ds4_cli.c index 6278c1bb2..ca01372d9 100644 --- a/ds4_cli.c +++ b/ds4_cli.c @@ -1937,6 +1937,19 @@ static cli_config parse_options(int argc, char **argv) { } int main(int argc, char **argv) { + /* Kernel self-tests that need no model file; handled before option + * parsing so no other flags are required. */ + for (int i = 1; i < argc; i++) { + if (!strcmp(argv[i], "--gqa-selftest")) { +#ifndef DS4_NO_GPU + extern int ds4_gpu_gqa_selftest(void); + return ds4_gpu_gqa_selftest() ? 0 : 1; +#else + fprintf(stderr, "ds4: --gqa-selftest requires a GPU build\n"); + return 2; +#endif + } + } cli_config cfg = parse_options(argc, argv); if (cfg.gen.dump_tokens) { if (cfg.gen.prompt == NULL) { diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 2e85e95bc..1d4a0a8e9 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -18,6 +18,29 @@ #include #include +#if defined(__linux__) && defined(DS4_USE_IO_URING) +#include +#endif + +/* CUDA < 12.2 lacks the cudaMemLocation overloads of cudaMemAdvise() and + * cudaMemPrefetchAsync(). Forward to the legacy int-device signatures. */ +#if CUDART_VERSION < 12020 +static inline cudaError_t cudaMemAdvise(const void *ptr, size_t count, + cudaMemoryAdvise advice, + cudaMemLocation loc) { + int dev = loc.type == cudaMemLocationTypeDevice ? loc.id : cudaCpuDeviceId; + return cudaMemAdvise(ptr, count, advice, dev); +} +static inline cudaError_t cudaMemPrefetchAsync(const void *ptr, size_t count, + cudaMemLocation loc, + unsigned int flags, + cudaStream_t stream) { + (void)flags; + int dev = loc.type == cudaMemLocationTypeDevice ? loc.id : cudaCpuDeviceId; + return cudaMemPrefetchAsync(ptr, count, dev, stream); +} +#endif + #include "ds4_gpu.h" #ifndef M_PI @@ -192,12 +215,19 @@ static std::unordered_map g_q8_f16_by_offset; static std::vector g_q8_f32_ranges; static std::unordered_map g_q8_f32_by_offset; static cuda_stream_selected_cache g_stream_selected_cache; -static cuda_stream_expert_cache g_stream_expert_cache; +/* Expert caches are kept per (gate,down) byte-size class: mixed-quant models + * (e.g. the q2-q4-imatrix Flash GGUF) interleave layers whose routed experts + * have different byte sizes. A single cache keyed on the last-seen size gets + * released and reallocated on every layer-size transition, so it never + * accumulates hits (decode collapses to ~1 t/s with permanent cap-log spam). + * Keeping one small cache per size class removes the thrash. */ +#define DS4_CUDA_STREAM_EXPERT_CLASSES 4 +static cuda_stream_expert_cache g_stream_expert_caches[DS4_CUDA_STREAM_EXPERT_CLASSES]; static uint32_t g_stream_expert_budget_override; -static uint32_t g_stream_expert_runtime_cap; -static uint32_t g_stream_expert_memory_cap_notice; -static uint64_t g_stream_expert_runtime_gate_bytes; -static uint64_t g_stream_expert_runtime_down_bytes; +static uint32_t g_stream_expert_runtime_caps[DS4_CUDA_STREAM_EXPERT_CLASSES]; +static uint32_t g_stream_expert_memory_cap_notices[DS4_CUDA_STREAM_EXPERT_CLASSES]; +static uint64_t g_stream_expert_class_gate_bytes[DS4_CUDA_STREAM_EXPERT_CLASSES]; +static uint64_t g_stream_expert_class_down_bytes[DS4_CUDA_STREAM_EXPERT_CLASSES]; static uint64_t g_model_range_bytes; static uint64_t g_q8_f16_bytes; static uint64_t g_q8_f32_bytes; @@ -1175,15 +1205,75 @@ static char *cuda_model_arena_alloc(uint64_t bytes, const char *what) { const uint64_t chunk = cuda_model_arena_chunk_bytes(aligned); void *dev = NULL; - cudaError_t err = cudaMalloc(&dev, (size_t)chunk); + cudaError_t err = cudaSuccess; + /* Optional VRAM headroom: skip device chunks once free VRAM drops below + * the reserve so batch/verify kernels can still launch mid-decode. The + * chunk then lands in the pinned-host fallback below. Default 0 keeps the + * greedy fill-VRAM-first behavior. */ + static int reserve_parsed = -1; + static uint64_t reserve_bytes = 0; + if (reserve_parsed < 0) { + reserve_parsed = 0; + const char *renv = getenv("DS4_CUDA_ARENA_VRAM_RESERVE_GB"); + if (renv && renv[0]) { + char *rend = NULL; + unsigned long long rv = strtoull(renv, &rend, 10); + if (rend != renv) reserve_bytes = (uint64_t)rv * 1073741824ull; + } + } + if (reserve_bytes != 0) { + size_t free_b = 0, total_b = 0; + if (cudaMemGetInfo(&free_b, &total_b) == cudaSuccess && + (uint64_t)free_b < reserve_bytes + chunk) { + err = cudaErrorMemoryAllocation; + } + } + if (err == cudaSuccess) 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)); + /* Device arena is full: fall back to pinned host memory mapped + * into the GPU's address space (zero-copy). This keeps models + * whose non-routed weights exceed available VRAM working, at + * the cost of PCIe latency on every access to this chunk. Set + * DS4_CUDA_NO_PINNED_ARENA_FALLBACK to restore the old + * fail-fast behavior. */ (void)cudaGetLastError(); - g_model_cache_full = 1; - return NULL; + if (getenv("DS4_CUDA_NO_PINNED_ARENA_FALLBACK") != NULL) { + fprintf(stderr, "ds4: CUDA model arena alloc failed for %s (%.2f MiB chunk): %s\n", + what ? what : "weights", + (double)chunk / 1048576.0, + cudaGetErrorString(err)); + g_model_cache_full = 1; + return NULL; + } + void *host_ptr = NULL; + cudaError_t herr = cudaHostAlloc(&host_ptr, (size_t)chunk, cudaHostAllocMapped); + if (herr != cudaSuccess) { + fprintf(stderr, "ds4: CUDA model arena alloc failed for %s (%.2f MiB chunk): %s " + "(pinned fallback also failed: %s)\n", + what ? what : "weights", + (double)chunk / 1048576.0, + cudaGetErrorString(err), + cudaGetErrorString(herr)); + (void)cudaGetLastError(); + g_model_cache_full = 1; + return NULL; + } + void *dev_ptr = NULL; + herr = cudaHostGetDevicePointer(&dev_ptr, host_ptr, 0); + if (herr != cudaSuccess || !dev_ptr) { + fprintf(stderr, "ds4: CUDA model arena pinned fallback device pointer lookup failed for %s: %s\n", + what ? what : "weights", + cudaGetErrorString(herr)); + (void)cudaGetLastError(); + (void)cudaFreeHost(host_ptr); + g_model_cache_full = 1; + return NULL; + } + fprintf(stderr, "ds4: CUDA model arena using pinned host RAM fallback for %s " + "(%.2f MiB chunk, zero-copy device access)\n", + what ? what : "weights", + (double)chunk / 1048576.0); + dev = dev_ptr; } g_model_arenas.push_back({(char *)dev, chunk, aligned}); if (getenv("DS4_CUDA_WEIGHT_CACHE_VERBOSE")) { @@ -1428,27 +1518,52 @@ static void cuda_stream_selected_cache_release(void) { memset(&g_stream_selected_cache, 0, sizeof(g_stream_selected_cache)); } -static void cuda_stream_expert_cache_release_all(void) { - if (g_stream_expert_cache.gate_ptr) { - (void)cudaFree(g_stream_expert_cache.gate_ptr); +static void cuda_stream_expert_cache_release_class(int class_idx) { + cuda_stream_expert_cache *cache = &g_stream_expert_caches[class_idx]; + if (cache->gate_ptr) { + (void)cudaFree(cache->gate_ptr); + } + if (cache->up_ptr) { + (void)cudaFree(cache->up_ptr); } - if (g_stream_expert_cache.up_ptr) { - (void)cudaFree(g_stream_expert_cache.up_ptr); + if (cache->down_ptr) { + (void)cudaFree(cache->down_ptr); } - if (g_stream_expert_cache.down_ptr) { - (void)cudaFree(g_stream_expert_cache.down_ptr); + std::vector().swap(cache->slots); + cache->valid = 0; + cache->capacity = 0; + cache->count = 0; + cache->tick = 0; + cache->gate_expert_bytes = 0; + cache->down_expert_bytes = 0; + cache->gate_ptr = NULL; + cache->up_ptr = NULL; + cache->down_ptr = NULL; + cache->gate_capacity = 0; + cache->up_capacity = 0; + cache->down_capacity = 0; +} + +static void cuda_stream_expert_cache_release_all(void) { + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + cuda_stream_expert_cache_release_class(i); + g_stream_expert_class_gate_bytes[i] = 0; + g_stream_expert_class_down_bytes[i] = 0; + g_stream_expert_runtime_caps[i] = 0; + g_stream_expert_memory_cap_notices[i] = 0; } - g_stream_expert_cache.slots.clear(); - memset(&g_stream_expert_cache, 0, sizeof(g_stream_expert_cache)); } static void cuda_stream_expert_cache_invalidate(void) { - for (cuda_stream_expert_cache_slot &slot : g_stream_expert_cache.slots) { - slot.valid = 0; + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + cuda_stream_expert_cache *cache = &g_stream_expert_caches[i]; + for (cuda_stream_expert_cache_slot &slot : cache->slots) { + slot.valid = 0; + } + cache->valid = 0; + cache->count = 0; + cache->tick = 0; } - g_stream_expert_cache.valid = 0; - g_stream_expert_cache.count = 0; - g_stream_expert_cache.tick = 0; } static uint32_t cuda_stream_expert_cache_requested_budget(void) { @@ -1471,10 +1586,26 @@ static uint32_t cuda_stream_expert_cache_requested_budget(void) { return cap; } +static uint32_t cuda_stream_expert_cache_configured_budget_class(int class_idx) { + uint32_t cap = cuda_stream_expert_cache_requested_budget(); + if (g_stream_expert_runtime_caps[class_idx] != 0 && + cap > g_stream_expert_runtime_caps[class_idx]) { + cap = g_stream_expert_runtime_caps[class_idx]; + } + return cap; +} + +/* Class-less view (reporting/shared queries): apply the most restrictive + * runtime cap across the active classes. */ static uint32_t cuda_stream_expert_cache_configured_budget(void) { uint32_t cap = cuda_stream_expert_cache_requested_budget(); - if (g_stream_expert_runtime_cap != 0 && cap > g_stream_expert_runtime_cap) { - cap = g_stream_expert_runtime_cap; + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + if ((g_stream_expert_class_gate_bytes[i] != 0 || + g_stream_expert_class_down_bytes[i] != 0) && + g_stream_expert_runtime_caps[i] != 0 && + cap > g_stream_expert_runtime_caps[i]) { + cap = g_stream_expert_runtime_caps[i]; + } } return cap; } @@ -1508,6 +1639,7 @@ static uint64_t cuda_stream_expert_cache_reserve_bytes(void) { } static uint32_t cuda_stream_expert_cache_live_budget( + int class_idx, uint32_t requested, uint64_t gate_expert_bytes, uint64_t down_expert_bytes, @@ -1547,13 +1679,13 @@ static uint32_t cuda_stream_expert_cache_live_budget( reserve = total_bytes / 2ull; } if (free_bytes <= reserve) { - if (report && g_stream_expert_memory_cap_notice != requested) { + if (report && g_stream_expert_memory_cap_notices[class_idx] != requested) { cuda_model_load_progress_finish(); fprintf(stderr, "ds4: CUDA streaming expert cache disabled: available %.2f GiB <= reserve %.2f GiB\n", (double)free_bytes / 1073741824.0, (double)reserve / 1073741824.0); - g_stream_expert_memory_cap_notice = requested; + g_stream_expert_memory_cap_notices[class_idx] = requested; } return 0; } @@ -1563,7 +1695,8 @@ static uint32_t cuda_stream_expert_cache_live_budget( if (max_slots64 > UINT32_MAX) max_slots64 = UINT32_MAX; uint32_t capped = requested; if ((uint64_t)capped > max_slots64) capped = (uint32_t)max_slots64; - if (report && capped != requested && g_stream_expert_memory_cap_notice != capped) { + if (report && capped != requested && + g_stream_expert_memory_cap_notices[class_idx] != capped) { cuda_model_load_progress_finish(); fprintf(stderr, "ds4: CUDA streaming expert cache capped from %u to %u experts " @@ -1573,7 +1706,7 @@ static uint32_t cuda_stream_expert_cache_live_budget( (double)free_bytes / 1073741824.0, (double)reserve / 1073741824.0, (double)per_expert_bytes / 1048576.0); - g_stream_expert_memory_cap_notice = capped; + g_stream_expert_memory_cap_notices[class_idx] = capped; } return capped; } @@ -1589,17 +1722,36 @@ static uint64_t cuda_stream_expert_cache_expert_bytes( return gate_expert_bytes * 2ull + down_expert_bytes; } -static void cuda_stream_expert_cache_note_size( +/* Return the size-class index for (gate,down), assigning a free slot on + * first sight. If distinct classes exceed DS4_CUDA_STREAM_EXPERT_CLASSES + * (never observed in official GGUFs), class 0 is recycled: that size falls + * back to the pre-patch behavior (cache rebuilt) without touching the + * others. */ +static int cuda_stream_expert_cache_class_index( uint64_t gate_expert_bytes, uint64_t down_expert_bytes) { - if (g_stream_expert_runtime_gate_bytes == gate_expert_bytes && - g_stream_expert_runtime_down_bytes == down_expert_bytes) { - return; - } - g_stream_expert_runtime_gate_bytes = gate_expert_bytes; - g_stream_expert_runtime_down_bytes = down_expert_bytes; - g_stream_expert_runtime_cap = 0; - g_stream_expert_memory_cap_notice = 0; + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + if (g_stream_expert_class_gate_bytes[i] == gate_expert_bytes && + g_stream_expert_class_down_bytes[i] == down_expert_bytes) { + return i; + } + } + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + if (g_stream_expert_class_gate_bytes[i] == 0 && + g_stream_expert_class_down_bytes[i] == 0) { + g_stream_expert_class_gate_bytes[i] = gate_expert_bytes; + g_stream_expert_class_down_bytes[i] = down_expert_bytes; + g_stream_expert_runtime_caps[i] = 0; + g_stream_expert_memory_cap_notices[i] = 0; + return i; + } + } + cuda_stream_expert_cache_release_class(0); + g_stream_expert_class_gate_bytes[0] = gate_expert_bytes; + g_stream_expert_class_down_bytes[0] = down_expert_bytes; + g_stream_expert_runtime_caps[0] = 0; + g_stream_expert_memory_cap_notices[0] = 0; + return 0; } static uint32_t cuda_stream_expert_cache_shrunken_cap(uint32_t cap) { @@ -1609,15 +1761,16 @@ static uint32_t cuda_stream_expert_cache_shrunken_cap(uint32_t cap) { } static void cuda_stream_expert_cache_note_oom_cap( + int class_idx, uint32_t failed_cap, uint32_t new_cap, uint64_t expert_bytes, const char *errstr) { - if (g_stream_expert_runtime_cap != 0 && - g_stream_expert_runtime_cap <= new_cap) { + if (g_stream_expert_runtime_caps[class_idx] != 0 && + g_stream_expert_runtime_caps[class_idx] <= new_cap) { return; } - g_stream_expert_runtime_cap = new_cap; + g_stream_expert_runtime_caps[class_idx] = new_cap; const uint32_t released = failed_cap > new_cap ? failed_cap - new_cap : 0; cuda_model_load_progress_finish(); @@ -1825,47 +1978,52 @@ static cuda_stream_expert_cache *cuda_stream_expert_cache_prepare( cuda_stream_expert_cache_expert_bytes(gate_expert_bytes, down_expert_bytes); if (expert_bytes == 0) return NULL; - cuda_stream_expert_cache_note_size(gate_expert_bytes, down_expert_bytes); + const int class_idx = + cuda_stream_expert_cache_class_index(gate_expert_bytes, + down_expert_bytes); + cuda_stream_expert_cache *cache = &g_stream_expert_caches[class_idx]; - const uint32_t requested_cap = cuda_stream_expert_cache_configured_budget(); + const uint32_t requested_cap = + cuda_stream_expert_cache_configured_budget_class(class_idx); if (requested_cap == 0) return NULL; if (target_cap == 0 || target_cap > requested_cap) target_cap = requested_cap; if (target_cap == 0) return NULL; const int same_dims = - g_stream_expert_cache.valid && - g_stream_expert_cache.gate_expert_bytes == gate_expert_bytes && - g_stream_expert_cache.down_expert_bytes == down_expert_bytes; - if (!same_dims && g_stream_expert_cache.valid) { - cuda_stream_expert_cache_release_all(); + cache->valid && + cache->gate_expert_bytes == gate_expert_bytes && + cache->down_expert_bytes == down_expert_bytes; + if (!same_dims && cache->valid) { + cuda_stream_expert_cache_release_class(class_idx); } if (same_dims && - g_stream_expert_cache.capacity != 0 && - g_stream_expert_cache.capacity >= target_cap && - g_stream_expert_cache.slots.size() == g_stream_expert_cache.capacity) { - return &g_stream_expert_cache; + cache->capacity != 0 && + cache->capacity >= target_cap && + cache->slots.size() == cache->capacity) { + return cache; } uint64_t reclaim_bytes = 0; if (same_dims && - g_stream_expert_cache.capacity != 0 && - (uint64_t)g_stream_expert_cache.capacity <= UINT64_MAX / expert_bytes) { - reclaim_bytes = (uint64_t)g_stream_expert_cache.capacity * expert_bytes; + cache->capacity != 0 && + (uint64_t)cache->capacity <= UINT64_MAX / expert_bytes) { + reclaim_bytes = (uint64_t)cache->capacity * expert_bytes; } uint32_t cap = - cuda_stream_expert_cache_live_budget(target_cap, + cuda_stream_expert_cache_live_budget(class_idx, + target_cap, gate_expert_bytes, down_expert_bytes, reclaim_bytes, reclaim_bytes == 0); if (cap == 0) return NULL; if (same_dims && - g_stream_expert_cache.capacity != 0 && - g_stream_expert_cache.capacity >= cap && - g_stream_expert_cache.slots.size() == g_stream_expert_cache.capacity) { - return &g_stream_expert_cache; + cache->capacity != 0 && + cache->capacity >= cap && + cache->slots.size() == cache->capacity) { + return cache; } - cuda_stream_expert_cache_release_all(); + cuda_stream_expert_cache_release_class(class_idx); while (cap != 0) { if ((uint64_t)cap > UINT64_MAX / gate_expert_bytes || (uint64_t)cap > UINT64_MAX / down_expert_bytes) { @@ -1886,13 +2044,15 @@ static cuda_stream_expert_cache *cuda_stream_expert_cache_prepare( &alloc_error)) { const uint32_t new_cap = cuda_stream_expert_cache_shrunken_cap(cap); - cuda_stream_expert_cache_note_oom_cap(cap, + cuda_stream_expert_cache_note_oom_cap(class_idx, + cap, new_cap, expert_bytes, alloc_error); cap = new_cap; if (cap != 0) { - cap = cuda_stream_expert_cache_live_budget(cap, + cap = cuda_stream_expert_cache_live_budget(class_idx, + cap, gate_expert_bytes, down_expert_bytes, 0, @@ -1902,32 +2062,32 @@ static cuda_stream_expert_cache *cuda_stream_expert_cache_prepare( } try { - g_stream_expert_cache.slots.resize(cap); + cache->slots.resize(cap); } catch (...) { fprintf(stderr, "ds4: CUDA streaming expert cache metadata allocation failed\n"); (void)cudaFree(gate_ptr); (void)cudaFree(up_ptr); (void)cudaFree(down_ptr); - cuda_stream_expert_cache_release_all(); + cuda_stream_expert_cache_release_class(class_idx); return NULL; } - g_stream_expert_cache.valid = 1; - g_stream_expert_cache.capacity = cap; - g_stream_expert_cache.count = 0; - g_stream_expert_cache.tick = 0; - g_stream_expert_cache.gate_expert_bytes = gate_expert_bytes; - g_stream_expert_cache.down_expert_bytes = down_expert_bytes; - g_stream_expert_cache.gate_ptr = gate_ptr; - g_stream_expert_cache.up_ptr = up_ptr; - g_stream_expert_cache.down_ptr = down_ptr; - g_stream_expert_cache.gate_capacity = + cache->valid = 1; + cache->capacity = cap; + cache->count = 0; + cache->tick = 0; + cache->gate_expert_bytes = gate_expert_bytes; + cache->down_expert_bytes = down_expert_bytes; + cache->gate_ptr = gate_ptr; + cache->up_ptr = up_ptr; + cache->down_ptr = down_ptr; + cache->gate_capacity = (uint64_t)cap * gate_expert_bytes; - g_stream_expert_cache.up_capacity = + cache->up_capacity = (uint64_t)cap * gate_expert_bytes; - g_stream_expert_cache.down_capacity = + cache->down_capacity = (uint64_t)cap * down_expert_bytes; - return &g_stream_expert_cache; + return cache; } return NULL; } @@ -2008,6 +2168,37 @@ static int cuda_stream_expert_cache_copy_to_compact( "streaming selected down cache copy"); } +/* Reverse of copy_to_compact: after a parallel fetch lands a missed expert in + * the compact buffers, mirror it into a cache slot device-to-device (cheap) + * instead of re-reading it from disk through the serial mmap path. */ +static int cuda_stream_expert_cache_backfill_from_compact( + cuda_stream_expert_cache *cache, + uint32_t cache_slot, + uint32_t compact_slot, + const char *compact_gate, + const char *compact_up, + const char *compact_down) { + const uint64_t gate_dst = (uint64_t)cache_slot * cache->gate_expert_bytes; + const uint64_t down_dst = (uint64_t)cache_slot * cache->down_expert_bytes; + const uint64_t gate_src = (uint64_t)compact_slot * cache->gate_expert_bytes; + const uint64_t down_src = (uint64_t)compact_slot * cache->down_expert_bytes; + return cuda_ok(cudaMemcpy(cache->gate_ptr + gate_dst, + compact_gate + gate_src, + (size_t)cache->gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming gate cache backfill") && + cuda_ok(cudaMemcpy(cache->up_ptr + gate_dst, + compact_up + gate_src, + (size_t)cache->gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming up cache backfill") && + cuda_ok(cudaMemcpy(cache->down_ptr + down_dst, + compact_down + down_src, + (size_t)cache->down_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming down cache backfill"); +} + static int cuda_stream_expert_cache_load_slot( cuda_stream_expert_cache *cache, const void *model_map, @@ -2149,6 +2340,788 @@ static int cuda_stream_layer_expert_ranges_valid( return 1; } +/* Host expert read cache. The Linux path reads experts with O_DIRECT, so + * the kernel page cache never retains them and every re-fetch of the same + * routed expert pays full disk latency. This optional cache keeps recently + * read small ranges (routed experts are ~MiB-sized) in host memory, keyed by + * file range. 16-way set-associative, LRU within the set. + * Enable with DS4_CUDA_HOST_EXPERT_CACHE_GB=N (default: off). + * ponytail: pageable malloc buffers, not pinned; upload from pageable is + * slower but allocation is cheap and it is still >4x faster than the disk. + * Single-threaded like the rest of the streaming globals. */ +#define CUDA_HOST_CACHE_WAYS 16u +#define CUDA_HOST_CACHE_MAX_ENTRY (64ull << 20) + +typedef struct { + uint64_t offset; + uint64_t bytes; + uint64_t age; + uint64_t uses; /* hit count; eviction is least-frequently-used, age breaks ties */ + void *buf; + uint32_t payload_off; /* payload start inside buf (O_DIRECT-aligned reads) */ + uint8_t aligned_buf; /* buf belongs to the aligned recycling pool */ + int valid; +} cuda_host_cache_slot; + +/* Forward decl: pool release lives with the fetch machinery below. */ +static void cuda_fetch_buf_put(void *ptr, uint64_t size); + +static void cuda_host_cache_release_buf(cuda_host_cache_slot *s) { + if (!s->buf) return; + if (s->aligned_buf) { + const uint64_t a = g_model_direct_align > 1 ? g_model_direct_align : 4096; + const uint64_t alloc = + ((s->payload_off + s->bytes + a - 1u) / a) * a; + cuda_fetch_buf_put(s->buf, alloc); + } else { + free(s->buf); + } + s->buf = NULL; +} + +static cuda_host_cache_slot *g_host_cache_slots; +static uint32_t g_host_cache_nslots; +static uint64_t g_host_cache_budget; +static uint64_t g_host_cache_used; +static uint64_t g_host_cache_tick; +static uint64_t g_host_cache_hits; +static uint64_t g_host_cache_misses; +static int g_host_cache_checked; + +static void cuda_host_cache_init(void) { + if (g_host_cache_checked) return; + g_host_cache_checked = 1; + const char *env = getenv("DS4_CUDA_HOST_EXPERT_CACHE_GB"); + if (!env || !env[0]) return; + char *end = NULL; + errno = 0; + unsigned long long gb = strtoull(env, &end, 10); + if (end == env || errno != 0 || gb == 0) return; + uint64_t budget = (uint64_t)gb << 30; + uint64_t n = budget / (1ull << 20); + if (n < CUDA_HOST_CACHE_WAYS) n = CUDA_HOST_CACHE_WAYS; + if (n > 262144) n = 262144; + g_host_cache_slots = + (cuda_host_cache_slot *)calloc((size_t)n, sizeof(cuda_host_cache_slot)); + if (!g_host_cache_slots) return; + g_host_cache_nslots = (uint32_t)n; + g_host_cache_budget = budget; + fprintf(stderr, + "ds4: CUDA host expert cache enabled: %llu GiB budget, %u slots\n", + gb, g_host_cache_nslots); +} + +static uint32_t cuda_host_cache_bucket(uint64_t offset) { + offset ^= offset >> 33; + offset *= 0xff51afd7ed558ccdull; + offset ^= offset >> 33; + return (uint32_t)(offset % g_host_cache_nslots); +} + +static void cuda_host_cache_report(void) { + const uint64_t total = g_host_cache_hits + g_host_cache_misses; + if (total != 0 && total % 2048 == 0) { + fprintf(stderr, + "ds4: CUDA host expert cache: %.1f%% hit rate (%llu lookups, %.2f GiB resident)\n", + 100.0 * (double)g_host_cache_hits / (double)total, + (unsigned long long)total, + (double)g_host_cache_used / 1073741824.0); + } +} + +static const void *cuda_host_cache_get(uint64_t offset, uint64_t bytes) { + if (g_host_cache_nslots == 0) return NULL; + const uint32_t h = cuda_host_cache_bucket(offset); + for (uint32_t i = 0; i < CUDA_HOST_CACHE_WAYS; i++) { + cuda_host_cache_slot *s = + &g_host_cache_slots[(h + i) % g_host_cache_nslots]; + if (s->valid && s->offset == offset && s->bytes == bytes) { + s->age = ++g_host_cache_tick; + s->uses++; + g_host_cache_hits++; + cuda_host_cache_report(); + return (const char *)s->buf + s->payload_off; + } + } + g_host_cache_misses++; + cuda_host_cache_report(); + return NULL; +} + +/* Quiet membership probe: no hit/miss accounting, no LRU/LFU bump. Used by + * the speculative prefetcher to skip experts that are already resident. */ +static int cuda_host_cache_contains(uint64_t offset, uint64_t bytes) { + if (g_host_cache_nslots == 0) return 0; + const uint32_t h = cuda_host_cache_bucket(offset); + for (uint32_t i = 0; i < CUDA_HOST_CACHE_WAYS; i++) { + const cuda_host_cache_slot *s = + &g_host_cache_slots[(h + i) % g_host_cache_nslots]; + if (s->valid && s->offset == offset && s->bytes == bytes) return 1; + } + return 0; +} + +static uint64_t g_host_cache_inserts; + +/* Takes ownership of buf; frees it if not inserted. The payload (bytes long) + * starts at buf + payload_off, which is nonzero for O_DIRECT-aligned reads. */ +static void cuda_host_cache_insert_owned(uint64_t offset, + uint64_t bytes, + void *buf, + uint32_t payload_off, + int aligned_buf) { + cuda_host_cache_slot incoming; + incoming.buf = buf; + incoming.bytes = bytes; + incoming.payload_off = payload_off; + incoming.aligned_buf = (uint8_t)(aligned_buf != 0); + if (g_host_cache_nslots == 0 || !buf) { + cuda_host_cache_release_buf(&incoming); + return; + } + /* Periodic decay so once-hot entries cannot squat forever. */ + if ((++g_host_cache_inserts & 4095u) == 0) { + for (uint32_t i = 0; i < g_host_cache_nslots; i++) { + g_host_cache_slots[i].uses >>= 1; + } + } + const uint32_t h = cuda_host_cache_bucket(offset); + cuda_host_cache_slot *victim = NULL; /* empty slot if available */ + cuda_host_cache_slot *lfu = NULL; /* least-used valid entry in window */ + for (uint32_t i = 0; i < CUDA_HOST_CACHE_WAYS; i++) { + cuda_host_cache_slot *s = + &g_host_cache_slots[(h + i) % g_host_cache_nslots]; + if (!s->valid) { + if (!victim) victim = s; + continue; + } + if (s->offset == offset && s->bytes == bytes) { + /* raced with an earlier insert of the same range */ + cuda_host_cache_release_buf(&incoming); + return; + } + if (!lfu || s->uses < lfu->uses || + (s->uses == lfu->uses && s->age < lfu->age)) { + lfu = s; + } + } + /* Under budget pressure, replace the window's least-used entry even when + * an empty slot exists: with large entries the budget fills long before + * the table does, and refusing inserts would freeze the cache contents. + * Frequency-weighted eviction protects repeatedly hit experts from being + * flushed by one-shot streams (pure LRU thrashes at these cache sizes). */ + if (g_host_cache_used + bytes > g_host_cache_budget && lfu) { + victim = lfu; + } else if (!victim) { + victim = lfu; + } + if (!victim) { + cuda_host_cache_release_buf(&incoming); + return; + } + if (victim->valid) { + g_host_cache_used -= victim->bytes; + cuda_host_cache_release_buf(victim); + victim->valid = 0; + } + if (g_host_cache_used + bytes > g_host_cache_budget) { + cuda_host_cache_release_buf(&incoming); /* window had nothing to reclaim */ + return; + } + victim->offset = offset; + victim->bytes = bytes; + victim->age = ++g_host_cache_tick; + victim->uses = 1; + victim->buf = buf; + victim->payload_off = payload_off; + victim->aligned_buf = incoming.aligned_buf; + victim->valid = 1; + g_host_cache_used += bytes; +} + +/* Parallel expert fetch. The selected-expert loader issues up to + * 3*n_experts serial read+upload+sync round trips per layer; on NVMe the + * drive sits idle between them. This batches the direct loads of one layer: + * host-cache hits upload immediately, misses are read by a small pread() + * thread pool on the buffered fd, then all uploads share one stream sync. + * Completed read buffers are donated to the host expert cache (no extra + * copy). Enable with DS4_CUDA_PARALLEL_FETCH_THREADS=N (default: off). + * ponytail: one-shot pthreads per layer batch and pageable upload buffers; + * persistent pool + pinned ring if profiling ever shows they matter. */ +typedef struct { + int kind; /* 0 = demand fetch, 1 = speculative prefetch */ + char *dst; + uint64_t offset; + uint64_t bytes; + void *host_buf; + int need_disk; + int ok; + int fd; /* source fd: model or bundle */ + /* bundle jobs carry one contiguous expert and fan out to 3 dsts */ + char *dst_up; + char *dst_down; + uint64_t part_gate; + uint64_t part_down; + /* O_DIRECT-aligned read bracket (io_uring path): read read_len bytes at + * aligned_off into host_buf; the payload starts at host_buf+payload_off. + * The pthread fallback reads buffered and leaves payload_off at 0. */ + uint64_t aligned_off; + uint64_t read_len; + uint64_t read_done; + uint32_t payload_off; + int read_fd; + int aligned_alloc; /* buffer came from the aligned recycling pool */ +} cuda_fetch_job; + +/* Recycling pool for the O_DIRECT-aligned read buffers. posix_memalign at + * ~3MB goes through mmap, so per-read alloc/free costs two page-table + * syscalls with TLB shootdowns per expert part; recycling makes buffer + * turnover free. Single-threaded like the rest of the streaming globals. */ +#define CUDA_FETCH_BUF_POOL_MAX 96 +static struct { void *ptr; uint64_t size; } g_fetch_buf_pool[CUDA_FETCH_BUF_POOL_MAX]; +static uint32_t g_fetch_buf_pool_n; + +static void *cuda_fetch_buf_get(uint64_t size, uint64_t align) { + for (uint32_t i = g_fetch_buf_pool_n; i-- > 0;) { + if (g_fetch_buf_pool[i].size == size) { + void *p = g_fetch_buf_pool[i].ptr; + g_fetch_buf_pool[i] = g_fetch_buf_pool[--g_fetch_buf_pool_n]; + return p; + } + } + void *p = NULL; + if (posix_memalign(&p, (size_t)align, (size_t)size) != 0) return NULL; + return p; +} + +static void cuda_fetch_buf_put(void *ptr, uint64_t size) { + if (!ptr) return; + if (g_fetch_buf_pool_n < CUDA_FETCH_BUF_POOL_MAX) { + g_fetch_buf_pool[g_fetch_buf_pool_n].ptr = ptr; + g_fetch_buf_pool[g_fetch_buf_pool_n].size = size; + g_fetch_buf_pool_n++; + return; + } + free(ptr); +} + +/* Release a fetch job's buffer to the right place. */ +static void cuda_fetch_job_release_buf(cuda_fetch_job *job) { + if (!job->host_buf) return; + if (job->aligned_alloc) { + cuda_fetch_buf_put(job->host_buf, job->read_len); + } else { + free(job->host_buf); + } + job->host_buf = NULL; +} + +typedef struct { + cuda_fetch_job *jobs; + uint32_t njobs; + volatile uint32_t next; +} cuda_fetch_pool; + +/* Expert bundle sidecar (CUDA port of the idea in upstream issue #491). + * Optional derived file where each routed expert's gate+up+down tensors are + * contiguous, so an expert miss is ONE pread instead of three scattered + * ones. Set DS4_CUDA_EXPERT_BUNDLE=/path/to/file. Only used by the parallel + * fetch path (DS4_CUDA_PARALLEL_FETCH_THREADS >= 1). */ +#define CUDA_BUNDLE_CACHE_KEY_BIT (1ull << 63) + +typedef struct { + uint32_t layer; + uint32_t n_experts; + uint64_t gate_bytes; + uint64_t down_bytes; + uint64_t base; +} cuda_bundle_layer; + +static int g_bundle_fd = -1; +static cuda_bundle_layer *g_bundle_layers; +static uint32_t g_bundle_nlayers; +static int g_bundle_checked; + +static void cuda_bundle_init(void) { + if (g_bundle_checked) return; + g_bundle_checked = 1; + const char *path = getenv("DS4_CUDA_EXPERT_BUNDLE"); + if (!path || !path[0]) return; + const int fd = open(path, O_RDONLY); + if (fd < 0) { + fprintf(stderr, "ds4: CUDA expert bundle open failed: %s\n", strerror(errno)); + return; + } + char magic[8]; + uint32_t n = 0; + if (pread(fd, magic, 8, 0) != 8 || memcmp(magic, "DS4BNDL1", 8) != 0 || + pread(fd, &n, 4, 8) != 4 || n == 0 || n > 512) { + fprintf(stderr, "ds4: CUDA expert bundle header invalid, ignoring\n"); + close(fd); + return; + } + cuda_bundle_layer *ls = + (cuda_bundle_layer *)calloc((size_t)n, sizeof(cuda_bundle_layer)); + if (!ls) { + close(fd); + return; + } + for (uint32_t i = 0; i < n; i++) { + unsigned char rec[32]; + if (pread(fd, rec, 32, 12 + (uint64_t)i * 32) != 32) { + fprintf(stderr, "ds4: CUDA expert bundle truncated header, ignoring\n"); + free(ls); + close(fd); + return; + } + memcpy(&ls[i].layer, rec, 4); + memcpy(&ls[i].n_experts, rec + 4, 4); + memcpy(&ls[i].gate_bytes, rec + 8, 8); + memcpy(&ls[i].down_bytes, rec + 16, 8); + memcpy(&ls[i].base, rec + 24, 8); + } + g_bundle_fd = fd; + g_bundle_layers = ls; + g_bundle_nlayers = n; + fprintf(stderr, "ds4: CUDA expert bundle enabled: %s (%u layers)\n", path, n); +} + +static const cuda_bundle_layer *cuda_bundle_layer_get(uint32_t layer, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (g_bundle_fd < 0) return NULL; + for (uint32_t i = 0; i < g_bundle_nlayers; i++) { + if (g_bundle_layers[i].layer == layer && + g_bundle_layers[i].gate_bytes == gate_expert_bytes && + g_bundle_layers[i].down_bytes == down_expert_bytes) { + return &g_bundle_layers[i]; + } + } + return NULL; +} + +static int cuda_parallel_fetch_threads(void) { + static int cached = -1; + if (cached >= 0) return cached; + cached = 0; + const char *env = getenv("DS4_CUDA_PARALLEL_FETCH_THREADS"); + if (env && env[0]) { + int v = atoi(env); + if (v > 0) cached = v > 64 ? 64 : v; + } + return cached; +} + +static void *cuda_fetch_worker(void *arg) { + cuda_fetch_pool *p = (cuda_fetch_pool *)arg; + for (;;) { + const uint32_t j = __sync_fetch_and_add(&p->next, 1u); + if (j >= p->njobs) return NULL; + cuda_fetch_job *job = &p->jobs[j]; + if (!job->need_disk) continue; + job->host_buf = malloc((size_t)job->bytes); + if (!job->host_buf) continue; + uint64_t done = 0; + while (done < job->bytes) { + const ssize_t r = pread(job->fd, + (char *)job->host_buf + done, + (size_t)(job->bytes - done), + (off_t)(job->offset + done)); + if (r <= 0) break; + done += (uint64_t)r; + } + if (done == job->bytes) { + job->ok = 1; + } else { + free(job->host_buf); + job->host_buf = NULL; + } + } +} + +static uint64_t cuda_fetch_cache_key(const cuda_fetch_job *job) { + return job->fd == g_bundle_fd ? (job->offset | CUDA_BUNDLE_CACHE_KEY_BIT) + : job->offset; +} + +/* io_uring fetch engine. Replaces the one-shot pthread pool with async + * submission at real queue depth, reading through the O_DIRECT fd with + * aligned brackets (same trick as cuda_model_stage_read) so misses bypass + * the page cache instead of churning it. Enabled when built with + * DS4_USE_IO_URING and the ring initializes; DS4_CUDA_FETCH_URING=0 falls + * back to the pthread pool, DS4_CUDA_FETCH_URING_QD sets queue depth. */ +#if defined(__linux__) && defined(DS4_USE_IO_URING) +static struct io_uring g_fetch_uring; +static int g_fetch_uring_state = -1; /* -1 unchecked, 0 off, 1 ready */ +static unsigned g_fetch_uring_qd = 64; + +static int cuda_fetch_uring_ready(void) { + if (g_fetch_uring_state >= 0) return g_fetch_uring_state; + g_fetch_uring_state = 0; + const char *env = getenv("DS4_CUDA_FETCH_URING"); + if (env && env[0] == '0') return 0; + const char *qd_env = getenv("DS4_CUDA_FETCH_URING_QD"); + if (qd_env && qd_env[0]) { + char *end = NULL; + unsigned long long v = strtoull(qd_env, &end, 10); + if (end != qd_env && v >= 8 && v <= 512) g_fetch_uring_qd = (unsigned)v; + } + if (io_uring_queue_init(g_fetch_uring_qd, &g_fetch_uring, 0) != 0) { + fprintf(stderr, "ds4: CUDA fetch io_uring init failed (%s); using thread pool\n", + strerror(errno)); + return 0; + } + fprintf(stderr, "ds4: CUDA fetch io_uring enabled (queue depth %u)\n", + g_fetch_uring_qd); + g_fetch_uring_state = 1; + return 1; +} +#endif + +/* Pick the read strategy for one disk job: an O_DIRECT-aligned bracket on + * the direct fd when the job targets the model file, else a plain buffered + * read. Allocates host_buf; returns 0 on allocation/geometry failure. */ +static int cuda_fetch_job_prepare_read(cuda_fetch_job *job) { + job->read_done = 0; + job->payload_off = 0; +#if defined(__linux__) && defined(O_DIRECT) + /* Buffered reads let the page cache hold part of the model, which wins + * when model size is within a small multiple of RAM (e.g. Flash 86GB on + * a 30GB box); O_DIRECT wins when the model dwarfs RAM (GLM 211GB) and + * cache churn is pure overhead. DS4_CUDA_FETCH_BUFFERED=1 forces the + * buffered path. */ + static int force_buffered = -1; + if (force_buffered < 0) { + const char *fb = getenv("DS4_CUDA_FETCH_BUFFERED"); + force_buffered = fb && fb[0] == '1'; + } + if (!force_buffered && + job->fd == g_model_fd && g_model_direct_fd >= 0 && + g_model_direct_align > 1 && g_model_file_size != 0) { + const uint64_t a = g_model_direct_align; + const uint64_t aligned_off = job->offset & ~(a - 1u); + const uint64_t payload = job->offset - aligned_off; + uint64_t read_len = ((payload + job->bytes + a - 1u) / a) * a; + if (aligned_off <= g_model_file_size && + read_len <= g_model_file_size - aligned_off && + payload <= UINT32_MAX) { + void *buf = cuda_fetch_buf_get(read_len, a); + if (buf) { + job->host_buf = buf; + job->read_fd = g_model_direct_fd; + job->aligned_off = aligned_off; + job->read_len = read_len; + job->payload_off = (uint32_t)payload; + job->aligned_alloc = 1; + return 1; + } + } + } +#endif + job->host_buf = malloc((size_t)job->bytes); + if (!job->host_buf) return 0; + job->read_fd = job->fd; + job->aligned_off = job->offset; + job->read_len = job->bytes; + job->aligned_alloc = 0; + return 1; +} + +#if defined(__linux__) && defined(DS4_USE_IO_URING) +/* Speculative prefetch pool: jobs launched from next-layer router predictions. + * Completions land in the host expert cache; slots recycle on completion. */ +#define CUDA_PREFETCH_SLOTS 48 +static cuda_fetch_job g_prefetch_jobs[CUDA_PREFETCH_SLOTS]; +static int g_prefetch_busy[CUDA_PREFETCH_SLOTS]; +static uint64_t g_prefetch_submitted; +static uint64_t g_prefetch_completed; +static unsigned g_fetch_uring_inflight; /* SQEs handed to the kernel, not yet reaped */ + +static int cuda_uring_submit_job(cuda_fetch_job *job) { + struct io_uring_sqe *sqe = io_uring_get_sqe(&g_fetch_uring); + if (!sqe) { + (void)io_uring_submit(&g_fetch_uring); + sqe = io_uring_get_sqe(&g_fetch_uring); + if (!sqe) return 0; + } + io_uring_prep_read(sqe, job->read_fd, + (char *)job->host_buf + job->read_done, + (unsigned)(job->read_len - job->read_done), + (off_t)(job->aligned_off + job->read_done)); + io_uring_sqe_set_data(sqe, job); + g_fetch_uring_inflight++; + return 1; +} + +/* Shared completion handler for demand and prefetch jobs. + * Returns 1 when the job left the ring (done or failed), 0 if resubmitted. */ +static int cuda_uring_handle_cqe(cuda_fetch_job *job, int res) { + int finished = 0; + if (res == -EINTR || res == -EAGAIN) { + /* transient: resubmit the same range */ + } else if (res <= 0) { + cuda_fetch_job_release_buf(job); + finished = 1; + } else { + job->read_done += (uint64_t)res; + if (job->read_done >= job->read_len) { + job->ok = 1; + finished = 1; + } + } + if (!finished && !cuda_uring_submit_job(job)) { + cuda_fetch_job_release_buf(job); + finished = 1; + } + if (finished && job->kind == 1) { + /* prefetch: donate a successful read to the host expert cache */ + if (job->ok && job->host_buf) { + cuda_host_cache_insert_owned(cuda_fetch_cache_key(job), + job->bytes, + job->host_buf, + job->payload_off, + job->aligned_alloc); + job->host_buf = NULL; + g_prefetch_completed++; + } + g_prefetch_busy[job - g_prefetch_jobs] = 0; + } + return finished; +} + +/* Block until a matching in-flight prefetch lands (or fails), then return + * the cached payload if it made it into the host cache. Runs only from the + * demand pre-pass, before any demand submissions, so every reaped CQE is a + * prefetch completion. Returns NULL when there is nothing to wait for. */ +static const void *cuda_prefetch_wait_for(uint64_t offset, uint64_t bytes) { + if (g_fetch_uring_state != 1) return NULL; + int slot = -1; + for (int s = 0; s < CUDA_PREFETCH_SLOTS; s++) { + if (g_prefetch_busy[s] && + g_prefetch_jobs[s].offset == offset && + g_prefetch_jobs[s].bytes == bytes) { + slot = s; + break; + } + } + if (slot < 0) return NULL; + while (g_prefetch_busy[slot]) { + (void)io_uring_submit(&g_fetch_uring); + struct io_uring_cqe *cqe = NULL; + if (io_uring_wait_cqe(&g_fetch_uring, &cqe) < 0) { + if (errno == EINTR) continue; + return NULL; + } + cuda_fetch_job *job = (cuda_fetch_job *)io_uring_cqe_get_data(cqe); + const int res = cqe->res; + io_uring_cqe_seen(&g_fetch_uring, cqe); + if (g_fetch_uring_inflight) g_fetch_uring_inflight--; + (void)cuda_uring_handle_cqe(job, res); + } + return cuda_host_cache_get(offset, bytes); +} + +/* Drain finished CQEs without blocking. Prefetch completions are absorbed + * here; demand jobs are never outstanding when this runs from the prefetch + * or pre-pass entry points, so nothing else consumes them. */ +static void cuda_fetch_uring_reap_nowait(void) { + if (g_fetch_uring_state != 1) return; + struct io_uring_cqe *cqe = NULL; + while (io_uring_peek_cqe(&g_fetch_uring, &cqe) == 0 && cqe) { + cuda_fetch_job *job = (cuda_fetch_job *)io_uring_cqe_get_data(cqe); + const int res = cqe->res; + io_uring_cqe_seen(&g_fetch_uring, cqe); + if (g_fetch_uring_inflight) g_fetch_uring_inflight--; + (void)cuda_uring_handle_cqe(job, res); + } + (void)io_uring_submit(&g_fetch_uring); /* flush any resubmitted SQEs */ +} + +/* Run all disk jobs through the ring with a sliding submission window. + * Short reads (O_DIRECT completes in block multiples) are resubmitted for + * the remainder. Failed jobs keep ok=0 and are reported by the caller. + * Prefetch completions arriving mid-drain are absorbed by the handler. */ +static void cuda_fetch_uring_execute(cuda_fetch_job *jobs, uint32_t njobs) { + uint32_t next = 0; + uint32_t remaining = 0; + for (uint32_t i = 0; i < njobs; i++) { + if (jobs[i].need_disk && jobs[i].host_buf) remaining++; + } + while (remaining != 0) { + unsigned queued = 0; + while (next < njobs) { + cuda_fetch_job *job = &jobs[next]; + if (!job->need_disk || !job->host_buf || job->ok) { + next++; + continue; + } + if (!cuda_uring_submit_job(job)) break; + next++; + queued++; + if (queued >= g_fetch_uring_qd) break; + } + const int wait_rc = io_uring_submit_and_wait(&g_fetch_uring, 1); + if (wait_rc < 0 && wait_rc != -EINTR) break; + struct io_uring_cqe *cqe = NULL; + int progressed = 0; + while (io_uring_peek_cqe(&g_fetch_uring, &cqe) == 0 && cqe) { + cuda_fetch_job *job = (cuda_fetch_job *)io_uring_cqe_get_data(cqe); + const int res = cqe->res; + io_uring_cqe_seen(&g_fetch_uring, cqe); + if (g_fetch_uring_inflight) g_fetch_uring_inflight--; + progressed = 1; + if (cuda_uring_handle_cqe(job, res) && job->kind == 0) { + remaining--; + } + } + if (!progressed && queued == 0 && g_fetch_uring_inflight == 0) break; + } + /* Never leave with the kernel still writing into job buffers the caller + * may free: drain every outstanding CQE before returning. */ + while (g_fetch_uring_inflight != 0) { + (void)io_uring_submit(&g_fetch_uring); /* flush resubmitted SQEs */ + struct io_uring_cqe *cqe = NULL; + if (io_uring_wait_cqe(&g_fetch_uring, &cqe) < 0) { + if (errno == EINTR) continue; + break; + } + cuda_fetch_job *job = (cuda_fetch_job *)io_uring_cqe_get_data(cqe); + const int res = cqe->res; + io_uring_cqe_seen(&g_fetch_uring, cqe); + g_fetch_uring_inflight--; + (void)cuda_uring_handle_cqe(job, res); + } +} +#endif + +/* Upload a job's payload to its device destination(s) on the shared stream. + * Bundle jobs fan one contiguous expert out to gate/up/down. */ +static int cuda_fetch_upload(const cuda_fetch_job *job, const void *payload) { + const char *src = (const char *)payload; + if (!job->dst_up) { + return cuda_ok(cudaMemcpyAsync(job->dst, src, (size_t)job->bytes, + cudaMemcpyHostToDevice, + g_stream_selected_upload_stream), + "parallel fetch upload"); + } + return cuda_ok(cudaMemcpyAsync(job->dst, src, (size_t)job->part_gate, + cudaMemcpyHostToDevice, + g_stream_selected_upload_stream), + "bundle gate upload") && + cuda_ok(cudaMemcpyAsync(job->dst_up, src + job->part_gate, + (size_t)job->part_gate, + cudaMemcpyHostToDevice, + g_stream_selected_upload_stream), + "bundle up upload") && + cuda_ok(cudaMemcpyAsync(job->dst_down, src + 2 * job->part_gate, + (size_t)job->part_down, + cudaMemcpyHostToDevice, + g_stream_selected_upload_stream), + "bundle down upload"); +} + +/* Runs a batch of expert fetch jobs. Returns 1 on success. Buffers of + * successful disk reads are donated to the host expert cache. */ +static int cuda_fetch_execute(cuda_fetch_job *jobs, uint32_t njobs) { + if (njobs == 0) return 1; + if (!g_stream_selected_upload_stream) { + if (!cuda_ok(cudaStreamCreateWithFlags(&g_stream_selected_upload_stream, + cudaStreamNonBlocking), + "parallel fetch stream creation")) { + return 0; + } + } + cuda_host_cache_init(); +#if defined(__linux__) && defined(DS4_USE_IO_URING) + if (g_fetch_uring_state == 1) cuda_fetch_uring_reap_nowait(); +#endif + + uint32_t disk_jobs = 0; + for (uint32_t i = 0; i < njobs; i++) { + cuda_fetch_job *job = &jobs[i]; + const void *cached = cuda_host_cache_get(cuda_fetch_cache_key(job), + job->bytes); +#if defined(__linux__) && defined(DS4_USE_IO_URING) + if (!cached && job->fd == g_model_fd) { + /* a speculative prefetch may already be reading this range: + * wait for it instead of issuing a duplicate disk read */ + cached = cuda_prefetch_wait_for(job->offset, job->bytes); + } +#endif + if (cached) { + if (!cuda_fetch_upload(job, cached)) return 0; + job->ok = 1; + } else { + job->need_disk = 1; + disk_jobs++; + } + } + + if (disk_jobs != 0) { + int ran_uring = 0; +#if defined(__linux__) && defined(DS4_USE_IO_URING) + if (cuda_fetch_uring_ready()) { + ran_uring = 1; + for (uint32_t i = 0; i < njobs; i++) { + cuda_fetch_job *job = &jobs[i]; + if (job->need_disk) (void)cuda_fetch_job_prepare_read(job); + } + cuda_fetch_uring_execute(jobs, njobs); + } +#endif + if (!ran_uring) { + cuda_fetch_pool pool; + pool.jobs = jobs; + pool.njobs = njobs; + pool.next = 0; + int nthreads = cuda_parallel_fetch_threads(); + if ((uint32_t)nthreads > disk_jobs) nthreads = (int)disk_jobs; + pthread_t tids[64]; + int spawned = 0; + for (int t = 0; t < nthreads; t++) { + if (pthread_create(&tids[t], NULL, cuda_fetch_worker, &pool) != 0) break; + spawned++; + } + if (spawned == 0) cuda_fetch_worker(&pool); /* degraded: run inline */ + for (int t = 0; t < spawned; t++) pthread_join(tids[t], NULL); + } + + for (uint32_t i = 0; i < njobs; i++) { + cuda_fetch_job *job = &jobs[i]; + if (!job->need_disk) continue; + if (!job->ok) { + fprintf(stderr, + "ds4: CUDA parallel expert fetch read failed at offset %.2f GiB\n", + (double)job->offset / 1073741824.0); + goto fail; + } + if (!cuda_fetch_upload(job, (const char *)job->host_buf + job->payload_off)) goto fail; + } + } + + if (!cuda_ok(cudaStreamSynchronize(g_stream_selected_upload_stream), + "parallel fetch sync")) { + goto fail; + } + for (uint32_t i = 0; i < njobs; i++) { + if (jobs[i].need_disk && jobs[i].ok && jobs[i].host_buf) { + cuda_host_cache_insert_owned(cuda_fetch_cache_key(&jobs[i]), + jobs[i].bytes, jobs[i].host_buf, + jobs[i].payload_off, + jobs[i].aligned_alloc); + jobs[i].host_buf = NULL; + } + } + return 1; + +fail: + (void)cudaStreamSynchronize(g_stream_selected_upload_stream); + for (uint32_t i = 0; i < njobs; i++) { + cuda_fetch_job_release_buf(&jobs[i]); + } + return 0; +} + static int cuda_model_copy_to_device_streamed( char *dst, const void *model_map, @@ -2169,9 +3142,26 @@ static int cuda_model_copy_to_device_streamed( what ? what : "stream selected expert copy"); } + cuda_host_cache_init(); + if (bytes <= CUDA_HOST_CACHE_MAX_ENTRY) { + const void *cached = cuda_host_cache_get(offset, bytes); + if (cached) { + return cuda_ok(cudaMemcpy(dst, cached, (size_t)bytes, + cudaMemcpyHostToDevice), + what ? what : "host cache expert copy"); + } + } + char *capture = NULL; + if (g_host_cache_nslots != 0 && bytes <= CUDA_HOST_CACHE_MAX_ENTRY) { + capture = (char *)malloc((size_t)bytes); /* NULL is fine: skip caching */ + } + 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; + if (!cuda_stream_selected_stage_pool_alloc(stage_bytes)) { + free(capture); + return 0; + } cudaError_t err = cudaSuccess; uint64_t copied = 0; @@ -2187,6 +3177,7 @@ static int cuda_model_copy_to_device_streamed( what ? what : "expert", cudaGetErrorString(err)); (void)cudaGetLastError(); + free(capture); return 0; } } @@ -2202,8 +3193,10 @@ static int cuda_model_copy_to_device_streamed( what ? what : "expert", (double)copied / 1048576.0, strerror(errno)); + free(capture); return 0; } + if (capture) memcpy(capture + copied, payload, (size_t)n); err = cudaMemcpyAsync(dst + copied, payload, (size_t)n, @@ -2216,6 +3209,7 @@ static int cuda_model_copy_to_device_streamed( (double)copied / 1048576.0, cudaGetErrorString(err)); (void)cudaGetLastError(); + free(capture); return 0; } err = cudaEventRecord(g_stream_selected_stage_event[bi], @@ -2226,6 +3220,7 @@ static int cuda_model_copy_to_device_streamed( what ? what : "expert", cudaGetErrorString(err)); (void)cudaGetLastError(); + free(capture); return 0; } cuda_model_drop_file_pages(offset + copied, n); @@ -2241,8 +3236,10 @@ static int cuda_model_copy_to_device_streamed( what ? what : "expert", cudaGetErrorString(err)); (void)cudaGetLastError(); + free(capture); return 0; } + if (capture) cuda_host_cache_insert_owned(offset, bytes, capture, 0, 0); return 1; } @@ -2761,10 +3758,10 @@ extern "C" void ds4_gpu_set_glm_model(bool enabled) { extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { g_ssd_streaming_mode = enabled ? 1 : 0; - g_stream_expert_runtime_cap = 0; - g_stream_expert_runtime_gate_bytes = 0; - g_stream_expert_runtime_down_bytes = 0; - g_stream_expert_memory_cap_notice = 0; + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + g_stream_expert_runtime_caps[i] = 0; + g_stream_expert_memory_cap_notices[i] = 0; + } if (!g_ssd_streaming_mode) { cuda_stream_selected_cache_release(); cuda_stream_expert_cache_release_all(); @@ -2777,11 +3774,8 @@ extern "C" void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) { extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { g_stream_expert_budget_override = experts; - g_stream_expert_runtime_cap = 0; - g_stream_expert_runtime_gate_bytes = 0; - g_stream_expert_runtime_down_bytes = 0; - g_stream_expert_memory_cap_notice = 0; cuda_stream_selected_cache_invalidate(); + /* release_all also clears the classes, runtime caps, and notices. */ cuda_stream_expert_cache_release_all(); } @@ -2799,7 +3793,11 @@ extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { } extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) { - return g_stream_expert_cache.count; + uint32_t count = 0; + for (int i = 0; i < DS4_CUDA_STREAM_EXPERT_CLASSES; i++) { + count += g_stream_expert_caches[i].count; + } + return count; } extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { @@ -2817,8 +3815,10 @@ extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( down_expert_bytes) == 0) { return 0; } - cuda_stream_expert_cache_note_size(gate_expert_bytes, down_expert_bytes); - return cuda_stream_expert_cache_configured_budget(); + const int class_idx = + cuda_stream_expert_cache_class_index(gate_expert_bytes, + down_expert_bytes); + return cuda_stream_expert_cache_configured_budget_class(class_idx); } extern "C" int ds4_gpu_stream_expert_cache_seed_selected( @@ -2881,6 +3881,29 @@ extern "C" int ds4_gpu_stream_expert_cache_seed_selected( return 1; } +/* Allocate (grow-only) the staging buffers for the selected experts. */ +static int cuda_stream_selected_cache_ensure_slots( + uint64_t compact_gate_bytes, + uint64_t compact_down_bytes, + uint32_t slot_count) { + return cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.gate_ptr, + &g_stream_selected_cache.gate_capacity, + compact_gate_bytes, + "selected gate experts") && + cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.up_ptr, + &g_stream_selected_cache.up_capacity, + compact_gate_bytes, + "selected up experts") && + cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.down_ptr, + &g_stream_selected_cache.down_capacity, + compact_down_bytes, + "selected down experts") && + cuda_stream_selected_ensure_i32(&g_stream_selected_cache.slot_selected_ptr, + &g_stream_selected_cache.slot_selected_capacity, + slot_count, + "selected expert slots"); +} + static int cuda_stream_selected_cache_begin_compact_load( const void *model_map, uint64_t model_size, @@ -2896,7 +3919,8 @@ static int cuda_stream_selected_cache_begin_compact_load( uint64_t gate_expert_bytes, uint64_t down_expert_bytes, int strict_failure, - int allow_global_cache) { + int allow_global_cache, + int allow_cache_evict) { cuda_stream_selected_cache_invalidate(); cuda_model_load_progress_finish(); @@ -2928,35 +3952,36 @@ static int cuda_stream_selected_cache_begin_compact_load( return 0; } - if (!allow_global_cache) { + /* + * Prefill batch loads (allow_global_cache=0) used to ALWAYS release the + * whole resident expert cache before allocating the staging buffers: + * long prompts need the VRAM, but for short (chat) prompts the + * unconditional release throws away GiBs of warm cache on EVERY request + * and forces a full decode re-warm. Strategy: try the allocation first; + * only if VRAM is really short, sacrifice the cache and retry once. + */ + int selected_ok = + cuda_stream_selected_cache_ensure_slots(compact_gate_bytes, + compact_down_bytes, + slot_count); + if (!selected_ok) { + /* VRAM davvero corta (prompt molto lunghi): sacrifica la cache + * esperti e riprova una volta. */ cuda_stream_expert_cache_release_all(); + selected_ok = + cuda_stream_selected_cache_ensure_slots(compact_gate_bytes, + compact_down_bytes, + slot_count); } - - if (!cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.gate_ptr, - &g_stream_selected_cache.gate_capacity, - compact_gate_bytes, - "selected gate experts") || - !cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.up_ptr, - &g_stream_selected_cache.up_capacity, - compact_gate_bytes, - "selected up experts") || - !cuda_stream_selected_ensure_bytes(&g_stream_selected_cache.down_ptr, - &g_stream_selected_cache.down_capacity, - compact_down_bytes, - "selected down experts") || - !cuda_stream_selected_ensure_i32(&g_stream_selected_cache.slot_selected_ptr, - &g_stream_selected_cache.slot_selected_capacity, - slot_count, - "selected expert slots")) { + if (!selected_ok) { return strict_failure ? 0 : 1; } - if (allow_global_cache) { - cuda_stream_expert_cache_note_size(gate_expert_bytes, - down_expert_bytes); - } - const uint32_t configured_cache_budget = - cuda_stream_expert_cache_configured_budget(); + const uint32_t configured_cache_budget = allow_global_cache ? + cuda_stream_expert_cache_configured_budget_class( + cuda_stream_expert_cache_class_index(gate_expert_bytes, + down_expert_bytes)) : + 0; const int use_global_cache = allow_global_cache && configured_cache_budget != 0; @@ -2972,6 +3997,30 @@ static int cuda_stream_selected_cache_begin_compact_load( uint32_t cache_misses = 0; uint32_t direct_loads = 0; + cuda_fetch_job *fetch_jobs = NULL; + uint32_t fetch_njobs = 0; + const int parallel_fetch = + cuda_parallel_fetch_threads() > 0 && + g_model_fd >= 0 && + (g_model_fd_host_base == NULL || model_map == g_model_fd_host_base); + if (parallel_fetch && compact_count != 0) { + fetch_jobs = (cuda_fetch_job *)calloc((size_t)compact_count * 3u, + sizeof(cuda_fetch_job)); + /* NULL is fine: falls back to the serial path below */ + } + /* Global-cache misses fetched by the pool, to mirror into cache slots + * after the batch lands: {compact index, expert id} pairs. */ + uint32_t *backfill = NULL; + uint32_t backfill_count = 0; + if (fetch_jobs) { + backfill = (uint32_t *)calloc((size_t)compact_count * 2u, + sizeof(uint32_t)); + } + cuda_bundle_init(); + const cuda_bundle_layer *bundle = fetch_jobs != NULL ? + cuda_bundle_layer_get(layer, gate_expert_bytes, down_expert_bytes) : + NULL; + for (uint32_t i = 0; i < compact_count; i++) { if (compact_ids[i] < 0 || (uint32_t)compact_ids[i] >= n_total_expert) { fprintf(stderr, @@ -2979,6 +4028,8 @@ static int cuda_stream_selected_cache_begin_compact_load( compact_ids[i], n_total_expert, layer); + free(fetch_jobs); + free(backfill); return 0; } @@ -3004,31 +4055,50 @@ static int cuda_stream_selected_cache_begin_compact_load( cache_hits++; expert_cache->slots[(uint32_t)cache_slot].age = ++expert_cache->tick; + } else if (backfill) { + /* Parallel fetch active: read the miss straight into the + * compact buffers via the thread pool (which also consults + * the host expert cache) and mirror it into a cache slot + * afterwards. The serial load_slot mmap path caps effective + * disk throughput at SATA-class rates. */ + cache_misses++; + backfill[backfill_count * 2u] = i; + backfill[backfill_count * 2u + 1u] = (uint32_t)expert; + backfill_count++; } else { cache_misses++; - const uint32_t load_slot = - cuda_stream_expert_cache_lru_slot(expert_cache); - const int append = !expert_cache->slots[load_slot].valid; - if (cuda_stream_expert_cache_load_slot(expert_cache, - model_map, - model_size, - load_slot, - layer, - n_total_expert, - (uint32_t)expert, - gate_offset, - up_offset, - down_offset, - gate_expert_bytes, - down_expert_bytes)) { - if (append && expert_cache->count < expert_cache->capacity) { - expert_cache->count++; + /* + * Prefill batch loads (allow_cache_evict=0) may append while + * free capacity remains but must NOT evict valid slots: a + * long prompt would cycle the whole LRU cache and destroy + * the decode's warm working set. + */ + if (allow_cache_evict || + expert_cache->count < expert_cache->capacity) { + const uint32_t load_slot = + cuda_stream_expert_cache_lru_slot(expert_cache); + const int append = !expert_cache->slots[load_slot].valid; + if (cuda_stream_expert_cache_load_slot(expert_cache, + model_map, + model_size, + load_slot, + layer, + n_total_expert, + (uint32_t)expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes)) { + if (append && expert_cache->count < expert_cache->capacity) { + expert_cache->count++; + } + cache_slot = (int)load_slot; + } else { + cuda_stream_expert_cache_invalidate(); + expert_cache_disabled = 1; + cache_slot = -1; } - cache_slot = (int)load_slot; - } else { - cuda_stream_expert_cache_invalidate(); - expert_cache_disabled = 1; - cache_slot = -1; } } @@ -3053,7 +4123,35 @@ static int cuda_stream_selected_cache_begin_compact_load( const uint64_t up_src = up_offset + expert * gate_expert_bytes; const uint64_t down_src = down_offset + expert * down_expert_bytes; direct_loads++; - if (!cuda_model_copy_to_device_streamed(g_stream_selected_cache.gate_ptr + gate_dst, + if (fetch_jobs && bundle && expert < bundle->n_experts) { + cuda_fetch_job *j = &fetch_jobs[fetch_njobs]; + const uint64_t per = + 2 * bundle->gate_bytes + bundle->down_bytes; + j[0].dst = g_stream_selected_cache.gate_ptr + gate_dst; + j[0].dst_up = g_stream_selected_cache.up_ptr + gate_dst; + j[0].dst_down = g_stream_selected_cache.down_ptr + down_dst; + j[0].part_gate = bundle->gate_bytes; + j[0].part_down = bundle->down_bytes; + j[0].offset = bundle->base + expert * per; + j[0].bytes = per; + j[0].fd = g_bundle_fd; + fetch_njobs += 1; + } else if (fetch_jobs) { + cuda_fetch_job *j = &fetch_jobs[fetch_njobs]; + j[0].dst = g_stream_selected_cache.gate_ptr + gate_dst; + j[0].offset = gate_src; + j[0].bytes = gate_expert_bytes; + j[0].fd = g_model_fd; + j[1].dst = g_stream_selected_cache.up_ptr + gate_dst; + j[1].offset = up_src; + j[1].bytes = gate_expert_bytes; + j[1].fd = g_model_fd; + j[2].dst = g_stream_selected_cache.down_ptr + down_dst; + j[2].offset = down_src; + j[2].bytes = down_expert_bytes; + j[2].fd = g_model_fd; + fetch_njobs += 3; + } else if (!cuda_model_copy_to_device_streamed(g_stream_selected_cache.gate_ptr + gate_dst, model_map, model_size, gate_src, @@ -3077,6 +4175,53 @@ static int cuda_stream_selected_cache_begin_compact_load( } } + if (fetch_jobs) { + const int fetch_ok = cuda_fetch_execute(fetch_jobs, fetch_njobs); + free(fetch_jobs); + if (!fetch_ok) { + free(backfill); + cuda_stream_selected_cache_invalidate(); + return strict_failure ? 0 : 1; + } + } + if (backfill && backfill_count != 0 && !expert_cache_disabled && + expert_cache && expert_cache->valid) { + for (uint32_t b = 0; b < backfill_count; b++) { + const uint32_t compact_slot = backfill[b * 2u]; + const uint32_t expert_id = backfill[b * 2u + 1u]; + const uint32_t slot = + cuda_stream_expert_cache_lru_slot(expert_cache); + const int append = !expert_cache->slots[slot].valid; + if (!cuda_stream_expert_cache_backfill_from_compact( + expert_cache, + slot, + compact_slot, + g_stream_selected_cache.gate_ptr, + g_stream_selected_cache.up_ptr, + g_stream_selected_cache.down_ptr)) { + cuda_stream_expert_cache_invalidate(); + break; + } + cuda_stream_expert_cache_slot &entry = expert_cache->slots[slot]; + entry.valid = 1; + entry.model_map = model_map; + entry.model_size = model_size; + entry.layer = layer; + entry.n_total_expert = n_total_expert; + entry.expert = expert_id; + entry.gate_offset = gate_offset; + entry.up_offset = up_offset; + entry.down_offset = down_offset; + entry.gate_expert_bytes = gate_expert_bytes; + entry.down_expert_bytes = down_expert_bytes; + entry.age = ++expert_cache->tick; + if (append && expert_cache->count < expert_cache->capacity) { + expert_cache->count++; + } + } + } + free(backfill); + if (!cuda_ok(cudaMemcpy(g_stream_selected_cache.slot_selected_ptr, slot_ids, (size_t)slot_count * sizeof(slot_ids[0]), @@ -3123,6 +4268,96 @@ static int cuda_stream_selected_cache_begin_compact_load( return 1; } +/* Speculative prefetch: read predicted next-layer experts into the host + * expert cache so the coming selected load turns misses into RAM hits. + * Predictions come from running layer N+1's router on layer N's hidden + * state (cross-layer gate approximation); mispredicts only cost idle disk + * bandwidth. Fire-and-forget: completions are reaped opportunistically at + * the next prefetch call or demand fetch. io_uring builds only; a no-op + * elsewhere. Keys use raw model offsets, so runs using an expert bundle + * sidecar get no benefit (harmless). */ +extern "C" int ds4_gpu_stream_expert_cache_prefetch( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_selected) { +#if defined(__linux__) && defined(DS4_USE_IO_URING) + if (!g_ssd_streaming_mode || !table || !selected || + n_selected == 0 || n_selected > 8u) return 1; + if (g_model_fd < 0) return 1; + if (!cuda_fetch_uring_ready()) return 1; + cuda_host_cache_init(); + if (g_host_cache_nslots == 0) return 1; /* nowhere for reads to land */ + cuda_fetch_uring_reap_nowait(); + int32_t ids[8]; + if (!ds4_gpu_tensor_read(selected, 0, ids, + (uint64_t)n_selected * sizeof(ids[0]))) { + return 1; + } + unsigned queued = 0; + for (uint32_t i = 0; i < n_selected; i++) { + if (ids[i] < 0 || (uint32_t)ids[i] >= table->n_total_expert) continue; + const uint64_t expert = (uint64_t)(uint32_t)ids[i]; + const uint64_t offs[3] = { + table->gate_offset + expert * table->gate_expert_bytes, + table->up_offset + expert * table->gate_expert_bytes, + table->down_offset + expert * table->down_expert_bytes, + }; + const uint64_t lens[3] = { + table->gate_expert_bytes, + table->gate_expert_bytes, + table->down_expert_bytes, + }; + for (int p = 0; p < 3; p++) { + if (lens[p] == 0 || offs[p] > table->model_size || + lens[p] > table->model_size - offs[p]) { + continue; + } + if (cuda_host_cache_contains(offs[p], lens[p])) continue; + int dup = 0; + int free_slot = -1; + for (int s = 0; s < CUDA_PREFETCH_SLOTS; s++) { + if (g_prefetch_busy[s]) { + if (g_prefetch_jobs[s].offset == offs[p]) { dup = 1; break; } + } else if (free_slot < 0) { + free_slot = s; + } + } + if (dup || free_slot < 0) continue; + cuda_fetch_job *job = &g_prefetch_jobs[free_slot]; + memset(job, 0, sizeof(*job)); + job->kind = 1; + job->fd = g_model_fd; + job->offset = offs[p]; + job->bytes = lens[p]; + job->need_disk = 1; + if (!cuda_fetch_job_prepare_read(job)) continue; + if (!cuda_uring_submit_job(job)) { + free(job->host_buf); + job->host_buf = NULL; + continue; + } + g_prefetch_busy[free_slot] = 1; + g_prefetch_submitted++; + queued++; + } + } + if (queued != 0) (void)io_uring_submit(&g_fetch_uring); + if (queued != 0 && getenv("DS4_CUDA_PREFETCH_VERBOSE")) { + fprintf(stderr, + "ds4: expert prefetch layer=%u queued=%u submitted=%llu landed=%llu\n", + table->layer, + queued, + (unsigned long long)g_prefetch_submitted, + (unsigned long long)g_prefetch_completed); + } +#else + (void)table; + (void)selected; + (void)n_selected; +#endif + return 1; +} + extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load( const ds4_gpu_stream_expert_table *table, const int32_t *selected_ids, @@ -3179,6 +4414,7 @@ extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load( gate_expert_bytes, down_expert_bytes, 0, + 1, 1); } @@ -3249,6 +4485,13 @@ extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch( gate_expert_bytes, down_expert_bytes, 1, + /* + * The prefill batch also reads the global expert cache: hits + * become device-to-device copies instead of host->device uploads + * (short prompts drop from tens of seconds to a few). No evict + * though: see allow_cache_evict in the loop. + */ + 1, 0); } @@ -10859,15 +12102,23 @@ __global__ static void moe_gate_up_mid_expert_tile8_row32_kernel( slot[np] = pair[np] - tok[np] * n_expert; xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; } + /* The IQ2 dequant LUTs are consumed unconditionally by the dot calls + * below, so they must load for every xq_blocks; only the xq staging is + * <=16-blocks specific. GLM's 7168 embd gives 28 blocks, and gating the + * LUT loads on the staging condition left s_iq2_grid/s_iq2_signs + * uninitialized there: the whole batch gate/up dequant ran on garbage + * shared memory (fast wrong prefill, corrupt MTP verify logits). */ + 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]; 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 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(); + } + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } if (row >= expert_mid_dim) return; @@ -10949,15 +12200,23 @@ __global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel( slot[np] = pair[np] - tok[np] * n_expert; xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; } + /* The IQ2 dequant LUTs are consumed unconditionally by the dot calls + * below, so they must load for every xq_blocks; only the xq staging is + * <=16-blocks specific. GLM's 7168 embd gives 28 blocks, and gating the + * LUT loads on the staging condition left s_iq2_grid/s_iq2_signs + * uninitialized there: the whole batch gate/up dequant ran on garbage + * shared memory (fast wrong prefill, corrupt MTP verify logits). */ + 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]; 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 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(); + } + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } for (uint32_t rr = 0; rr < 64u; rr++) { @@ -11043,15 +12302,23 @@ __global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel( slot[np] = pair[np] - tok[np] * n_expert; xqb[np] = xq + (uint64_t)tok[np] * xq_blocks; } + /* The IQ2 dequant LUTs are consumed unconditionally by the dot calls + * below, so they must load for every xq_blocks; only the xq staging is + * <=16-blocks specific. GLM's 7168 embd gives 28 blocks, and gating the + * LUT loads on the staging condition left s_iq2_grid/s_iq2_signs + * uninitialized there: the whole batch gate/up dequant ran on garbage + * shared memory (fast wrong prefill, corrupt MTP verify logits). */ + 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]; 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 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(); + } + __syncthreads(); + if (xq_blocks <= 16u) { for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p]; } for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) { @@ -11677,6 +12944,61 @@ __global__ static void moe_down_sorted_qwarp32_kernel( if (lane == 0) down_out[(uint64_t)pair * out_dim + row] = acc; } +/* IQ2_XXS twins of moe_down_qwarp32 / moe_down_sorted_qwarp32 for quants + * that keep the routed down projection in IQ2_XXS (GLM UD RoutedIQ2XXS). */ +__global__ static void moe_down_iq2xxs_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_iq2_xxs *wr = (const cuda_block_iq2_xxs *)(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_iq2_xxs_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_iq2xxs_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_iq2_xxs *wr = (const cuda_block_iq2_xxs *)(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_iq2_xxs_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, @@ -12276,7 +13598,8 @@ static int routed_moe_launch( float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, - uint32_t n_tokens) { + uint32_t n_tokens, + bool force_resident) { 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 || @@ -12292,7 +13615,10 @@ static int routed_moe_launch( return 0; } const int q4k_path = (gate_type == 12u && down_type == 12u); - if (!q4k_path && (gate_type != 16u || down_type != 10u)) return 0; + /* IQ2_XXS gate pairs with a Q2_K down (DeepSeek quants) or an IQ2_XXS + * down (GLM UD RoutedIQ2XXS quants). */ + const int iq2xxs_down = (!q4k_path && down_type == 16u); + if (!q4k_path && (gate_type != 16u || (down_type != 10u && down_type != 16u))) return 0; 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 || @@ -12302,6 +13628,7 @@ static int routed_moe_launch( } const uint64_t required_slot_count = (uint64_t)n_tokens * n_expert; const int use_stream_selected_cache = + !force_resident && g_ssd_streaming_mode && g_stream_selected_cache.valid && g_stream_selected_cache.model_map == model_map && @@ -12363,7 +13690,7 @@ static int routed_moe_launch( const uint32_t expert_tile_m = (!q4k_path && getenv("DS4_CUDA_MOE_TILE4")) ? 4u : 8u; const uint32_t write_gate_up = getenv("DS4_CUDA_MOE_WRITE_GATE_UP") != NULL; const uint32_t use_p2_sorted = use_sorted_pairs && !q4k_path && getenv("DS4_CUDA_MOE_NO_P2") == NULL; - const uint32_t use_atomic_down = use_expert_tiles && + const uint32_t use_atomic_down = use_expert_tiles && !iq2xxs_down && getenv("DS4_CUDA_MOE_NO_ATOMIC_DOWN") == NULL && (getenv("DS4_CUDA_MOE_ATOMIC_DOWN") != NULL || (!q4k_path && n_tokens >= 128u)); @@ -12398,7 +13725,7 @@ static int routed_moe_launch( getenv("DS4_CUDA_MOE_NO_DOWN_ROW128") == NULL && getenv("DS4_CUDA_MOE_NO_DOWN_ROW64") == NULL)); const uint32_t use_direct_down_sum6 = - n_tokens == 1u && n_expert == 6u && + n_tokens == 1u && n_expert == 6u && !iq2xxs_down && getenv("DS4_CUDA_MOE_NO_DIRECT_DOWN_SUM6") == NULL; uint32_t *sorted_pairs = NULL; uint32_t *sorted_offsets = NULL; @@ -12724,7 +14051,8 @@ static int routed_moe_launch( } if (use_direct_down_sum6) { /* The direct decode kernel writes the final token row. */ - } else if (sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && + } else if (!iq2xxs_down && + sorted_pairs && use_expert_tiles && sorted_offsets && sorted_counts && down_tile_total && down_tile_experts && down_tile_starts) { if (q4k_path) { if (use_down_row2048) { @@ -12834,7 +14162,7 @@ static int routed_moe_launch( 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) { + } else if (!iq2xxs_down && sorted_pairs && use_p2_sorted) { dim3 p2_dgrid((out_dim + 15u) / 16u, (pair_count + 1u) / 2u, 1); moe_down_sorted_p2_qwarp32_kernel<<>>( (float *)down->ptr, @@ -12849,17 +14177,31 @@ static int routed_moe_launch( n_expert, pair_count); } else if (!q4k_path && sorted_pairs) { - moe_down_sorted_qwarp32_kernel<<>>( - (float *)down->ptr, - down_w, - midq, - sorted_pairs, - selected_ptr, - down_expert_bytes, - down_row_bytes, - midq_blocks, - out_dim, - n_expert); + if (iq2xxs_down) { + moe_down_iq2xxs_sorted_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + sorted_pairs, + selected_ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } else { + moe_down_sorted_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + sorted_pairs, + selected_ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); + } } else { if (q4k_path) { moe_down_q4K_qwarp32_kernel<<>>( @@ -12872,6 +14214,17 @@ static int routed_moe_launch( midq_blocks, out_dim, n_expert); + } else if (iq2xxs_down) { + moe_down_iq2xxs_qwarp32_kernel<<>>( + (float *)down->ptr, + down_w, + midq, + selected_ptr, + down_expert_bytes, + down_row_bytes, + midq_blocks, + out_dim, + n_expert); } else { moe_down_qwarp32_kernel<<>>( (float *)down->ptr, @@ -12960,7 +14313,7 @@ extern "C" int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, return 1; } -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, uint32_t layer_index) { +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, uint32_t layer_index, bool force_resident) { return routed_moe_launch(out, gate, up, mid, down, model_map, model_size, gate_offset, up_offset, down_offset, gate_type, down_type, @@ -12968,7 +14321,7 @@ extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor 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); + layer_index, 1, force_resident); } 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) { if (mid_is_f16) *mid_is_f16 = false; @@ -12979,7 +14332,7 @@ extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tens 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); + layer_index, n_tokens, false); } 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; @@ -13297,3 +14650,107 @@ extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor( ds4_gpu_hc_expand_split_tensor(out_hc, block_out, residual_hc, split, n_embd, n_hc); } + +/* --- GLM CUDA port: q8_0 token embedding lookups ------------------------- */ + +__global__ static void embed_token_q8_0_kernel( + float *out, + const unsigned char *w, /* q8_0 row base for the token */ + uint32_t n_embd) { + uint32_t e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= n_embd) return; + const uint32_t blk = e >> 5; /* 32 elements per q8_0 block */ + const uint32_t idx = e & 31u; + const unsigned char *b = w + (uint64_t)blk * 34u; + const float d = __half2float(*reinterpret_cast(b)); + out[e] = d * (float)((const signed char *)(b + 2))[idx]; +} + +__global__ static void embed_tokens_q8_0_kernel( + float *out, + const int32_t *tokens, + const unsigned char *w, /* q8_0 embedding table base */ + uint32_t n_vocab, + uint32_t n_tokens, + uint32_t n_embd, + uint64_t row_bytes) { + uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t n = (uint64_t)n_tokens * n_embd; + if (gid >= n) return; + const uint32_t e = (uint32_t)(gid % n_embd); + const uint32_t t = (uint32_t)(gid / n_embd); + int32_t tok_i = tokens[t]; + uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i; + if (tok >= n_vocab) tok = 0; + const unsigned char *row = w + (uint64_t)tok * row_bytes; + const uint32_t blk = e >> 5; + const uint32_t idx = e & 31u; + const unsigned char *b = row + (uint64_t)blk * 34u; + const float d = __half2float(*reinterpret_cast(b)); + out[gid] = d * (float)((const signed char *)(b + 2))[idx]; +} + +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) { + if (!out || !model_map || n_embd == 0 || (n_embd & 31u) != 0) return 0; + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + const uint64_t table_bytes = (uint64_t)n_vocab * row_bytes; + if (weight_offset > model_size || table_bytes > model_size - weight_offset || + out->bytes < (uint64_t)n_embd * sizeof(float)) { + return 0; + } + if (token >= n_vocab) token = 0; + const char *wptr = cuda_model_range_ptr(model_map, + weight_offset + (uint64_t)token * row_bytes, + row_bytes, + "glm token_embd row"); + if (!wptr) return 0; + embed_token_q8_0_kernel<<<(n_embd + 255) / 256, 256>>>( + (float *)out->ptr, (const unsigned char *)wptr, n_embd); + return cuda_ok(cudaGetLastError(), "glm embed token q8_0 launch"); +} + +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) { + if (!out || !tokens || !model_map || n_embd == 0 || (n_embd & 31u) != 0 || + n_tokens == 0) { + return 0; + } + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + const uint64_t table_bytes = (uint64_t)n_vocab * row_bytes; + if (weight_offset > model_size || table_bytes > model_size - weight_offset || + tokens->bytes < (uint64_t)n_tokens * sizeof(int32_t) || + out->bytes < (uint64_t)n_tokens * n_embd * sizeof(float)) { + return 0; + } + const char *wptr = cuda_model_range_ptr(model_map, weight_offset, + table_bytes, "glm token_embd"); + if (!wptr) return 0; + const uint64_t n = (uint64_t)n_tokens * n_embd; + embed_tokens_q8_0_kernel<<<(uint32_t)((n + 255) / 256), 256>>>( + (float *)out->ptr, + (const int32_t *)tokens->ptr, + (const unsigned char *)wptr, + n_vocab, n_tokens, n_embd, row_bytes); + return cuda_ok(cudaGetLastError(), "glm embed tokens q8_0 launch"); +} + +#include "ds4_cuda_glm_kv.inc" +#include "ds4_cuda_glm_indexer.inc" +#include "ds4_cuda_glm_attn.inc" +#include "ds4_cuda_glm_moe.inc" +#include "ds4_cuda_glm_stubs.inc" +#include "ds4_cuda_gqa.inc" diff --git a/ds4_cuda_glm_attn.inc b/ds4_cuda_glm_attn.inc new file mode 100644 index 000000000..ba1d144b4 --- /dev/null +++ b/ds4_cuda_glm_attn.inc @@ -0,0 +1,1415 @@ +/* ds4_cuda_glm_attn.inc - CUDA port of the GLM-5.2 attention core. + * + * Included near the end of ds4_cuda.cu: relies on cuda_ok(), + * cuda_model_range_ptr(), struct ds4_gpu_tensor {ptr,bytes} and the standard + * CUDA headers that file already includes. + * + * Ported from the Metal launchers in ds4_metal.m and the shaders in + * metal/dsv4_misc.metal (kernel_glm_attention_indexed_decode, + * kernel_glm_attention_indexed_decode_split_group8_*, + * kernel_glm_attention_indexed_batch, the *_lora_group8_vec[_causal] family, + * kernel_glm_attention_full and kernel_glm_fill_selected_range[_batch]). + * + * ponytail: the Metal backend ships several tuned variants per entry point + * (group2/group8/q2_group4/vec/fullheads). Each CUDA entry point below is a + * single kernel implementing the same math; add tuned variants only if + * profiling demands it. + */ + +#include + +/* ------------------------------------------------------------------------ + * Device helpers (glm_attn_ prefix to stay clash-free inside ds4_cuda.cu). + * ------------------------------------------------------------------------ */ + +__device__ static float glm_attn_cache_load( + const char *base, uint64_t index, uint32_t cache_f16) { + if (cache_f16) return __half2float(((const __half *)base)[index]); + return ((const float *)base)[index]; +} + +/* Port of glm_rope_yarn_corr_dims() (metal/dsv4_misc.metal:514). */ +__device__ static void glm_attn_rope_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, + float beta_fast, float beta_slow, float *c0, float *c1) { + const float denom = 2.0f * logf(freq_base); + const float f_fast = (float)n_dims * + logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom; + const float f_slow = (float)n_dims * + logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom; + *c0 = fmaxf(0.0f, floorf(f_fast)); + *c1 = fminf((float)n_dims - 1.0f, ceilf(f_slow)); +} + +/* Port of glm_rope_yarn() (metal/dsv4_misc.metal:490). */ +__device__ static void glm_attn_rope_yarn( + float theta_extrap, float freq_scale, float c0, float c1, int i0, + float ext_factor, float mscale, float *cos_theta, float *sin_theta) { + const float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + const float y = ((float)(i0 / 2) - c0) / fmaxf(0.001f, c1 - c0); + const float ramp_mix = + (1.0f - fminf(1.0f, fmaxf(0.0f, y))) * 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; +} + +/* Port of glm_cache_load_rotated_rope_pair() (metal/dsv4_misc.metal:1455): + * loads K-rope cache pair (r, r+1) of `row` and rotates it by the row + * position at attention time. */ +__device__ static float2 glm_attn_rope_pair( + const char *base, uint64_t rope_base, uint32_t r, uint32_t row, + uint32_t qk_rope, uint32_t cache_f16, float freq_base, + float freq_scale, float ext_factor, float attn_factor, + float c0, float c1) { + const float inv_ndims = -1.0f / (float)qk_rope; + const float theta = + (float)row * powf(freq_base, inv_ndims * (float)r); + float cos_theta, sin_theta; + glm_attn_rope_yarn(theta, freq_scale, c0, c1, (int)r, + ext_factor, attn_factor, &cos_theta, &sin_theta); + const float x0 = glm_attn_cache_load(base, rope_base + r, cache_f16); + const float x1 = glm_attn_cache_load(base, rope_base + r + 1u, cache_f16); + return make_float2(x0 * cos_theta - x1 * sin_theta, + x0 * sin_theta + x1 * cos_theta); +} + +/* Port of glm_q8_0_dot_row_tg_f32() (metal/dsv4_misc.metal:1539). `x` may + * point to shared or global memory. */ +__device__ static float glm_attn_q8_dot_row( + const char *row, const float *x, uint32_t n_cols) { + float acc = 0.0f; + const uint32_t n_blocks = (n_cols + 31u) >> 5; + for (uint32_t block = 0; block < n_blocks; block++) { + const char *block_base = row + (uint64_t)block * 34u; + const float d = __half2float(*(const __half *)block_base); + const int8_t *qs = (const int8_t *)(block_base + 2u); + const uint32_t base = block << 5; + const uint32_t count = min(32u, n_cols - base); + for (uint32_t qi = 0; qi < count; qi++) { + acc += d * (float)qs[qi] * x[base + qi]; + } + } + return acc; +} + +/* Dynamic shared memory above the 48KB static limit needs an opt-in. */ +static int glm_attn_smem_ok(const void *func, uint64_t smem_bytes, + const char *what) { + if (smem_bytes <= 48u * 1024u) return 1; + if (smem_bytes > (uint64_t)INT_MAX) return 0; + return cuda_ok(cudaFuncSetAttribute(func, + cudaFuncAttributeMaxDynamicSharedMemorySize, + (int)smem_bytes), + what); +} + +/* ------------------------------------------------------------------------ + * Selection-buffer fills (kernel_glm_fill_selected_range[_batch]). + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_fill_selected_range_kernel( + uint32_t *selected, uint32_t n_selected) { + const uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + if (gid < n_selected) selected[gid] = gid; +} + +__global__ static void glm_attn_fill_selected_range_batch_kernel( + uint32_t *selected, uint32_t n_tokens, uint32_t pos0, + uint32_t n_selected, uint32_t pad_row) { + const 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_tensor( + ds4_gpu_tensor *selected, + uint32_t n_selected) { + if (!selected || n_selected == 0) return 0; + if (selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { + fprintf(stderr, "ds4: CUDA GLM selected range received undersized buffer\n"); + return 0; + } + glm_attn_fill_selected_range_kernel<<<(n_selected + 255u) / 256u, 256>>>( + (uint32_t *)selected->ptr, n_selected); + return cuda_ok(cudaGetLastError(), "glm fill selected range launch"); +} + +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) return 0; + const uint64_t total = (uint64_t)n_tokens * n_selected; + if (total / n_tokens != n_selected) return 0; + if (total > UINT64_MAX / sizeof(uint32_t)) return 0; + if (total > UINT32_MAX) return 0; + if (selected->bytes < total * sizeof(uint32_t)) { + fprintf(stderr, "ds4: CUDA GLM selected range batch received undersized buffer\n"); + return 0; + } + glm_attn_fill_selected_range_batch_kernel<<< + (uint32_t)((total + 255u) / 256u), 256>>>( + (uint32_t *)selected->ptr, n_tokens, pos0, n_selected, pad_row); + return cuda_ok(cudaGetLastError(), "glm fill selected range batch launch"); +} + +/* ------------------------------------------------------------------------ + * Dense fallback attention (kernel_glm_attention_full, mode 2, which is the + * only mode the Metal launcher dispatches). + * + * Grid (n_tokens, n_head), 256 threads. Dynamic shared memory holds the + * reduction scratch (256 floats) followed by one score per visible row. + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_full_kernel( + float *heads, + const float *q, + const char *key_cache, + const char *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + uint32_t cache_f16, + float scale) { + const uint32_t token = blockIdx.x; + const uint32_t head = blockIdx.y; + if (token >= n_tokens || head >= n_head) return; + + const uint32_t nth = blockDim.x; + const uint32_t tid = threadIdx.x; + const uint32_t visible = min(cache_len, pos0 + token + 1u); + extern __shared__ float glm_attn_full_sm[]; + float *red = glm_attn_full_sm; + float *scores = glm_attn_full_sm + 256u; + + const float *qh = q + ((uint64_t)token * n_head + head) * qk_dim; + + for (uint32_t s = tid; s < visible; s += nth) { + const uint64_t kbase = ((uint64_t)s * n_head + head) * qk_dim; + float dotv = 0.0f; + for (uint32_t i = 0; i < qk_dim; i++) { + dotv += qh[i] * glm_attn_cache_load(key_cache, kbase + i, cache_f16); + } + scores[s] = dotv * scale; + } + __syncthreads(); + + /* Metal's mode-2 kernel runs the softmax normalization serially on one + * thread; keep that for exact-order parity. */ + if (tid == 0u) { + float max_score = -INFINITY; + for (uint32_t s = 0; s < visible; s++) { + max_score = fmaxf(max_score, scores[s]); + } + float sum = 0.0f; + for (uint32_t s = 0; s < visible; s++) { + const float w = expf(scores[s] - max_score); + scores[s] = w; + sum += w; + } + red[0] = fmaxf(sum, 1.0e-20f); + } + __syncthreads(); + + const float denom = red[0]; + float *out = heads + ((uint64_t)token * n_head + head) * value_dim; + for (uint32_t d = tid; d < value_dim; d += nth) { + float acc = 0.0f; + for (uint32_t s = 0; s < visible; s++) { + const uint64_t vbase = ((uint64_t)s * n_head + head) * value_dim; + acc += scores[s] * + glm_attn_cache_load(value_cache, vbase + d, cache_f16); + } + out[d] = acc / denom; + } +} + +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) { + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim == 0 || value_dim == 0 || + (qk_dim & 3u) != 0 || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || + key_cache->bytes < (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes || + value_cache->bytes < (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes) { + fprintf(stderr, "ds4: CUDA GLM attention received undersized buffers\n"); + return 0; + } + const uint64_t smem = (256u + (uint64_t)cache_len) * sizeof(float); + if (!glm_attn_smem_ok((const void *)glm_attn_full_kernel, smem, + "glm full attention smem opt-in")) { + return 0; + } + dim3 grid(n_tokens, n_head, 1); + glm_attn_full_kernel<<>>( + (float *)heads->ptr, + (const float *)q->ptr, + (const char *)key_cache->ptr, + (const char *)value_cache->ptr, + pos0, n_tokens, cache_len, n_head, qk_dim, value_dim, + cache_f16 ? 1u : 0u, + 1.0f / sqrtf((float)qk_dim)); + return cuda_ok(cudaGetLastError(), "glm full attention launch"); +} + +/* ------------------------------------------------------------------------ + * FlashAttention fallback. + * + * The Metal backend stages an f16 copy of the dense KV caches and runs the + * llama.cpp-style kernel_flash_attn_ext_f16_dk256_dv256 with a causal + * prefill mask. Mathematically that is dense causal attention over + * cache_len rows with scale 1/sqrt(qk_dim). Here it is computed with a + * chunked online softmax so no cache_len-sized score buffer is needed + * (hence no Metal-style 8192-row cap). KV is read straight from the dense + * caches instead of an f16 staging copy, which only changes rounding when + * the cache is f32. + * + * Grid (n_tokens, n_head), 256 threads; thread d owns output dim d, which + * requires value_dim == 256 (the launcher enforces qk_dim == value_dim == + * 256, same as Metal). + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_flash_kernel( + float *heads, + const float *q, + const char *key_cache, + const char *value_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_len, + uint32_t n_head, + uint32_t qk_dim, + uint32_t value_dim, + uint32_t cache_f16, + float scale) { + const uint32_t token = blockIdx.x; + const uint32_t head = blockIdx.y; + if (token >= n_tokens || head >= n_head) return; + + const uint32_t tid = threadIdx.x; + const uint32_t visible = min(cache_len, pos0 + token + 1u); + const float *qh = q + ((uint64_t)token * n_head + head) * qk_dim; + + __shared__ float w[256]; + __shared__ float red[256]; + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float acc = 0.0f; /* output accumulator for dim d == tid */ + + for (uint32_t base = 0; base < visible; base += 256u) { + const uint32_t cn = min(256u, visible - base); + + float score = -FLT_MAX / 2.0f; + if (tid < cn) { + const uint32_t s = base + tid; + const uint64_t kbase = ((uint64_t)s * n_head + head) * qk_dim; + float dotv = 0.0f; + for (uint32_t i = 0; i < qk_dim; i++) { + dotv += qh[i] * + glm_attn_cache_load(key_cache, kbase + i, cache_f16); + } + score = dotv * scale; + } + + red[tid] = score; + __syncthreads(); + for (uint32_t step = 128u; step > 0; step >>= 1) { + if (tid < step) red[tid] = fmaxf(red[tid], red[tid + step]); + __syncthreads(); + } + const float chunk_max = red[0]; + __syncthreads(); + + const float new_m = fmaxf(M, chunk_max); + const float wv = (tid < cn) ? expf(score - new_m) : 0.0f; + w[tid] = wv; + red[tid] = wv; + __syncthreads(); + for (uint32_t step = 128u; step > 0; step >>= 1) { + if (tid < step) red[tid] += red[tid + step]; + __syncthreads(); + } + const float chunk_sum = red[0]; + + const float old_scale = expf(M - new_m); + float a = 0.0f; + for (uint32_t r = 0; r < cn; r++) { + const uint64_t vbase = + ((uint64_t)(base + r) * n_head + head) * value_dim; + a += w[r] * glm_attn_cache_load(value_cache, vbase + tid, cache_f16); + } + acc = acc * old_scale + a; + S = S * old_scale + chunk_sum; + M = new_m; + __syncthreads(); + } + + heads[((uint64_t)token * n_head + head) * value_dim + tid] = + acc / fmaxf(S, 1.0e-20f); +} + +static int glm_attn_flash_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) { + if (!heads || !q || !key_cache || !value_cache || + n_tokens == 0 || cache_len == 0 || cache_cap == 0 || + n_head == 0 || qk_dim != 256u || value_dim != 256u || + cache_len > cache_cap || + pos0 > cache_len || n_tokens > cache_len - pos0) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float) || + key_cache->bytes < (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes || + value_cache->bytes < (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes) { + fprintf(stderr, "ds4: CUDA GLM FlashAttention received undersized buffers\n"); + return 0; + } + dim3 grid(n_tokens, n_head, 1); + glm_attn_flash_kernel<<>>( + (float *)heads->ptr, + (const float *)q->ptr, + (const char *)key_cache->ptr, + (const char *)value_cache->ptr, + pos0, n_tokens, cache_len, n_head, qk_dim, value_dim, + cache_f16 ? 1u : 0u, + 1.0f / sqrtf((float)qk_dim)); + return cuda_ok(cudaGetLastError(), "glm flash attention launch"); +} + +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_attn_flash_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_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) { + /* Metal's staged variant skips the KV->f16 staging copy because + * kernel_glm_build_kv_cache_flash already produced it. That builder + * also writes the dense key/value caches (see dsv4_misc.metal:1111), so + * on CUDA the staged entry can read the caches directly: same data, no + * staging buffer required. */ + if (pos0 != 0 || n_tokens != cache_len) return 0; + return glm_attn_flash_launch(heads, q, key_cache, value_cache, + pos0, n_tokens, cache_len, cache_cap, + n_head, qk_dim, value_dim, cache_f16); +} + +/* ------------------------------------------------------------------------ + * Indexed (DSA) attention over selected compact-KV rows, single fused + * kernel: scores + softmax + lora accumulation + q8_0 value projection. + * + * Port of kernel_glm_attention_indexed_batch (which is the decode kernel + * kernel_glm_attention_indexed_decode plus a token grid dimension; the + * decode entry launches it with n_tokens == 1). Handles f16 and f32 + * caches; the Metal f16 fast path is just a vectorized version of the same + * math. + * + * Grid (n_head, n_tokens), 256 threads. Dynamic shared memory: + * red[256] + scores[n_selected] + lora_sum[kv_lora_dim]. + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_indexed_kernel( + float *heads, + const float *q, + const float *qk_low, + const char *kv_lora_cache, + const char *k_rope_cache, + const char *value_weight, + const uint32_t *selected, + uint32_t n_tokens, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t 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 value_row_bytes, + float scale, + 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; + if (head >= n_head || token >= n_tokens || n_selected == 0u) return; + const uint32_t nth = blockDim.x; + const uint32_t tid = threadIdx.x; + const uint32_t qk_dim = qk_nope + qk_rope; + + extern __shared__ float glm_attn_indexed_sm[]; + float *red = glm_attn_indexed_sm; + float *scores = glm_attn_indexed_sm + 256u; + float *lora_sum = scores + n_selected; + + 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; + const uint32_t *token_selected = + selected + (uint64_t)token * n_selected; + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + glm_attn_rope_corr_dims((int)qk_rope, (int)n_ctx_orig, freq_base, + beta_fast, beta_slow, &corr0, &corr1); + } + + float local_max = -INFINITY; + for (uint32_t s = tid; s < n_selected; s += nth) { + const uint32_t row = token_selected[s]; + float score = -INFINITY; + if (row < cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * kv_lora_dim; + for (uint32_t j = 0; j < kv_lora_dim; j++) { + dotv += low[j] * + glm_attn_cache_load(kv_lora_cache, lora_base + j, + cache_f16); + } + const uint64_t rope_base = (uint64_t)row * qk_rope; + for (uint32_t r = 0; r < qk_rope; r += 2u) { + const float2 y = glm_attn_rope_pair(k_rope_cache, rope_base, r, + row, qk_rope, cache_f16, + freq_base, freq_scale, + ext_factor, attn_factor, + corr0, corr1); + 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]; + + float local_sum = 0.0f; + for (uint32_t s = tid; s < n_selected; 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(); + + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint32_t s = 0; s < n_selected; s++) { + const uint32_t row = token_selected[s]; + if (row < cache_cap) { + acc += scores[s] * + glm_attn_cache_load(kv_lora_cache, + (uint64_t)row * kv_lora_dim + j, + cache_f16); + } + } + lora_sum[j] = acc / denom; + } + __syncthreads(); + + float *out = + heads + ((uint64_t)token * n_head + head) * value_dim; + for (uint32_t d = tid; d < value_dim; d += nth) { + const char *row = + value_weight + ((uint64_t)head * value_dim + d) * value_row_bytes; + out[d] = glm_attn_q8_dot_row(row, lora_sum, kv_lora_dim); + } +} + +static int glm_attn_indexed_launch( + 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, + 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, + const char *label) { + const uint32_t qk_dim = qk_nope + qk_rope; + if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + !model_map || !selected || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || + n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + value_dim == 0 || qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t value_row_bytes = (uint64_t)((kv_lora_dim + 31u) / 32u) * 34u; + const uint64_t value_weight_bytes = + (uint64_t)n_head * value_dim * value_row_bytes; + if (heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float) || + 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_bytes || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem_bytes || + selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t)) { + fprintf(stderr, "ds4: CUDA GLM indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: CUDA GLM indexed attention value range is outside the mapped model\n"); + return 0; + } + const char *value_weight = cuda_model_range_ptr( + model_map, value_weight_offset, value_weight_bytes, + "glm attn v_b weights"); + if (!value_weight) return 0; + + const uint64_t smem = + (256u + (uint64_t)n_selected + kv_lora_dim) * sizeof(float); + if (!glm_attn_smem_ok((const void *)glm_attn_indexed_kernel, smem, + "glm indexed attention smem opt-in")) { + return 0; + } + dim3 grid(n_head, n_tokens, 1); + glm_attn_indexed_kernel<<>>( + (float *)heads->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const char *)kv_lora_cache->ptr, + (const char *)k_rope_cache->ptr, + value_weight, + (const uint32_t *)selected->ptr, + n_tokens, n_selected, cache_cap, + cache_f16 ? 1u : 0u, + n_head, kv_lora_dim, qk_nope, qk_rope, value_dim, n_ctx_orig, + (uint32_t)value_row_bytes, + 1.0f / sqrtf((float)qk_dim), + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), label); +} + +extern "C" int ds4_gpu_glm_attention_indexed_decode_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, + 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) { + return glm_attn_indexed_launch(heads, q, qk_low, kv_lora_cache, + k_rope_cache, model_map, model_size, + value_weight_offset, selected, + 1u /* n_tokens */, n_selected, cache_cap, + cache_f16, n_head, kv_lora_dim, qk_nope, + qk_rope, value_dim, n_ctx_orig, + freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, + "glm indexed attention decode launch"); +} + +extern "C" int ds4_gpu_glm_attention_indexed_batch_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, + 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) { + return glm_attn_indexed_launch(heads, q, qk_low, kv_lora_cache, + k_rope_cache, model_map, model_size, + value_weight_offset, selected, + n_tokens, n_selected, cache_cap, + cache_f16, n_head, kv_lora_dim, qk_nope, + qk_rope, value_dim, n_ctx_orig, + freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, + "glm indexed attention batch launch"); +} + +/* ------------------------------------------------------------------------ + * Split (two-phase) grouped indexed decode. + * + * Phase 1 (partial): one warp per (head, selected-row block) computes an + * online-softmax partial over its rows and stores the unnormalized lora + * accumulator plus its (max, sum) pair. Port of + * kernel_glm_attention_indexed_decode_split_group8_partial_impl with the + * same fixed shape (kv_lora_dim == 512, qk_rope == 64, f16 cache): + * lane l owns lora elements {c*128 + l*4 + k : c in 0..3, k in 0..3}. + * ponytail: Metal groups 8 heads per threadgroup to share a staged KV tile; + * here each warp reads global memory directly - same math, add staging only + * if profiling demands it. + * + * Phase 2 (reduce): port of + * kernel_glm_attention_indexed_decode_split_group8_reduce_impl - merges the + * per-block partials with softmax rescaling and applies the q8_0 value + * projection. + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_split_g8_partial_kernel( + float *partial_lora, + float *partial_ms, + const float *q, + const float *qk_low, + const char *kv_lora_cache, /* f16 */ + const char *k_rope_cache, /* f16 */ + const uint32_t *selected, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t n_head, + uint32_t qk_nope, + uint32_t n_ctx_orig, + uint32_t block_rows, + float scale, + 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 block = blockIdx.y; + const uint32_t lane = threadIdx.x; /* 32 threads per block */ + if (head >= n_head || n_selected == 0u || block_rows == 0u) return; + + const uint32_t kv_lora_dim = 512u; + const uint32_t qk_rope = 64u; + const uint32_t qk_dim = qk_nope + qk_rope; + const uint32_t block_start = block * block_rows; + const uint32_t block_end = min(n_selected, block_start + block_rows); + + const float *qh = q + (uint64_t)head * qk_dim; + const float *low = qk_low + (uint64_t)head * kv_lora_dim; + + float lowv[16]; +#pragma unroll + for (uint32_t c = 0; c < 4u; c++) { +#pragma unroll + for (uint32_t k = 0; k < 4u; k++) { + lowv[c * 4u + k] = low[c * 128u + lane * 4u + k]; + } + } + float qrope[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + if (lane < 16u) { +#pragma unroll + for (uint32_t k = 0; k < 4u; k++) { + qrope[k] = qh[qk_nope + lane * 4u + k]; + } + } + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + glm_attn_rope_corr_dims((int)qk_rope, (int)n_ctx_orig, freq_base, + beta_fast, beta_slow, &corr0, &corr1); + } + + float M = -FLT_MAX / 2.0f; + float S = 0.0f; + float o[16]; +#pragma unroll + for (uint32_t i = 0; i < 16u; i++) o[i] = 0.0f; + + for (uint32_t idx = block_start; idx < block_end; idx++) { + const uint32_t row = selected[idx]; + const bool valid_row = row < cache_cap; + float kvv[16]; + float part = 0.0f; + if (valid_row) { + const __half *kv_row = + (const __half *)kv_lora_cache + (uint64_t)row * kv_lora_dim; +#pragma unroll + for (uint32_t c = 0; c < 4u; c++) { +#pragma unroll + for (uint32_t k = 0; k < 4u; k++) { + const float v = + __half2float(kv_row[c * 128u + lane * 4u + k]); + kvv[c * 4u + k] = v; + part += lowv[c * 4u + k] * v; + } + } + if (lane < 16u) { + const uint64_t rope_base = (uint64_t)row * qk_rope; + const uint32_t r = lane * 4u; + const float2 y0 = glm_attn_rope_pair(k_rope_cache, rope_base, + r, row, qk_rope, 1u, + freq_base, freq_scale, + ext_factor, attn_factor, + corr0, corr1); + const float2 y1 = glm_attn_rope_pair(k_rope_cache, rope_base, + r + 2u, row, qk_rope, 1u, + freq_base, freq_scale, + ext_factor, attn_factor, + corr0, corr1); + part += qrope[0] * y0.x + qrope[1] * y0.y + + qrope[2] * y1.x + qrope[3] * y1.y; + } + } else { +#pragma unroll + for (uint32_t i = 0; i < 16u; i++) kvv[i] = 0.0f; + } + /* warp all-reduce so every lane sees the full dot product */ + for (uint32_t off = 16u; off > 0; off >>= 1) { + part += __shfl_xor_sync(0xffffffffu, part, off); + } + if (valid_row) { + const float score = part * scale; + const float new_m = fmaxf(M, score); + const float old_scale = expf(M - new_m); + const float row_scale = expf(score - new_m); +#pragma unroll + for (uint32_t i = 0; i < 16u; i++) { + o[i] = o[i] * old_scale + kvv[i] * row_scale; + } + S = S * old_scale + row_scale; + M = new_m; + } + } + + float *out = + partial_lora + ((uint64_t)block * n_head + head) * kv_lora_dim; +#pragma unroll + for (uint32_t c = 0; c < 4u; c++) { +#pragma unroll + for (uint32_t k = 0; k < 4u; k++) { + out[c * 128u + lane * 4u + k] = o[c * 4u + k]; + } + } + if (lane == 0u) { + float *ms = partial_ms + ((uint64_t)block * n_head + head) * 2u; + ms[0] = M; + ms[1] = S; + } +} + +__global__ static void glm_attn_split_g8_reduce_kernel( + float *heads, + const float *partial_lora, + const float *partial_ms, + const char *value_weight, + uint32_t n_head, + uint32_t n_blocks, + uint32_t value_dim, + uint32_t value_row_bytes) { + const uint32_t head = blockIdx.x; + if (head >= n_head || n_blocks == 0u || n_blocks > 64u) return; + + const uint32_t kv_lora_dim = 512u; + const uint32_t nth = blockDim.x; + const uint32_t tid = threadIdx.x; + + __shared__ float red[256]; + __shared__ float block_scale[64]; + __shared__ float lora_sum[512]; + + float local_m = -FLT_MAX / 2.0f; + if (tid < n_blocks) { + const float *ms = partial_ms + ((uint64_t)tid * n_head + head) * 2u; + local_m = ms[1] > 0.0f ? ms[0] : -FLT_MAX / 2.0f; + } + red[tid] = local_m; + __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_m = red[0]; + __syncthreads(); + + float local_denom = 0.0f; + if (tid < n_blocks) { + const float *ms = partial_ms + ((uint64_t)tid * n_head + head) * 2u; + const float s = ms[1]; + const float e = s > 0.0f ? expf(ms[0] - max_m) : 0.0f; + block_scale[tid] = e; + local_denom = s * e; + } + red[tid] = local_denom; + __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); + + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint32_t b = 0; b < n_blocks; b++) { + const float *src = + partial_lora + ((uint64_t)b * n_head + head) * kv_lora_dim; + acc += src[j] * block_scale[b]; + } + lora_sum[j] = acc / denom; + } + __syncthreads(); + + float *out = heads + (uint64_t)head * value_dim; + for (uint32_t d = tid; d < value_dim; d += nth) { + const char *row = + value_weight + ((uint64_t)head * value_dim + d) * value_row_bytes; + out[d] = glm_attn_q8_dot_row(row, lora_sum, kv_lora_dim); + } +} + +extern "C" int ds4_gpu_glm_attention_indexed_decode_split_group8_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, + 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) { + /* selected_rows_valid only permits skipping the row bounds check; the + * kernel always checks, which is correct in both cases. */ + (void)selected_rows_valid; + const uint32_t qk_dim = qk_nope + qk_rope; + const uint32_t needed_blocks = + block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u; + if (!heads || !partial_lora || !partial_ms || !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 || (n_head % 8u) != 0 || + kv_lora_dim != 512u || + qk_nope == 0 || qk_rope != 64u || + value_dim == 0 || qk_dim < qk_nope || + block_rows == 0u || needed_blocks == 0u || + n_blocks < needed_blocks || n_blocks > 64u || + !cache_f16 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + 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 (heads->bytes < (uint64_t)n_head * value_dim * sizeof(float) || + partial_lora->bytes < + (uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float) || + partial_ms->bytes < (uint64_t)n_blocks * n_head * 2u * 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 * sizeof(uint16_t) || + k_rope_cache->bytes < + (uint64_t)cache_cap * qk_rope * sizeof(uint16_t) || + selected->bytes < (uint64_t)n_selected * sizeof(uint32_t)) { + fprintf(stderr, "ds4: CUDA GLM split grouped indexed attention received undersized buffers\n"); + return 0; + } + if (value_weight_offset > model_size || + value_weight_bytes > model_size - value_weight_offset) { + fprintf(stderr, "ds4: CUDA GLM split grouped indexed attention value range is outside the mapped model\n"); + return 0; + } + const char *value_weight = cuda_model_range_ptr( + model_map, value_weight_offset, value_weight_bytes, + "glm attn v_b weights"); + if (!value_weight) return 0; + + dim3 partial_grid(n_head, n_blocks, 1); + glm_attn_split_g8_partial_kernel<<>>( + (float *)partial_lora->ptr, + (float *)partial_ms->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const char *)kv_lora_cache->ptr, + (const char *)k_rope_cache->ptr, + (const uint32_t *)selected->ptr, + n_selected, cache_cap, n_head, qk_nope, n_ctx_orig, block_rows, + 1.0f / sqrtf((float)qk_dim), + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + if (!cuda_ok(cudaGetLastError(), + "glm split grouped indexed attention partial launch")) { + return 0; + } + + glm_attn_split_g8_reduce_kernel<<>>( + (float *)heads->ptr, + (const float *)partial_lora->ptr, + (const float *)partial_ms->ptr, + value_weight, + n_head, n_blocks, value_dim, (uint32_t)value_row_bytes); + return cuda_ok(cudaGetLastError(), + "glm split grouped indexed attention reduce launch"); +} + +/* ------------------------------------------------------------------------ + * Indexed batch attention with lora output: same scores as the indexed + * kernels but the softmax-weighted sum of kv_lora rows is written directly + * (no q8_0 value projection). One kernel covers all three entry points: + * - selected list (lora / lora_valid): row = selected[token][s] + * - causal (lora_causal, selected == NULL): row = s, s < pos0 + token + 1 + * Port of kernel_glm_attention_indexed_batch_lora_group8_vec_impl and + * ..._vec_causal_impl (also covering the group8 pad0==1 fallback layout). + * The "valid" variant only relaxes the row bounds check, which the kernel + * always performs. + * + * Grid (n_head, n_tokens), 256 threads. Dynamic shared memory: + * red[256] + scores[n_selected]. + * ------------------------------------------------------------------------ */ + +__global__ static void glm_attn_batch_lora_kernel( + float *lora_out, + const float *q, + const float *qk_low, + const char *kv_lora_cache, + const char *k_rope_cache, + const uint32_t *selected, /* NULL => causal row = s */ + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_selected, + uint32_t cache_cap, + uint32_t 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 scale, + 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; + if (head >= n_head || token >= n_tokens || n_selected == 0u) return; + const uint32_t nth = blockDim.x; + const uint32_t tid = threadIdx.x; + const uint32_t qk_dim = qk_nope + qk_rope; + const uint32_t count = + selected ? n_selected : min(n_selected, pos0 + token + 1u); + if (count == 0u) return; + + extern __shared__ float glm_attn_lora_sm[]; + float *red = glm_attn_lora_sm; + float *scores = glm_attn_lora_sm + 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; + const uint32_t *token_selected = + selected ? selected + (uint64_t)token * n_selected : NULL; + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + glm_attn_rope_corr_dims((int)qk_rope, (int)n_ctx_orig, freq_base, + beta_fast, beta_slow, &corr0, &corr1); + } + + float local_max = -FLT_MAX / 2.0f; + for (uint32_t s = tid; s < count; s += nth) { + const uint32_t row = token_selected ? token_selected[s] : s; + float score = -FLT_MAX / 2.0f; + if (row < cache_cap) { + float dotv = 0.0f; + const uint64_t lora_base = (uint64_t)row * kv_lora_dim; + for (uint32_t j = 0; j < kv_lora_dim; j++) { + dotv += low[j] * + glm_attn_cache_load(kv_lora_cache, lora_base + j, + cache_f16); + } + const uint64_t rope_base = (uint64_t)row * qk_rope; + for (uint32_t r = 0; r < qk_rope; r += 2u) { + const float2 y = glm_attn_rope_pair(k_rope_cache, rope_base, r, + row, qk_rope, cache_f16, + freq_base, freq_scale, + ext_factor, attn_factor, + corr0, corr1); + 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]; + + float local_sum = 0.0f; + for (uint32_t s = tid; s < count; s += nth) { + const uint32_t row = token_selected ? token_selected[s] : s; + /* Invalid rows contribute nothing to the softmax sum, matching the + * Metal vec kernels' online accumulation. */ + const float w = + row < cache_cap ? expf(scores[s] - max_score) : 0.0f; + 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 sum = red[0]; + const float inv_s = sum > 0.0f ? 1.0f / sum : 0.0f; + __syncthreads(); + + float *out = + lora_out + ((uint64_t)token * n_head + head) * kv_lora_dim; + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + float acc = 0.0f; + for (uint32_t s = 0; s < count; s++) { + const uint32_t row = token_selected ? token_selected[s] : s; + if (row < cache_cap) { + acc += scores[s] * + glm_attn_cache_load(kv_lora_cache, + (uint64_t)row * kv_lora_dim + j, + cache_f16); + } + } + out[j] = acc * inv_s; + } +} + +static int glm_attn_batch_lora_launch( + 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, /* NULL => causal */ + 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, + const char *label) { + const uint32_t qk_dim = qk_nope + qk_rope; + if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || + n_tokens == 0 || n_selected == 0 || cache_cap == 0 || + n_selected > cache_cap || + n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || (qk_rope & 1u) != 0 || + qk_dim < qk_nope || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (lora_out->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || + 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_bytes || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem_bytes || + (selected && + selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(uint32_t))) { + fprintf(stderr, "ds4: CUDA GLM indexed batch attention-lora received undersized buffers\n"); + return 0; + } + const uint64_t smem = (256u + (uint64_t)n_selected) * sizeof(float); + if (!glm_attn_smem_ok((const void *)glm_attn_batch_lora_kernel, smem, + "glm batch attention-lora smem opt-in")) { + return 0; + } + dim3 grid(n_head, n_tokens, 1); + glm_attn_batch_lora_kernel<<>>( + (float *)lora_out->ptr, + (const float *)q->ptr, + (const float *)qk_low->ptr, + (const char *)kv_lora_cache->ptr, + (const char *)k_rope_cache->ptr, + selected ? (const uint32_t *)selected->ptr : NULL, + n_tokens, pos0, n_selected, cache_cap, + cache_f16 ? 1u : 0u, + n_head, kv_lora_dim, qk_nope, qk_rope, n_ctx_orig, + 1.0f / sqrtf((float)qk_dim), + freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), label); +} + +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) { + if (!selected) return 0; + return glm_attn_batch_lora_launch(lora_out, q, qk_low, kv_lora_cache, + k_rope_cache, selected, + n_tokens, 0u /* pos0 */, 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, + "glm indexed batch attention-lora 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) { + /* "valid" asserts every selected row is in-bounds; the shared kernel + * bounds-checks anyway, which is correct for that case too. */ + if (!selected) return 0; + return glm_attn_batch_lora_launch(lora_out, q, qk_low, kv_lora_cache, + k_rope_cache, selected, + n_tokens, 0u /* pos0 */, 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, + "glm indexed batch attention-lora valid launch"); +} + +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) { + /* The Metal causal launcher additionally pins the fast-path shape. */ + if (kv_lora_dim != 512u || qk_rope != 64u || !cache_f16 || + pos0 > n_selected || n_tokens > n_selected - pos0) { + return 0; + } + return glm_attn_batch_lora_launch(lora_out, q, qk_low, kv_lora_cache, + k_rope_cache, NULL /* causal */, + n_tokens, pos0, 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, + "glm causal indexed batch attention-lora launch"); +} diff --git a/ds4_cuda_glm_indexer.inc b/ds4_cuda_glm_indexer.inc new file mode 100644 index 000000000..519649bcc --- /dev/null +++ b/ds4_cuda_glm_indexer.inc @@ -0,0 +1,604 @@ +/* ds4_cuda_glm_indexer.inc: CUDA port of the GLM-5.2 indexer / rope-tail / + * low-rank QK Metal kernels (launchers in ds4_metal.m, shaders in + * metal/dsv4_misc.metal). This file is #included near the end of ds4_cuda.cu + * and relies on its internals: cuda_ok(), cuda_model_range_ptr(), + * warp_sum_f32(), struct ds4_gpu_tensor {ptr,bytes}. + * + * Entry points implemented here: + * ds4_gpu_glm_rope_tail_tensor + * ds4_gpu_glm_indexer_rope_tail_tensor + * ds4_gpu_glm_indexer_score_one_tensor + * ds4_gpu_glm_indexer_scores_batch_tensor + * ds4_gpu_glm_qk_lowrank_q8_0_tensor + * ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor + * ds4_gpu_glm_value_project_q8_0_batch_heads_tensor + * + * ponytail: only the scalar/generic Metal kernels are ported (plus the + * decode-hot score-one direct kernel). The Metal tiled-simdgroup batch + * scores, glm52-specialized qk-lowrank and the value-project MMA fast + * paths are perf variants with identical math; add WMMA versions here if + * profiling shows these launches on the critical path. */ + +/* --------------------------------------------------------------------------- + * Device helpers. + * ------------------------------------------------------------------------- */ + +/* Port of glm_cache_load_f32_or_f16 (metal/dsv4_misc.metal). */ +__device__ static float glm_idx_cache_load( + const char *base, + uint64_t index, + uint32_t cache_f16) { + if (cache_f16 != 0u) { + return __half2float(((const __half *)base)[index]); + } + return ((const float *)base)[index]; +} + +/* Port of glm_q8_0_dot_row_dev_f32 (metal/dsv4_misc.metal). One q8_0 block + * is 34 bytes: __half scale followed by 32 int8 quants. */ +__device__ static float glm_idx_q8_0_dot_row( + const char *row, + const float *x, + uint32_t n_cols) { + float acc = 0.0f; + const uint32_t n_blocks = (n_cols + 31u) >> 5; + for (uint32_t block = 0; block < n_blocks; block++) { + const char *block_base = row + (uint64_t)block * 34u; + const float d = __half2float(*(const __half *)block_base); + const int8_t *qs = (const int8_t *)(block_base + 2u); + const uint32_t base = block << 5; + const uint32_t count = n_cols - base < 32u ? n_cols - base : 32u; + for (uint32_t qi = 0; qi < count; qi++) { + acc += d * (float)qs[qi] * x[base + qi]; + } + } + return acc; +} + +static uint64_t glm_idx_q8_0_row_bytes(uint32_t n_cols) { + return (((uint64_t)n_cols + 31u) / 32u) * 34u; +} + +/* --------------------------------------------------------------------------- + * GLM RoPE tail (kernel_glm_indexer_rope_tail_f32). + * + * Rotates a rot_dim slice starting at rot_offset inside each head vector, + * with contiguous (i, i+1) pairs and YaRN corrections. Both + * ds4_gpu_glm_rope_tail_tensor (rot_offset = head_dim - rot_dim) and + * ds4_gpu_glm_indexer_rope_tail_tensor (rot_offset = 0) route here, exactly + * like the shared Metal launcher ds4_gpu_glm_rope_tail_offset_tensor. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_rope_tail_offset_kernel( + float *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 uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t half_rot = rot_dim / 2u; + const uint32_t pairs = n_tokens * n_head * half_rot; + if (gid >= pairs) return; + const uint32_t pair = gid % half_rot; + const uint32_t tmp = gid / half_rot; + const uint32_t head = tmp % n_head; + const uint32_t token = tmp / n_head; + 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)rot_dim * + logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom); + corr1 = ceilf((float)rot_dim * + logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom); + corr0 = fmaxf(0.0f, corr0); + corr1 = fminf((float)(rot_dim - 1u), corr1); + } + + const float theta_extrap = (float)(pos0 + token) * + powf(freq_base, -((float)i) / (float)rot_dim); + 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; + const float s = sinf(theta) * mscale; + + float *row = x + ((uint64_t)token * n_head + head) * head_dim + rot_offset; + const float x0 = row[i]; + const float x1 = row[i + 1u]; + row[i] = x0 * c - x1 * s; + row[i + 1u] = x0 * s + x1 * c; +} + +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 *label) { + if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 || + rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset || + (rot_dim & 1u) != 0 || + pos0 > UINT32_MAX - n_tokens || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow) || + x->bytes < (uint64_t)n_tokens * n_head * head_dim * sizeof(float)) { + return 0; + } + const uint32_t pairs = n_tokens * n_head * (rot_dim / 2u); + glm_rope_tail_offset_kernel<<<(pairs + 255u) / 256u, 256>>>( + (float *)x->ptr, + n_tokens, 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(), label); +} + +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 launch"); +} + +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 launch"); +} + +/* --------------------------------------------------------------------------- + * GLM indexer scores, single token (kernel_glm_indexer_score_one + + * kernel_glm_indexer_score_one_direct). + * + * scores[row] = sum_h relu(dot(q[h], key_cache[row]) * scale) * weights[h] + * Note: unlike the DeepSeek indexer, relu is applied AFTER scaling and the + * final sum is NOT rescaled. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_indexer_score_one_kernel( + float *scores, + const float *q, + const float *weights, + const char *indexer_key_cache, + uint32_t n_rows, + uint32_t n_head, + uint32_t head_dim, + uint32_t cache_f16, + float scale) { + const uint32_t row = blockIdx.x; + if (row >= n_rows) return; + __shared__ float partial[256]; + float score = 0.0f; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + (uint64_t)h * head_dim; + float dot = 0.0f; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + const float k = glm_idx_cache_load(indexer_key_cache, + (uint64_t)row * head_dim + d, + cache_f16); + dot += qh[d] * k; + } + 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(); + } + if (threadIdx.x == 0) { + score += fmaxf(partial[0] * scale, 0.0f) * weights[h]; + } + __syncthreads(); + } + if (threadIdx.x == 0) scores[row] = score; +} + +/* Fast path for the GLM indexer decode shape (n_head=32, head_dim=128). + * One block of 128 threads per cache row, one warp per head, four heads + * per pass. Mirrors kernel_glm_indexer_score_one_direct. */ +__global__ static void glm_indexer_score_one_direct32_kernel( + float *scores, + const float *q, + const float *weights, + const char *indexer_key_cache, + uint32_t n_rows, + uint32_t cache_f16, + float scale) { + const uint32_t row = blockIdx.x; + const uint32_t tid = threadIdx.x; + const uint32_t lane = tid & 31u; + const uint32_t warp = tid >> 5u; + if (row >= n_rows || tid >= 128u) return; + + __shared__ float krow[128]; + __shared__ float psum[4]; + krow[tid] = glm_idx_cache_load(indexer_key_cache, + (uint64_t)row * 128u + tid, + cache_f16); + __syncthreads(); + + float acc = 0.0f; + for (uint32_t head0 = 0; head0 < 32u; head0 += 4u) { + const uint32_t head = head0 + warp; + const float4 qv = ((const float4 *)(q + (uint64_t)head * 128u))[lane]; + const float4 kv = ((const float4 *)krow)[lane]; + float s = qv.x * kv.x + qv.y * kv.y + qv.z * kv.z + qv.w * kv.w; + s = warp_sum_f32(s); + if (lane == 0) psum[warp] = fmaxf(s * scale, 0.0f) * weights[head]; + __syncthreads(); + if (tid == 0) acc += psum[0] + psum[1] + psum[2] + psum[3]; + __syncthreads(); + } + if (tid == 0) scores[row] = acc; +} + +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) { + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_head == 0 || head_dim == 0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (scores->bytes < (uint64_t)n_rows * sizeof(float) || + q->bytes < (uint64_t)n_head * head_dim * sizeof(float) || + weights->bytes < (uint64_t)n_head * sizeof(float) || + indexer_key_cache->bytes < (uint64_t)n_rows * head_dim * cache_elem_bytes) { + return 0; + } + if (n_head == 32u && head_dim == 128u) { + glm_indexer_score_one_direct32_kernel<<>>( + (float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const char *)indexer_key_cache->ptr, + n_rows, cache_f16 ? 1u : 0u, scale); + return cuda_ok(cudaGetLastError(), "glm indexer score one direct launch"); + } + glm_indexer_score_one_kernel<<>>( + (float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const char *)indexer_key_cache->ptr, + n_rows, n_head, head_dim, cache_f16 ? 1u : 0u, scale); + return cuda_ok(cudaGetLastError(), "glm indexer score one launch"); +} + +/* --------------------------------------------------------------------------- + * GLM indexer scores, token batch (kernel_glm_indexer_scores_batch). + * + * scores[token][row] as above, with causal masking: rows at or beyond + * pos0 + token + 1 get -INF. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_indexer_scores_batch_kernel( + float *scores, + const float *q, + const float *weights, + const char *indexer_key_cache, + uint32_t n_rows, + uint32_t n_tokens, + uint32_t pos0, + uint32_t n_head, + uint32_t head_dim, + uint32_t cache_f16, + float scale) { + const uint32_t row = blockIdx.x; + const uint32_t token = blockIdx.y; + if (row >= n_rows || token >= n_tokens) return; + + float *dst = scores + (uint64_t)token * n_rows + row; + const uint32_t visible = min(pos0 + token + 1u, n_rows); + if (row >= visible) { + if (threadIdx.x == 0) *dst = -INFINITY; + return; + } + + __shared__ float partial[256]; + float score = 0.0f; + for (uint32_t h = 0; h < n_head; h++) { + const float *qh = q + ((uint64_t)token * n_head + h) * head_dim; + float dot = 0.0f; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + const float k = glm_idx_cache_load(indexer_key_cache, + (uint64_t)row * head_dim + d, + cache_f16); + dot += qh[d] * k; + } + 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(); + } + if (threadIdx.x == 0) { + const float *w = weights + (uint64_t)token * n_head; + score += fmaxf(partial[0] * scale, 0.0f) * w[h]; + } + __syncthreads(); + } + if (threadIdx.x == 0) *dst = score; +} + +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) { + /* head_dim must be 128, matching the Metal launcher's contract. */ + if (!scores || !q || !weights || !indexer_key_cache || + n_rows == 0 || n_tokens == 0 || n_head == 0 || head_dim != 128u || + pos0 >= n_rows || n_tokens > n_rows - pos0 || + !isfinite(scale) || scale <= 0.0f) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (scores->bytes < (uint64_t)n_rows * n_tokens * sizeof(float) || + 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_bytes) { + return 0; + } + dim3 grid(n_rows, n_tokens, 1); + glm_indexer_scores_batch_kernel<<>>( + (float *)scores->ptr, + (const float *)q->ptr, + (const float *)weights->ptr, + (const char *)indexer_key_cache->ptr, + n_rows, n_tokens, pos0, n_head, head_dim, + cache_f16 ? 1u : 0u, scale); + return cuda_ok(cudaGetLastError(), "glm indexer batch scores launch"); +} + +/* --------------------------------------------------------------------------- + * GLM low-rank QK projection over q8_0 weights + * (kernel_glm_qk_lowrank_q8_0 / _batch). + * + * Per head h and token t: + * qk_low[t][h][j] = dot(weight_row(h * kv_lora_dim + j)[0..qk_nope), + * q[t][h][0..qk_nope)) + * The single-token entry point is the batch kernel with n_tokens = 1. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_qk_lowrank_q8_0_kernel( + float *qk_low, + const float *q, + const char *weight, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim, + uint32_t row_bytes) { + const uint32_t head = blockIdx.x; + const uint32_t token = blockIdx.y; + if (head >= n_head || token >= n_tokens) return; + const float *qh = q + ((uint64_t)token * n_head + head) * qk_dim; + float *out = qk_low + ((uint64_t)token * n_head + head) * kv_lora_dim; + for (uint32_t j = threadIdx.x; j < kv_lora_dim; j += blockDim.x) { + const char *row = weight + ((uint64_t)head * kv_lora_dim + j) * row_bytes; + out[j] = glm_idx_q8_0_dot_row(row, qh, qk_nope); + } +} + +static int glm_qk_lowrank_q8_0_launch( + ds4_gpu_tensor *qk_low, + const ds4_gpu_tensor *q, + const void *model_map, + uint64_t model_size, + uint64_t weight_offset, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim, + const char *label) { + if (!qk_low || !q || !model_map || + n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_nope > qk_dim) { + return 0; + } + const uint64_t row_bytes = glm_idx_q8_0_row_bytes(qk_nope); + const uint64_t weight_bytes = (uint64_t)n_head * kv_lora_dim * row_bytes; + if (qk_low->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float) || + q->bytes < (uint64_t)n_tokens * n_head * qk_dim * sizeof(float)) { + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: CUDA GLM qk lowrank range is outside the mapped model\n"); + return 0; + } + const char *wptr = cuda_model_range_ptr(model_map, weight_offset, + weight_bytes, "glm qk lowrank"); + if (!wptr) return 0; + dim3 grid(n_head, n_tokens, 1); + glm_qk_lowrank_q8_0_kernel<<>>( + (float *)qk_low->ptr, + (const float *)q->ptr, + wptr, + n_tokens, n_head, kv_lora_dim, qk_nope, qk_dim, + (uint32_t)row_bytes); + return cuda_ok(cudaGetLastError(), label); +} + +extern "C" int ds4_gpu_glm_qk_lowrank_q8_0_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 n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return glm_qk_lowrank_q8_0_launch(qk_low, q, model_map, model_size, + weight_offset, 1, n_head, kv_lora_dim, + qk_nope, qk_dim, + "glm qk lowrank launch"); +} + +extern "C" int ds4_gpu_glm_qk_lowrank_q8_0_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 n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t qk_dim) { + return glm_qk_lowrank_q8_0_launch(qk_low, q, model_map, model_size, + weight_offset, n_tokens, n_head, + kv_lora_dim, qk_nope, qk_dim, + "glm batch qk lowrank launch"); +} + +/* --------------------------------------------------------------------------- + * GLM value projection over q8_0 weights, per (token, head) + * (kernel_glm_value_project_q8_0_batch_heads). + * + * heads[t][h][d] = dot(weight_row(h * value_dim + d), + * lora[t][h][0..kv_lora_dim)) + * The lora vector is staged in dynamic shared memory like the Metal + * threadgroup buffer. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_value_project_q8_0_batch_heads_kernel( + float *heads, + const float *lora, + const char *weight, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t value_dim, + uint32_t row_bytes) { + extern __shared__ float glm_idx_vp_x[]; + const uint32_t head = blockIdx.x; + const uint32_t token = blockIdx.y; + if (head >= n_head || token >= n_tokens) return; + const float *src = lora + ((uint64_t)token * n_head + head) * kv_lora_dim; + float *out = heads + ((uint64_t)token * n_head + head) * value_dim; + + for (uint32_t j = threadIdx.x; j < kv_lora_dim; j += blockDim.x) { + glm_idx_vp_x[j] = src[j]; + } + __syncthreads(); + + for (uint32_t d = threadIdx.x; d < value_dim; d += blockDim.x) { + const char *row = weight + ((uint64_t)head * value_dim + d) * row_bytes; + out[d] = glm_idx_q8_0_dot_row(row, glm_idx_vp_x, kv_lora_dim); + } +} + +extern "C" int ds4_gpu_glm_value_project_q8_0_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 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 || value_dim == 0) { + return 0; + } + const uint64_t row_bytes = glm_idx_q8_0_row_bytes(kv_lora_dim); + const uint64_t weight_bytes = (uint64_t)n_head * value_dim * row_bytes; + if (heads->bytes < (uint64_t)n_tokens * n_head * value_dim * sizeof(float) || + lora->bytes < (uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float)) { + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) { + fprintf(stderr, "ds4: CUDA GLM batch value project range is outside the mapped model\n"); + return 0; + } + const char *wptr = cuda_model_range_ptr(model_map, weight_offset, + weight_bytes, "glm value project"); + if (!wptr) return 0; + dim3 grid(n_head, n_tokens, 1); + glm_value_project_q8_0_batch_heads_kernel + <<>>( + (float *)heads->ptr, + (const float *)lora->ptr, + wptr, + n_tokens, n_head, kv_lora_dim, value_dim, + (uint32_t)row_bytes); + return cuda_ok(cudaGetLastError(), "glm batch value project launch"); +} diff --git a/ds4_cuda_glm_kv.inc b/ds4_cuda_glm_kv.inc new file mode 100644 index 000000000..071b081b9 --- /dev/null +++ b/ds4_cuda_glm_kv.inc @@ -0,0 +1,833 @@ +/* ============================================================================ + * GLM-5.2 KV-write path kernels, ported from metal/dsv4_misc.metal + * (kernel_glm_kv_lora_rms_norm, kernel_glm_k_b_project_q8_0, + * kernel_glm_store_compact_kv, kernel_glm_qkv_norm_store_compact_kv, + * kernel_glm_store_indexer_k, kernel_glm_build_kv_cache[_flash]). + * + * This file is #included near the end of ds4_cuda.cu and relies on helpers + * defined there: cuda_ok(), cuda_model_range_ptr(), rope_yarn_ramp_dev(), + * struct ds4_gpu_tensor. + * ============================================================================ */ + +/* --------------------------------------------------------------------------- + * Device helpers. + * ------------------------------------------------------------------------- */ + +__device__ static void glm_cache_store_f32_or_f16_dev( + char *base, + uint64_t index, + uint32_t cache_f16, + float x) { + if (cache_f16 != 0u) { + ((__half *)base)[index] = __float2half(x); + } else { + ((float *)base)[index] = x; + } +} + +/* YaRN correction dims, Metal glm_rope_yarn_corr_dims equivalent. */ +__device__ static void glm_rope_yarn_corr_dims_dev( + uint32_t n_dims, + uint32_t n_ctx_orig, + float freq_base, + float beta_fast, + float beta_slow, + float *corr0, + float *corr1) { + const float denom = 2.0f * logf(freq_base); + const float c0 = floorf((float)n_dims * + logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom); + const float c1 = ceilf((float)n_dims * + logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom); + *corr0 = fmaxf(0.0f, c0); + *corr1 = fminf((float)n_dims - 1.0f, c1); +} + +/* Metal glm_rope_yarn equivalent (reuses rope_yarn_ramp_dev from ds4_cuda.cu). */ +__device__ static void glm_rope_yarn_dev( + float theta_extrap, + float freq_scale, + float corr0, + float corr1, + int i0, + float ext_factor, + float mscale, + float *cos_theta, + float *sin_theta) { + const float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + const float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, 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; +} + +/* --------------------------------------------------------------------------- + * kernel_glm_kv_lora_rms_norm: RMS norm of the first kv_lora_dim elements of + * each kv_raw row (row stride kv_raw_dim) into a dense [n_tokens, kv_lora_dim] + * output, scaled by a norm weight. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_kv_lora_rms_norm_kernel( + float *out, + const float *kv_raw, + 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 float *x = kv_raw + (uint64_t)row * kv_raw_dim; + float *o = out + (uint64_t)row * kv_lora_dim; + + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < kv_lora_dim; i += blockDim.x) { + const float v = x[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 inv = rsqrtf(partial[0] / (float)kv_lora_dim + eps); + for (uint32_t i = threadIdx.x; i < kv_lora_dim; i += blockDim.x) { + o[i] = x[i] * inv * w[i]; + } +} + +extern "C" int ds4_gpu_glm_kv_lora_rms_norm_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *kv_raw, + 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 (!out || !kv_raw || !model_map || + n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || + kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 || + !isfinite(eps) || eps < 0.0f) { + return 0; + } + const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + if (kv_raw->bytes < raw_bytes || out->bytes < out_bytes) return 0; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0; + const float *w = (const float *)cuda_model_range_ptr(model_map, weight_offset, + weight_bytes, + "glm kv_lora rms weight"); + if (!w) return 0; + glm_kv_lora_rms_norm_kernel<<>>( + (float *)out->ptr, (const float *)kv_raw->ptr, w, + n_tokens, kv_raw_dim, kv_lora_dim, eps); + return cuda_ok(cudaGetLastError(), "glm kv_lora rms norm launch"); +} + +/* --------------------------------------------------------------------------- + * kernel_glm_k_b_project_q8_0: per (token, head) projection of the normalized + * kv_lora row through a q8_0 weight laid out as [n_head * kv_lora_dim] rows of + * qk_nope elements each (34-byte blocks of 32). out[token][head][qk_nope]. + * Block shape mirrors Metal: (32 lanes, q_blocks) with the kv row staged in + * shared memory. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_k_b_project_q8_0_kernel( + float *out, + const unsigned char *weight, + const float *kv_norm, + uint32_t n_tokens, + uint32_t n_head, + uint32_t kv_lora_dim, + uint32_t qk_nope, + uint32_t row_bytes) { + const uint32_t token = blockIdx.x; + const uint32_t head = blockIdx.y; + if (token >= n_tokens || head >= n_head) return; + + extern __shared__ float glm_kb_kv_scratch[]; + const uint32_t nth = blockDim.x * blockDim.y; + const uint32_t tid = threadIdx.y * blockDim.x + threadIdx.x; + const float *kv = kv_norm + (uint64_t)token * kv_lora_dim; + float *o = out + ((uint64_t)token * n_head + head) * qk_nope; + + for (uint32_t j = tid; j < kv_lora_dim; j += nth) { + glm_kb_kv_scratch[j] = kv[j]; + } + __syncthreads(); + + const uint32_t block = threadIdx.y; + const uint32_t lane = threadIdx.x; + const uint32_t q = (block << 5) + lane; + if (q < qk_nope) { + float acc = 0.0f; + const unsigned char *hbase = weight + (uint64_t)head * kv_lora_dim * row_bytes; + for (uint32_t j = 0; j < kv_lora_dim; j++) { + const unsigned char *row = hbase + (uint64_t)j * row_bytes; + const __half *dptr = (const __half *)(row + (uint64_t)block * 34u); + const int8_t *qs = (const int8_t *)(row + (uint64_t)block * 34u + 2u); + acc += __half2float(*dptr) * (float)qs[lane] * glm_kb_kv_scratch[j]; + } + o[q] = acc; + } +} + +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) { + if (!out || !kv_norm || !model_map || + n_tokens == 0 || kv_lora_dim == 0 || qk_nope == 0 || n_head == 0) { + return 0; + } + const uint64_t row_bytes = (((uint64_t)qk_nope + 31u) / 32u) * 34u; + const uint64_t weight_bytes = (uint64_t)n_head * kv_lora_dim * row_bytes; + const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float); + const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float); + if (kv_norm->bytes < kv_bytes || out->bytes < out_bytes) return 0; + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0; + const uint32_t q_blocks = (qk_nope + 31u) / 32u; + if (q_blocks > 8u) return 0; /* mirrors the Metal tiled-kernel limit */ + const unsigned char *w = (const unsigned char *) + cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "glm k_b q8_0"); + if (!w) return 0; + dim3 grid(n_tokens, n_head, 1); + dim3 block(32, q_blocks, 1); + glm_k_b_project_q8_0_kernel<<>>( + (float *)out->ptr, w, (const float *)kv_norm->ptr, + n_tokens, n_head, kv_lora_dim, qk_nope, (uint32_t)row_bytes); + return cuda_ok(cudaGetLastError(), "glm k_b project launch"); +} + +/* --------------------------------------------------------------------------- + * kernel_glm_store_compact_kv: copies the normalized kv_lora row and the raw + * rope tail into the compact caches at pos0+token. blockIdx.y selects the + * part (0 = kv_lora, 1 = rope tail), as in the Metal dispatch. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_store_compact_kv_kernel( + const float *kv_norm, + const float *kv_raw, + char *kv_lora_cache, + char *k_rope_cache, + 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; + + if (part == 0u) { + const float *src = kv_norm + (uint64_t)token * kv_lora_dim; + for (uint32_t i = threadIdx.x; i < kv_lora_dim; i += blockDim.x) { + glm_cache_store_f32_or_f16_dev(kv_lora_cache, + (uint64_t)pos * kv_lora_dim + i, + cache_f16, src[i]); + } + } else { + const float *src = kv_raw + (uint64_t)token * kv_raw_dim + kv_lora_dim; + for (uint32_t i = threadIdx.x; i < qk_rope; i += blockDim.x) { + glm_cache_store_f32_or_f16_dev(k_rope_cache, + (uint64_t)pos * qk_rope + i, + cache_f16, 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 || + kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem_bytes || + kv_norm->bytes < (uint64_t)n_tokens * kv_lora_dim * sizeof(float) || + kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float)) { + return 0; + } + dim3 grid(n_tokens, 2, 1); + glm_store_compact_kv_kernel<<>>( + (const float *)kv_norm->ptr, (const float *)kv_raw->ptr, + (char *)kv_lora_cache->ptr, (char *)k_rope_cache->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"); +} + +/* --------------------------------------------------------------------------- + * kernel_glm_qkv_norm_store_compact_kv: fused RMS norm of the q row (part 0) + * and of the kv_lora prefix of the raw row (part 1, stored to the compact + * cache), plus a rope-tail copy (part 2). blockIdx.y selects the part. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_qkv_norm_store_compact_kv_kernel( + float *q_dst, + const float *q_src, + const float *q_weight, + const float *kv_raw, + const float *kv_weight, + char *kv_lora_cache, + char *k_rope_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t q_n, + uint32_t kv_raw_dim, + uint32_t kv_lora_dim, + uint32_t qk_rope, + uint32_t cache_f16, + float eps) { + const uint32_t token = blockIdx.x; + const uint32_t part = blockIdx.y; + if (token >= n_tokens || part > 2u) return; + + if (part == 2u) { + const uint32_t pos = pos0 + token; + if (pos >= cache_cap) return; + const float *src = kv_raw + (uint64_t)token * kv_raw_dim + kv_lora_dim; + for (uint32_t i = threadIdx.x; i < qk_rope; i += blockDim.x) { + glm_cache_store_f32_or_f16_dev(k_rope_cache, + (uint64_t)pos * qk_rope + i, + cache_f16, 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_weight : q_weight; + + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { + const float v = x[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 (!kv_task) { + float *y = q_dst + (uint64_t)token * q_n; + for (uint32_t i = threadIdx.x; i < q_n; i += blockDim.x) { + y[i] = (x[i] * scale) * w[i]; + } + return; + } + + const uint32_t pos = pos0 + token; + if (pos >= cache_cap) return; + for (uint32_t i = threadIdx.x; i < kv_lora_dim; i += blockDim.x) { + glm_cache_store_f32_or_f16_dev(kv_lora_cache, + (uint64_t)pos * kv_lora_dim + i, + cache_f16, (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 || cache_cap == 0 || + q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 || qk_rope == 0 || + (q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 || + kv_lora_dim > kv_raw_dim || + qk_rope > kv_raw_dim - kv_lora_dim || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float); + const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float); + if (q->bytes < q_row_bytes * n_tokens || + q_out->bytes < q_row_bytes * n_tokens || + kv_lora_cache->bytes < (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes || + k_rope_cache->bytes < (uint64_t)cache_cap * qk_rope * cache_elem_bytes || + kv_raw->bytes < (uint64_t)n_tokens * kv_raw_dim * sizeof(float)) { + return 0; + } + if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) { + return 0; + } + const float *q_w = (const float *)cuda_model_range_ptr(model_map, q_weight_offset, + q_row_bytes, + "glm q rms weight"); + const float *kv_w = (const float *)cuda_model_range_ptr(model_map, kv_weight_offset, + kv_weight_bytes, + "glm kv rms weight"); + if (!q_w || !kv_w) return 0; + dim3 grid(n_tokens, 3, 1); + glm_qkv_norm_store_compact_kv_kernel<<>>( + (float *)q_out->ptr, (const float *)q->ptr, q_w, + (const float *)kv_raw->ptr, kv_w, + (char *)kv_lora_cache->ptr, (char *)k_rope_cache->ptr, + pos0, n_tokens, cache_cap, q_n, kv_raw_dim, kv_lora_dim, qk_rope, + cache_f16 ? 1u : 0u, eps); + return cuda_ok(cudaGetLastError(), "glm qkv norm compact store launch"); +} + +/* --------------------------------------------------------------------------- + * kernel_glm_store_indexer_k: LayerNorm (mean/variance, weight + bias) over + * the raw indexer K row, partial rope over the first rot_dim dims in adjacent + * pairs, stored to the indexer key cache at pos0+token. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_store_indexer_k_kernel( + const float *raw_k, + const float *w, + const float *b, + char *indexer_key_cache, + uint32_t pos0, + uint32_t n_tokens, + uint32_t cache_cap, + uint32_t head_dim, + uint32_t rot_dim, + uint32_t n_ctx_orig, + uint32_t cache_f16, + float eps, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + const uint32_t token = blockIdx.x; + if (token >= n_tokens) return; + const uint32_t pos = pos0 + token; + if (pos >= cache_cap) return; + + const float *src = raw_k + (uint64_t)token * head_dim; + __shared__ float partial[256]; + + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) { + sum += src[i]; + } + 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 mean = partial[0] / (float)head_dim; + __syncthreads(); + + float ss = 0.0f; + for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) { + const float d = src[i] - mean; + ss += d * d; + } + partial[threadIdx.x] = ss; + __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 inv = rsqrtf(partial[0] / (float)head_dim + eps); + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + glm_rope_yarn_corr_dims_dev(rot_dim, n_ctx_orig, freq_base, + beta_fast, beta_slow, &corr0, &corr1); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)rot_dim; + + for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) { + if (i < rot_dim) { + if ((i & 1u) != 0u) continue; + const float theta = theta_base * powf(freq_base, inv_ndims * (float)i); + float cos_theta, sin_theta; + glm_rope_yarn_dev(theta, freq_scale, corr0, corr1, (int)i, + ext_factor, attn_factor, &cos_theta, &sin_theta); + const float x0 = (src[i] - mean) * inv * w[i] + b[i]; + const uint32_t j = i + 1u; + const float x1 = (src[j] - mean) * inv * w[j] + b[j]; + glm_cache_store_f32_or_f16_dev(indexer_key_cache, + (uint64_t)pos * head_dim + i, + cache_f16, + x0 * cos_theta - x1 * sin_theta); + glm_cache_store_f32_or_f16_dev(indexer_key_cache, + (uint64_t)pos * head_dim + j, + cache_f16, + x0 * sin_theta + x1 * cos_theta); + } else { + const float x = (src[i] - mean) * inv * w[i] + b[i]; + glm_cache_store_f32_or_f16_dev(indexer_key_cache, + (uint64_t)pos * head_dim + i, + cache_f16, x); + } + } +} + +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 || cache_cap == 0 || + head_dim == 0 || rot_dim == 0 || + rot_dim > head_dim || (rot_dim & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float); + if (indexer_key_cache->bytes < (uint64_t)cache_cap * head_dim * cache_elem_bytes || + raw_k->bytes < (uint64_t)n_tokens * head_dim * sizeof(float)) { + return 0; + } + if (weight_offset > model_size || norm_bytes > model_size - weight_offset || + bias_offset > model_size || norm_bytes > model_size - bias_offset) { + return 0; + } + const float *w = (const float *)cuda_model_range_ptr(model_map, weight_offset, + norm_bytes, + "glm indexer k norm weight"); + const float *b = (const float *)cuda_model_range_ptr(model_map, bias_offset, + norm_bytes, + "glm indexer k norm bias"); + if (!w || !b) return 0; + glm_store_indexer_k_kernel<<>>( + (const float *)raw_k->ptr, w, b, (char *)indexer_key_cache->ptr, + pos0, n_tokens, cache_cap, head_dim, rot_dim, n_ctx_orig, + cache_f16 ? 1u : 0u, eps, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "glm store indexer k launch"); +} + +/* --------------------------------------------------------------------------- + * kernel_glm_build_kv_cache[_flash]: per (token, head) dense KV cache build. + * Copies the per-head k_nope, ropes the shared kv_raw rope tail into the key + * row after qk_nope, and copies the per-head value row. The flash variant + * additionally stages f16 copies of K and V laid out [head][token][dim] for + * the staged flash-attention prefill kernel. + * + * ponytail: the Metal decode_group4 variant is a perf-only specialization + * with identical math; the generic kernel covers decode too. Add it back if + * decode KV-build shows up in profiles. + * ------------------------------------------------------------------------- */ + +__global__ static void glm_build_kv_cache_kernel( + const float *kv_raw, + const float *k_nope, + const float *value, + char *key_cache, + char *value_cache, + __half *key_f16, /* optional staging, may be NULL */ + __half *value_f16, /* optional staging, may be NULL */ + 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, + uint32_t cache_f16, + 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; + if (token >= n_tokens || head >= n_head) return; + + const uint32_t nth = blockDim.x; + 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; + __half *kdst_f16 = key_f16 ? + key_f16 + ((uint64_t)head * n_tokens + token) * qk_dim : NULL; + __half *vdst_f16 = value_f16 ? + value_f16 + ((uint64_t)head * n_tokens + token) * value_dim : NULL; + + for (uint32_t i = threadIdx.x; i < qk_nope; i += nth) { + const float x = kn[i]; + glm_cache_store_f32_or_f16_dev(key_cache, kbase + i, cache_f16, x); + if (kdst_f16) kdst_f16[i] = __float2half(x); + } + + float corr0 = 0.0f, corr1 = 0.0f; + if (ext_factor != 0.0f) { + glm_rope_yarn_corr_dims_dev(qk_rope, n_ctx_orig, freq_base, + beta_fast, beta_slow, &corr0, &corr1); + } + const float theta_base = (float)pos; + const float inv_ndims = -1.0f / (float)qk_rope; + for (uint32_t r = threadIdx.x * 2u; r < qk_rope; r += nth * 2u) { + const float theta = theta_base * powf(freq_base, inv_ndims * (float)r); + float cos_theta, sin_theta; + glm_rope_yarn_dev(theta, freq_scale, corr0, corr1, (int)r, + ext_factor, attn_factor, &cos_theta, &sin_theta); + const uint32_t src0 = kv_lora_dim + r; + const float x0 = raw[src0]; + const float x1 = raw[src0 + 1u]; + const uint32_t dst0 = qk_nope + r; + const float y0 = x0 * cos_theta - x1 * sin_theta; + const float y1 = x0 * sin_theta + x1 * cos_theta; + glm_cache_store_f32_or_f16_dev(key_cache, kbase + dst0, cache_f16, y0); + glm_cache_store_f32_or_f16_dev(key_cache, kbase + dst0 + 1u, cache_f16, y1); + if (kdst_f16) { + kdst_f16[dst0] = __float2half(y0); + kdst_f16[dst0 + 1u] = __float2half(y1); + } + } + + for (uint32_t i = threadIdx.x; i < value_dim; i += nth) { + const float x = val[i]; + glm_cache_store_f32_or_f16_dev(value_cache, vbase + i, cache_f16, x); + if (vdst_f16) vdst_f16[i] = __float2half(x); + } +} + +/* Shared host-side validation for both build variants. */ +static int glm_build_kv_cache_args_ok( + const ds4_gpu_tensor *key_cache, + const 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, + 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; + if (!key_cache || !value_cache || !kv_raw || !k_nope || !value || + n_tokens == 0 || cache_cap == 0 || n_head == 0 || + kv_raw_dim == 0 || kv_lora_dim == 0 || + qk_nope == 0 || qk_rope == 0 || value_dim == 0 || + kv_lora_dim + qk_rope > kv_raw_dim || + qk_dim < qk_nope || (qk_rope & 1u) != 0 || + pos0 > cache_cap || n_tokens > cache_cap - pos0 || + !isfinite(freq_base) || freq_base <= 0.0f || + !isfinite(freq_scale) || freq_scale <= 0.0f || + !isfinite(ext_factor) || !isfinite(attn_factor) || + !isfinite(beta_fast) || !isfinite(beta_slow)) { + return 0; + } + const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float); + if (key_cache->bytes < (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes || + value_cache->bytes < (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes || + 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; + } + return 1; +} + +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) { + if (!glm_build_kv_cache_args_ok(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, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, cache_f16)) { + return 0; + } + dim3 grid(n_tokens, n_head, 1); + glm_build_kv_cache_kernel<<>>( + (const float *)kv_raw->ptr, (const float *)k_nope->ptr, + (const float *)value->ptr, + (char *)key_cache->ptr, (char *)value_cache->ptr, + NULL, NULL, + pos0, n_tokens, n_head, kv_raw_dim, kv_lora_dim, + qk_nope, qk_rope, value_dim, n_ctx_orig, cache_f16 ? 1u : 0u, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "glm build kv cache launch"); +} + +/* Persistent f16 K/V staging scratch for the staged flash-attention prefill + * path (CUDA mirror of Metal's g_flash_attn_kv_buffer). Layout: keys at + * offset 0 as [n_head][n_tokens][qk_dim] halves, values right after the keys + * as [n_head][n_tokens][value_dim] halves. + * TODO(port): ds4_gpu_glm_attention_flash_staged_tensor (attention cluster) + * must consume this same buffer; it is intentionally non-static so the + * attention .inc can reference it. */ +void *g_glm_flash_attn_kv_buffer = NULL; +uint64_t g_glm_flash_attn_kv_bytes = 0; + +static int glm_flash_attn_kv_scratch_ensure(uint64_t bytes) { + if (g_glm_flash_attn_kv_bytes >= bytes) return 1; + if (g_glm_flash_attn_kv_buffer) { + (void)cudaFree(g_glm_flash_attn_kv_buffer); + g_glm_flash_attn_kv_buffer = NULL; + g_glm_flash_attn_kv_bytes = 0; + } + void *ptr = NULL; + if (!cuda_ok(cudaMalloc(&ptr, (size_t)bytes), "glm flash kv scratch alloc")) { + return 0; + } + g_glm_flash_attn_kv_buffer = ptr; + g_glm_flash_attn_kv_bytes = bytes; + return 1; +} + +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 != 0) return 0; /* staged flash path is prefill-from-zero only */ + if (!glm_build_kv_cache_args_ok(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, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow, cache_f16)) { + return 0; + } + const uint32_t qk_dim = qk_nope + qk_rope; + const uint64_t key_f16_bytes = + (uint64_t)n_tokens * n_head * qk_dim * sizeof(uint16_t); + const uint64_t value_f16_bytes = + (uint64_t)n_tokens * n_head * value_dim * sizeof(uint16_t); + if (!glm_flash_attn_kv_scratch_ensure(key_f16_bytes + value_f16_bytes)) return 0; + __half *key_f16 = (__half *)g_glm_flash_attn_kv_buffer; + __half *value_f16 = (__half *)((char *)g_glm_flash_attn_kv_buffer + key_f16_bytes); + dim3 grid(n_tokens, n_head, 1); + glm_build_kv_cache_kernel<<>>( + (const float *)kv_raw->ptr, (const float *)k_nope->ptr, + (const float *)value->ptr, + (char *)key_cache->ptr, (char *)value_cache->ptr, + key_f16, value_f16, + pos0, n_tokens, n_head, kv_raw_dim, kv_lora_dim, + qk_nope, qk_rope, value_dim, n_ctx_orig, cache_f16 ? 1u : 0u, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + return cuda_ok(cudaGetLastError(), "glm build kv cache flash launch"); +} diff --git a/ds4_cuda_glm_moe.inc b/ds4_cuda_glm_moe.inc new file mode 100644 index 000000000..dbd1731bd --- /dev/null +++ b/ds4_cuda_glm_moe.inc @@ -0,0 +1,1518 @@ +/* CUDA port of the Metal GLM-5.2 MoE and general utility backend functions. + * + * This file is included at the end of ds4_cuda.cu and uses its internals + * (cuda_ok, cuda_model_range_ptr, cuda_model_copy_to_device_streamed, + * g_ssd_streaming_mode, g_stream_selected_cache, g_stream_expert_cache, ...). + * + * Ports: + * - kernel_glm_router_select_one (metal/dsv4_misc.metal) + * - kernel_glm_qX_K_pair_swiglu / down (metal/moe.metal, scalar variants) + * - kernel_add3_f32, kernel_add_rms_norm_mul_f32_4, kernel_dsv4_sort_i32_rows_asc + * - streaming full-layer expert cache (rocm/ds4_rocm_runtime.cuh design, + * rebuilt on ds4_cuda.cu's own streamed-copy + expert-cache machinery) + */ + +/* ========================================================================= + * K-quant block layouts missing from ds4_cuda.cu (q2/q4 already exist). + * ========================================================================= */ + +typedef struct { + uint16_t d; + uint16_t dmin; + uint8_t scales[12]; + uint8_t qh[CUDA_QK_K / 8]; + uint8_t qs[CUDA_QK_K / 2]; +} cuda_block_q5_K; + +typedef struct { + uint8_t ql[CUDA_QK_K / 2]; + uint8_t qh[CUDA_QK_K / 4]; + int8_t scales[CUDA_QK_K / 16]; + uint16_t d; +} cuda_block_q6_K; + +/* ========================================================================= + * Per-element dequantization (faithful port of metal/moe.metal + * ds4_glm_qX_K_value helpers). + * ========================================================================= */ + +__device__ __forceinline__ static float glm_half_bits_to_f32(uint16_t h) { + return __half2float(__ushort_as_half(h)); +} + +__device__ __forceinline__ static void glm_get_scale_min_k4_just2( + uint32_t j, const uint8_t *q, uint32_t *sc, uint32_t *mn) { + if (j < 4u) { + *sc = q[j] & 63u; + *mn = q[j + 4u] & 63u; + } else { + *sc = (q[j + 4u] & 0x0Fu) | ((q[j - 4u] & 0xC0u) >> 2u); + *mn = (q[j + 4u] >> 4u) | ((q[j] & 0xC0u) >> 2u); + } +} + +struct glm_deq_q2_K { + __device__ __forceinline__ static float value(const char *blocks, uint32_t k) { + const cuda_block_q2_K *xb = + (const cuda_block_q2_K *)blocks + (k / CUDA_QK_K); + const uint32_t idx = k & (CUDA_QK_K - 1u); + const uint32_t group = idx / 16u; + const uint32_t l = idx & 15u; + const uint32_t q_base = 32u * (group / 8u) + 16u * (group & 1u); + const uint32_t shift = ((group / 2u) & 3u) * 2u; + const uint32_t q = ((uint32_t)xb->qs[q_base + l] >> shift) & 0x03u; + const uint32_t sc = (uint32_t)xb->scales[group]; + return glm_half_bits_to_f32(xb->d) * (float)(sc & 0x0Fu) * (float)q - + glm_half_bits_to_f32(xb->dmin) * (float)(sc >> 4u); + } +}; + +/* IQ2_XXS scalar dequant. Layout: {half d; uint16 qs[QK_K/8];}; each 32-element + * sub-block packs 4 grid indices (aux0) + a 4-bit scale and four 7-bit sign + * indices (aux1). value = 0.125 * d * ls * grid[j] * sign, matching the dp4a + * path dev_dot_iq2_xxs_q8_K_block_lut. Uses ds4_cuda.cu's cuda_iq2xxs_grid / + * cuda_ksigns_iq2xs tables (in scope: this file is included at end of ds4_cuda.cu). */ +struct glm_deq_iq2_xxs { + __device__ __forceinline__ static float value(const char *blocks, uint32_t k) { + const cuda_block_iq2_xxs *xb = + (const cuda_block_iq2_xxs *)blocks + (k >> 8); + const uint32_t within = k & 255u; + const uint32_t ib32 = within >> 5; + const uint32_t kk = within & 31u; + const uint32_t l = kk >> 3; + const uint32_t j = kk & 7u; + const uint16_t *q2 = xb->qs + 4u * 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); + const uint32_t ls = 2u * (aux1 >> 28) + 1u; + const uint32_t grid_idx = (aux0 >> (8u * l)) & 0xffu; + const uint32_t sign_idx = (aux1 >> (7u * l)) & 127u; + const uint64_t g = cuda_iq2xxs_grid[grid_idx]; + const int32_t gj = (int32_t)(uint32_t)(uint8_t)(g >> (8u * j)); + const uint8_t signs = cuda_ksigns_iq2xs[sign_idx]; + const float sign = (signs & (1u << j)) ? -1.0f : 1.0f; + return 0.125f * dev_f16_to_f32(xb->d) * (float)ls * (float)gj * sign; + } +}; + +struct glm_deq_q4_K { + __device__ __forceinline__ static float value(const char *blocks, uint32_t k) { + const cuda_block_q4_K *xb = + (const cuda_block_q4_K *)blocks + (k / CUDA_QK_K); + const uint32_t idx = k & (CUDA_QK_K - 1u); + const uint32_t group = idx / 32u; + const uint32_t l = idx & 31u; + uint32_t sc, mn; + glm_get_scale_min_k4_just2(group, xb->scales, &sc, &mn); + const uint32_t byte_off = (group >> 1u) * 32u + l; + const uint32_t shift = (group & 1u) * 4u; + const uint32_t q = ((uint32_t)xb->qs[byte_off] >> shift) & 0x0Fu; + return glm_half_bits_to_f32(xb->d) * (float)sc * (float)q - + glm_half_bits_to_f32(xb->dmin) * (float)mn; + } +}; + +struct glm_deq_q5_K { + __device__ __forceinline__ static float value(const char *blocks, uint32_t k) { + const cuda_block_q5_K *xb = + (const cuda_block_q5_K *)blocks + (k / CUDA_QK_K); + const uint32_t idx = k & (CUDA_QK_K - 1u); + const uint32_t group = idx / 32u; + const uint32_t l = idx & 31u; + uint32_t sc, mn; + glm_get_scale_min_k4_just2(group, xb->scales, &sc, &mn); + const uint32_t ql_base = (group >> 1u) * 32u + l; + const uint32_t shift = (group & 1u) * 4u; + uint32_t q = ((uint32_t)xb->qs[ql_base] >> shift) & 0x0Fu; + q += (xb->qh[l] & (uint8_t)(1u << group)) ? 16u : 0u; + return glm_half_bits_to_f32(xb->d) * (float)sc * (float)q - + glm_half_bits_to_f32(xb->dmin) * (float)mn; + } +}; + +struct glm_deq_q6_K { + __device__ __forceinline__ static float value(const char *blocks, uint32_t k) { + const cuda_block_q6_K *xb = + (const cuda_block_q6_K *)blocks + (k / CUDA_QK_K); + const uint32_t idx = k & (CUDA_QK_K - 1u); + const uint32_t n128 = idx >> 7u; + const uint32_t r = idx & 127u; + const uint32_t l = r & 31u; + const uint32_t quarter = r >> 5u; + const uint32_t ql_base = n128 * 64u; + const uint32_t qh_base = n128 * 32u; + const uint32_t sc_base = n128 * 8u; + uint32_t q; + int sc; + if (quarter == 0u) { + q = (xb->ql[ql_base + l] & 0x0Fu) | + (((xb->qh[qh_base + l] >> 0u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 0u]; + } else if (quarter == 1u) { + q = (xb->ql[ql_base + 32u + l] & 0x0Fu) | + (((xb->qh[qh_base + l] >> 2u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 2u]; + } else if (quarter == 2u) { + q = (xb->ql[ql_base + l] >> 4u) | + (((xb->qh[qh_base + l] >> 4u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 4u]; + } else { + q = (xb->ql[ql_base + 32u + l] >> 4u) | + (((xb->qh[qh_base + l] >> 6u) & 3u) << 4u); + sc = (int)xb->scales[sc_base + l / 16u + 6u]; + } + return glm_half_bits_to_f32(xb->d) * (float)sc * (float)((int)q - 32); + } +}; + +/* Tensor type ids (GGUF), mirrors DS4_METAL_TENSOR_* in ds4_metal.m. */ +enum { + GLM_CUDA_TENSOR_Q2_K = 10, + GLM_CUDA_TENSOR_Q4_K = 12, + GLM_CUDA_TENSOR_Q5_K = 13, + GLM_CUDA_TENSOR_Q6_K = 14, + GLM_CUDA_TENSOR_IQ2_XXS = 16 +}; + +/* ========================================================================= + * Small utility kernels. + * ========================================================================= */ + +__global__ static void glm_add3_kernel( + float *out, const float *a, const float *b, const float *c, uint32_t n) { + const uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + out[i] = a[i] + b[i] + c[i]; +} + +/* sum_out = a + b; norm_out = rms_norm(sum_out) * w. Single row, mirrors + * kernel_add_rms_norm_mul_f32_4 in metal/norm.metal (rows == 1 dispatch). */ +__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) { + float sum = 0.0f; + for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { + const float v = a[i] + b[i]; + sum_out[i] = v; + sum += v * v; + } + __shared__ float partial[256]; + partial[threadIdx.x] = sum; + __syncthreads(); + for (uint32_t stride = blockDim.x >> 1u; stride > 0u; stride >>= 1u) { + 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) { + norm_out[i] = sum_out[i] * scale * w[i]; + } +} + +/* Per-row ascending bitonic sort of int32 ids, port of + * kernel_dsv4_sort_i32_rows_asc (metal/dsv4_misc.metal). row_width must be a + * power of two (enforced by the launcher, as on Metal). */ +__global__ static void glm_sort_i32_rows_asc_kernel( + int32_t *dst, const int32_t *src, uint32_t row_width, uint32_t n_rows) { + extern __shared__ int32_t glm_sort_row_tmp[]; + const uint32_t row = blockIdx.x; + if (row >= n_rows) return; + const uint32_t n_threads = blockDim.x; + const uint32_t tid = threadIdx.x; + const uint64_t base = (uint64_t)row * row_width; + + for (uint32_t i = tid; i < row_width; i += n_threads) { + glm_sort_row_tmp[i] = src[base + i]; + } + __syncthreads(); + + for (uint32_t k = 2u; k <= row_width; k <<= 1u) { + for (uint32_t j = k >> 1u; j > 0u; j >>= 1u) { + for (uint32_t i = tid; i < row_width; i += n_threads) { + const uint32_t other = i ^ j; + if (other > i && other < row_width) { + const int32_t va = glm_sort_row_tmp[i]; + const int32_t vb = glm_sort_row_tmp[other]; + const bool up = (i & k) == 0u; + if ((up && va > vb) || (!up && va < vb)) { + glm_sort_row_tmp[i] = vb; + glm_sort_row_tmp[other] = va; + } + } + } + __syncthreads(); + } + } + + for (uint32_t i = tid; i < row_width; i += n_threads) { + dst[base + i] = glm_sort_row_tmp[i]; + } +} + +/* ========================================================================= + * GLM router: sigmoid probabilities, selection score = prob + bias, top-k, + * route weights = selected prob / sum(selected probs) * scale. + * Port of kernel_glm_router_select_one (metal/dsv4_misc.metal). + * One warp per token; n_expert <= 256, k_used <= n_expert. + * ========================================================================= */ + +__device__ __forceinline__ static float glm_router_sigmoid_dev(float x) { + if (x >= 0.0f) { + const float e = expf(-x); + return 1.0f / (1.0f + e); + } + const float e = expf(x); + return e / (1.0f + e); +} + +__device__ __forceinline__ static bool glm_router_better_dev( + float av, uint32_t ai, float bv, uint32_t bi) { + return av > bv || (av == bv && ai < bi); +} + +__global__ static void glm_router_select_kernel( + int32_t *selected, + float *weights, + float *probs, + const float *bias, + const float *logits, + uint32_t n_expert, + uint32_t k_used, + float expert_weight_scale, + uint32_t n_tokens) { + const uint32_t lane = threadIdx.x; + const uint32_t token = blockIdx.x * blockDim.y + threadIdx.y; + if (token >= n_tokens || lane >= 32u) return; + + const float *log = logits + (uint64_t)token * n_expert; + float *prob = probs + (uint64_t)token * n_expert; + int32_t *sel = selected + (uint64_t)token * k_used; + float *w = weights + (uint64_t)token * k_used; + + float local_prob[8]; + float local_score[8]; + #pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const uint32_t e = lane + j * 32u; + if (e < n_expert) { + const float p = glm_router_sigmoid_dev(log[e]); + local_prob[j] = p; + local_score[j] = p + bias[e]; + prob[e] = p; + } else { + local_prob[j] = 0.0f; + local_score[j] = -INFINITY; + } + } + __syncwarp(); + + float sum = 0.0f; + for (uint32_t k = 0; k < k_used; k++) { + float best_score = -INFINITY; + float best_prob = 0.0f; + uint32_t best_idx = UINT32_MAX; + #pragma unroll + for (uint32_t j = 0; j < 8u; j++) { + const uint32_t e = lane + j * 32u; + if (glm_router_better_dev(local_score[j], e, best_score, best_idx)) { + best_score = local_score[j]; + 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 (glm_router_better_dev(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 < 8u; j++) { + if (lane + j * 32u == best_idx) local_score[j] = -INFINITY; + } + if (lane == 0) { + sel[k] = (int32_t)best_idx; + w[k] = best_prob; + } + sum += best_prob; + } + + if (lane == 0) { + sum = fmaxf(sum, 6.103515625e-5f); + for (uint32_t k = 0; k < k_used; k++) { + w[k] = w[k] / sum * expert_weight_scale; + } + } +} + +static int glm_router_select_launch( + 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_tokens == 0 || + n_expert == 0 || n_expert > 256u || + n_expert_used == 0 || n_expert_used > n_expert) { + return 0; + } + if (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)) { + fprintf(stderr, "ds4: CUDA GLM router received undersized buffers\n"); + return 0; + } + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (bias_offset > model_size || bias_bytes > model_size - bias_offset) { + fprintf(stderr, "ds4: CUDA GLM router bias range is outside the mapped model\n"); + return 0; + } + const float *bias = (const float *)cuda_model_range_ptr(model_map, + bias_offset, bias_bytes, "glm_router_bias"); + if (!bias) return 0; + + dim3 block(32, 4, 1); + glm_router_select_kernel<<<(n_tokens + 3u) / 4u, block>>>( + (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 launch"); +} + +/* ========================================================================= + * GLM routed MoE kernels. + * + * pair swiglu: mid[token][slot][row] = + * silu(gate_row . x) * (up_row . x) * route_weight[token][slot] + * down: out[token][row] = sum_slot(down_row . mid[token][slot]) + * Route weights are folded into mid, exactly as in metal/moe.metal. + * One warp per output row. + * ========================================================================= */ + +template +__global__ static void glm_moe_pair_swiglu_kernel( + float *mid, + const char *gate, + const char *up, + const float *x, + const int32_t *selected, + const float *weights, + uint32_t in_dim, + uint32_t mid_dim, + uint32_t n_total_expert, + uint32_t n_expert_used, + uint32_t n_tokens, + uint32_t mid_token_stride, + uint64_t gate_expert_bytes, + uint64_t gate_row_bytes, + uint64_t up_expert_bytes, + uint64_t up_row_bytes) { + const uint32_t lane = threadIdx.x; + const uint32_t row = blockIdx.x * blockDim.y + threadIdx.y; + const uint32_t slot = blockIdx.y; + const uint32_t token = blockIdx.z; + if (row >= mid_dim || slot >= n_expert_used || token >= n_tokens) return; + + const uint64_t selected_off = (uint64_t)token * n_expert_used + slot; + const uint64_t mid_off = (uint64_t)token * mid_token_stride + + (uint64_t)slot * mid_dim + row; + const int expert = selected[selected_off]; + if (expert < 0 || (uint32_t)expert >= n_total_expert) { + if (lane == 0) mid[mid_off] = 0.0f; + return; + } + + const char *gate_row = gate + + (uint64_t)(uint32_t)expert * gate_expert_bytes + + (uint64_t)row * gate_row_bytes; + const char *up_row = up + + (uint64_t)(uint32_t)expert * up_expert_bytes + + (uint64_t)row * up_row_bytes; + const float *token_x = x + (uint64_t)token * in_dim; + + float acc_gate = 0.0f; + float acc_up = 0.0f; + for (uint32_t k = lane; k < in_dim; k += 32u) { + const float xv = token_x[k]; + acc_gate += DEQ::value(gate_row, k) * xv; + acc_up += DEQ::value(up_row, k) * xv; + } + #pragma unroll + for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) { + acc_gate += __shfl_xor_sync(0xffffffffu, acc_gate, mask); + acc_up += __shfl_xor_sync(0xffffffffu, acc_up, mask); + } + if (lane == 0) { + const float g = acc_gate; + const float sw = g / (1.0f + expf(-g)); + mid[mid_off] = sw * acc_up * weights[selected_off]; + } +} + +template +__global__ static void glm_moe_down_kernel( + float *out, + const char *down, + const int32_t *selected, + const float *mid, + uint32_t mid_dim, + uint32_t out_dim, + uint32_t n_total_expert, + uint32_t n_expert_used, + uint32_t n_tokens, + uint32_t mid_token_stride, + uint64_t down_expert_bytes, + uint64_t down_row_bytes) { + const uint32_t lane = threadIdx.x; + const uint32_t row = blockIdx.x * blockDim.y + threadIdx.y; + const uint32_t token = blockIdx.y; + if (row >= out_dim || token >= n_tokens) return; + + const uint64_t selected_base = (uint64_t)token * n_expert_used; + const uint64_t mid_base = (uint64_t)token * mid_token_stride; + float acc = 0.0f; + for (uint32_t slot = 0; slot < n_expert_used; slot++) { + const int expert = selected[selected_base + slot]; + if (expert < 0 || (uint32_t)expert >= n_total_expert) continue; + const char *down_row = down + + (uint64_t)(uint32_t)expert * down_expert_bytes + + (uint64_t)row * down_row_bytes; + const float *slot_mid = mid + mid_base + (uint64_t)slot * mid_dim; + for (uint32_t k = lane; k < mid_dim; k += 32u) { + acc += DEQ::value(down_row, k) * slot_mid[k]; + } + } + #pragma unroll + for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) { + acc += __shfl_xor_sync(0xffffffffu, acc, mask); + } + if (lane == 0) { + out[(uint64_t)token * out_dim + row] = acc; + } +} + +static bool glm_cuda_gate_pair_type_supported(uint32_t gate_type, uint32_t up_type) { + return gate_type == up_type && + (gate_type == GLM_CUDA_TENSOR_Q2_K || + gate_type == GLM_CUDA_TENSOR_Q4_K || + gate_type == GLM_CUDA_TENSOR_Q5_K || + gate_type == GLM_CUDA_TENSOR_IQ2_XXS); +} + +static bool glm_cuda_down_type_supported(uint32_t down_type) { + return down_type == GLM_CUDA_TENSOR_IQ2_XXS || + down_type == GLM_CUDA_TENSOR_Q2_K || + down_type == GLM_CUDA_TENSOR_Q4_K || + down_type == GLM_CUDA_TENSOR_Q5_K || + down_type == GLM_CUDA_TENSOR_Q6_K; +} + +/* ========================================================================= + * Streaming full-layer expert cache (GLM prefill). Two ping-pong slots so + * the next layer can load while the current one computes; port of the ROCm + * design in rocm/ds4_rocm_runtime.cuh rebuilt on ds4_cuda.cu's streamed copy. + * ========================================================================= */ + +struct glm_stream_layer_expert_cache { + int active; + const void *model_map; + uint32_t layer; + uint32_t n_total_expert; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + uint64_t capacity; + char *base; + char *gate; + char *up; + char *down; +}; + +static glm_stream_layer_expert_cache g_glm_stream_layer_expert_cache[2]; + +static void glm_stream_layer_expert_cache_release(void) { + for (uint32_t i = 0; i < 2u; i++) { + glm_stream_layer_expert_cache &c = g_glm_stream_layer_expert_cache[i]; + if (c.base) (void)cudaFree(c.base); + memset(&c, 0, sizeof(c)); + } +} + +static int glm_stream_layer_expert_cache_apply( + const void *model_map, + uint32_t layer, + uint32_t n_total_expert, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + const char **gate_w, + const char **up_w, + const char **down_w) { + if (!g_ssd_streaming_mode || !gate_w || !up_w || !down_w) return 0; + for (uint32_t i = 0; i < 2u; i++) { + const glm_stream_layer_expert_cache &c = g_glm_stream_layer_expert_cache[i]; + if (c.active && + c.model_map == model_map && + c.layer == layer && + c.n_total_expert == n_total_expert && + c.gate_offset == gate_offset && + c.up_offset == up_offset && + c.down_offset == down_offset && + c.gate_expert_bytes == gate_expert_bytes && + c.down_expert_bytes == down_expert_bytes && + c.gate && c.up && c.down) { + *gate_w = c.gate; + *up_w = c.up; + *down_w = c.down; + return 1; + } + } + return 0; +} + +extern "C" int ds4_gpu_stream_expert_cache_load_layer( + const ds4_gpu_stream_expert_table *table) { + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + + if (!g_ssd_streaming_mode || !model_map || model_size == 0 || + n_total_expert == 0 || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) { + return 0; + } + 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 > UINT64_MAX / 2u || + 2u * gate_bytes > UINT64_MAX - down_bytes) { + return 0; + } + const uint64_t total_bytes = 2u * gate_bytes + down_bytes; + if (gate_offset > model_size || up_offset > model_size || + down_offset > model_size || + gate_bytes > model_size - gate_offset || + gate_bytes > model_size - up_offset || + down_bytes > model_size - down_offset) { + return 0; + } + + glm_stream_layer_expert_cache &slot = + g_glm_stream_layer_expert_cache[layer & 1u]; + slot.active = 0; + if (slot.capacity < total_bytes) { + if (slot.base) { + (void)cudaFree(slot.base); + memset(&slot, 0, sizeof(slot)); + } + void *base = NULL; + cudaError_t err = cudaMalloc(&base, (size_t)total_bytes); + if (err != cudaSuccess) { + /* Try again after releasing the per-expert resident cache. */ + (void)cudaGetLastError(); + cuda_stream_expert_cache_release_all(); + err = cudaMalloc(&base, (size_t)total_bytes); + } + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming full-layer expert cache allocation " + "failed for layer=%u (%.2f GiB): %s\n", + layer, + (double)total_bytes / 1073741824.0, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + slot.base = (char *)base; + slot.capacity = total_bytes; + } + + slot.gate = slot.base; + slot.up = slot.base + gate_bytes; + slot.down = slot.base + 2u * gate_bytes; + if (!cuda_model_copy_to_device_streamed(slot.gate, model_map, model_size, + gate_offset, gate_bytes, + "layer moe_gate") || + !cuda_model_copy_to_device_streamed(slot.up, model_map, model_size, + up_offset, gate_bytes, + "layer moe_up") || + !cuda_model_copy_to_device_streamed(slot.down, model_map, model_size, + down_offset, down_bytes, + "layer moe_down")) { + return 0; + } + + slot.active = 1; + slot.model_map = model_map; + slot.layer = layer; + slot.n_total_expert = n_total_expert; + slot.gate_offset = gate_offset; + slot.up_offset = up_offset; + slot.down_offset = down_offset; + slot.gate_expert_bytes = gate_expert_bytes; + slot.down_expert_bytes = down_expert_bytes; + return 1; +} + +extern "C" int ds4_gpu_stream_expert_cache_release_layer_cache(void) { + glm_stream_layer_expert_cache_release(); + return 1; +} + +/* Seeds the per-expert resident streaming cache from a resident full-layer + * cache using the LAST n_seed_tokens tokens of the batch-selected ids + * (device tensor, n_tokens x n_selected int32), so decode starts warm. */ +extern "C" int ds4_gpu_stream_expert_cache_seed_from_layer_selected( + const ds4_gpu_stream_expert_table *table, + const ds4_gpu_tensor *selected, + uint32_t n_tokens, + uint32_t n_seed_tokens, + uint32_t n_selected) { + if (!table) return 0; + const void *model_map = table->model_map; + const uint64_t model_size = table->model_size; + const uint32_t layer = table->layer; + const uint32_t n_total_expert = table->n_total_expert; + const uint64_t gate_offset = table->gate_offset; + const uint64_t up_offset = table->up_offset; + const uint64_t down_offset = table->down_offset; + const uint64_t gate_expert_bytes = table->gate_expert_bytes; + const uint64_t down_expert_bytes = table->down_expert_bytes; + + if (!g_ssd_streaming_mode || !model_map || !selected || !selected->ptr || + n_tokens == 0 || n_seed_tokens == 0 || + n_total_expert == 0 || n_selected == 0 || + gate_expert_bytes == 0 || down_expert_bytes == 0 || + (uint64_t)n_tokens > UINT64_MAX / n_selected || + selected->bytes < (uint64_t)n_tokens * n_selected * sizeof(int32_t)) { + return 0; + } + + const char *layer_gate = NULL; + const char *layer_up = NULL; + const char *layer_down = NULL; + if (!glm_stream_layer_expert_cache_apply(model_map, layer, n_total_expert, + gate_offset, up_offset, down_offset, + gate_expert_bytes, down_expert_bytes, + &layer_gate, &layer_up, &layer_down)) { + return 0; + } + + if (n_seed_tokens > n_tokens) n_seed_tokens = n_tokens; + const uint64_t n_ids = (uint64_t)n_seed_tokens * n_selected; + std::vector ids((size_t)n_ids); + const uint64_t src_off = + (uint64_t)(n_tokens - n_seed_tokens) * n_selected * sizeof(int32_t); + if (!cuda_ok(cudaMemcpy(ids.data(), + (const char *)selected->ptr + src_off, + (size_t)n_ids * sizeof(int32_t), + cudaMemcpyDeviceToHost), + "glm layer seed selected ids copy")) { + return 0; + } + + std::vector seen(n_total_expert, 0); + std::vector unique; + unique.reserve((size_t)n_ids); + for (uint64_t i = 0; i < n_ids; i++) { + const int32_t expert = ids[(size_t)i]; + if (expert < 0 || (uint32_t)expert >= n_total_expert) { + fprintf(stderr, + "ds4: CUDA streaming full-layer seed expert id %d " + "outside 0..%u (layer=%u)\n", + expert, n_total_expert, layer); + return 0; + } + if (seen[(uint32_t)expert]) continue; + seen[(uint32_t)expert] = 1; + unique.push_back(expert); + } + if (unique.empty()) return 1; + + cuda_stream_expert_cache *cache = + cuda_stream_expert_cache_prepare(gate_expert_bytes, + down_expert_bytes, + 0); + if (!cache || !cache->valid) return 0; + + for (size_t u = 0; u < unique.size(); u++) { + const uint32_t expert = (uint32_t)unique[u]; + int cache_slot = cuda_stream_expert_cache_find(cache, + model_map, + model_size, + layer, + n_total_expert, + expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes); + if (cache_slot >= 0) { + cache->slots[(uint32_t)cache_slot].age = ++cache->tick; + continue; + } + const uint32_t load_slot = cuda_stream_expert_cache_lru_slot(cache); + const int append = !cache->slots[load_slot].valid; + const uint64_t gate_src = (uint64_t)expert * gate_expert_bytes; + const uint64_t down_src = (uint64_t)expert * down_expert_bytes; + const uint64_t gate_dst = (uint64_t)load_slot * gate_expert_bytes; + const uint64_t down_dst = (uint64_t)load_slot * down_expert_bytes; + if (!cuda_ok(cudaMemcpy(cache->gate_ptr + gate_dst, + layer_gate + gate_src, + (size_t)gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "glm layer seed gate copy") || + !cuda_ok(cudaMemcpy(cache->up_ptr + gate_dst, + layer_up + gate_src, + (size_t)gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "glm layer seed up copy") || + !cuda_ok(cudaMemcpy(cache->down_ptr + down_dst, + layer_down + down_src, + (size_t)down_expert_bytes, + cudaMemcpyDeviceToDevice), + "glm layer seed down copy")) { + return 0; + } + cuda_stream_expert_cache_slot &entry = cache->slots[load_slot]; + entry.valid = 1; + entry.model_map = model_map; + entry.model_size = model_size; + entry.layer = layer; + entry.n_total_expert = n_total_expert; + entry.expert = expert; + entry.gate_offset = gate_offset; + entry.up_offset = up_offset; + entry.down_offset = down_offset; + entry.gate_expert_bytes = gate_expert_bytes; + entry.down_expert_bytes = down_expert_bytes; + entry.age = ++cache->tick; + if (append && cache->count < cache->capacity) cache->count++; + } + return 1; +} + +/* Device-tensor variant of the early selected-expert load: reads the router's + * selected ids back to the host and delegates to the existing host-id loader + * that fills g_stream_selected_cache. Port of the Metal launcher at + * ds4_metal.m:12595 (no command-batch juggling needed on CUDA: the stream is + * in-order and cudaMemcpy synchronizes). */ +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 || + getenv("DS4_CUDA_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL || + getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) { + return 1; + } + if (!table || !selected || n_selected == 0 || n_selected > 8u) { + return 1; + } + /* Full layer already resident: nothing to load. */ + const char *g = NULL, *u = NULL, *d = NULL; + if (glm_stream_layer_expert_cache_apply(table->model_map, + table->layer, + table->n_total_expert, + table->gate_offset, + table->up_offset, + table->down_offset, + table->gate_expert_bytes, + table->down_expert_bytes, + &g, &u, &d)) { + return 1; + } + + int32_t selected_ids[8] = {-1, -1, -1, -1, -1, -1, -1, -1}; + if (!ds4_gpu_tensor_read(selected, 0, selected_ids, + (uint64_t)n_selected * sizeof(selected_ids[0]))) { + return 0; + } + return ds4_gpu_stream_expert_cache_begin_selected_load(table, + selected_ids, + n_selected); +} + +/* ========================================================================= + * GLM routed MoE launcher (decode + batch). + * + * Note vs the Metal launcher (ds4_metal.m:28085): the Metal-side streaming + * bookkeeping (note_decode_token at layer 3, selected-hotness notes, + * per-layer/global pruning, split resident/missing dispatch) belongs to the + * Metal per-expert buffer cache and has no CUDA counterpart; on CUDA the + * equivalent residency management lives inside + * ds4_gpu_stream_expert_cache_begin_selected_load / g_stream_expert_cache LRU. + * ========================================================================= */ + +static int glm_routed_moe_launch( + 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, + bool force_resident) { + if (!out || !mid || !model_map || !selected || !weights || !x || + n_tokens == 0 || n_tokens > 65535u || + n_total_expert == 0 || n_expert == 0 || n_expert > 256u || + n_expert > n_total_expert || + expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 || + gate_expert_bytes == 0 || gate_row_bytes == 0 || + up_expert_bytes == 0 || up_row_bytes == 0 || + down_expert_bytes == 0 || down_row_bytes == 0 || + (expert_in_dim % CUDA_QK_K) != 0 || + (expert_mid_dim % CUDA_QK_K) != 0 || + !glm_cuda_gate_pair_type_supported(gate_type, up_type) || + !glm_cuda_down_type_supported(down_type)) { + return 0; + } + const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim; + if ((uint64_t)mid_token_stride < per_token_mid || + (uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes || + (uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes || + gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes || + up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes || + down_expert_bytes != (uint64_t)out_dim * down_row_bytes) { + fprintf(stderr, "ds4: CUDA GLM routed MoE received inconsistent expert strides\n"); + return 0; + } + const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t up_tensor_bytes = (uint64_t)n_total_expert * up_expert_bytes; + const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes; + if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset || + up_offset > model_size || up_tensor_bytes > model_size - up_offset || + down_offset > model_size || down_tensor_bytes > model_size - down_offset) { + fprintf(stderr, "ds4: CUDA GLM routed MoE tensor range is outside the mapped model\n"); + return 0; + } + const uint64_t mid_values = + (uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid; + const uint64_t slot_values = (uint64_t)n_tokens * n_expert; + if (x->bytes < (uint64_t)n_tokens * expert_in_dim * sizeof(float) || + mid->bytes < mid_values * sizeof(float) || + out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float) || + selected->bytes < slot_values * sizeof(int32_t) || + weights->bytes < slot_values * sizeof(float)) { + fprintf(stderr, "ds4: CUDA GLM routed MoE received undersized activation buffers\n"); + return 0; + } + + /* Weight source resolution: + * 1. compact selected-expert streaming cache (this layer's routed + * experts already on device; selected ids remapped to compact slots) + * 2. resident full-layer streaming cache (prefill) + * 3. resident/mapped full expert tables. */ + const ds4_gpu_tensor *selected_exec = selected; + const char *gate_w = NULL; + const char *up_w = NULL; + const char *down_w = NULL; + const int streaming = g_ssd_streaming_mode && !force_resident; + + /* The shared streaming caches store one gate/up stride; GLM gate and up + * always share type and dims so the strides match by construction. */ + const int uniform_gate_up = gate_expert_bytes == up_expert_bytes; + + int use_selected_cache = 0; + for (int attempt = 0; attempt < 2 && !use_selected_cache; attempt++) { + use_selected_cache = + streaming && uniform_gate_up && + g_stream_selected_cache.valid && + 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_values && + 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_values * sizeof(int32_t); + if (use_selected_cache || attempt == 1 || !streaming || !uniform_gate_up) break; + /* Safety net: decode reached the MoE without an early selected load + * (e.g. the cache was invalidated). Load the selected experts now + * instead of pulling the full expert tables through the map. */ + const char *lg = NULL, *lu = NULL, *ld = NULL; + if (n_tokens != 1u || n_expert > 8u || + glm_stream_layer_expert_cache_apply(model_map, layer_index, + n_total_expert, + gate_offset, up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &lg, &lu, &ld)) { + break; + } + int32_t ids[8]; + if (!ds4_gpu_tensor_read(selected, 0, ids, + (uint64_t)n_expert * sizeof(ids[0]))) { + return 0; + } + 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 (!ds4_gpu_stream_expert_cache_begin_selected_load(&table, ids, n_expert)) { + return 0; + } + } + + if (use_selected_cache) { + selected_exec = &g_stream_selected_cache.slot_selected_tensor; + gate_w = g_stream_selected_cache.gate_ptr; + up_w = g_stream_selected_cache.up_ptr; + down_w = g_stream_selected_cache.down_ptr; + } else if (streaming && uniform_gate_up && + glm_stream_layer_expert_cache_apply(model_map, layer_index, + n_total_expert, + gate_offset, up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes, + &gate_w, &up_w, &down_w)) { + /* full-layer cache hit; original selected ids remain valid */ + } else { + gate_w = cuda_model_range_ptr(model_map, gate_offset, gate_tensor_bytes, + "glm_moe_gate"); + up_w = cuda_model_range_ptr(model_map, up_offset, up_tensor_bytes, + "glm_moe_up"); + down_w = cuda_model_range_ptr(model_map, down_offset, down_tensor_bytes, + "glm_moe_down"); + } + if (!gate_w || !up_w || !down_w) return 0; + + const uint32_t rows_per_block = 4u; + dim3 block(32, rows_per_block, 1); + dim3 pair_grid((expert_mid_dim + rows_per_block - 1u) / rows_per_block, + n_expert, + n_tokens); + dim3 down_grid((out_dim + rows_per_block - 1u) / rows_per_block, + n_tokens, + 1); + const int32_t *sel_ptr = (const int32_t *)selected_exec->ptr; + const float *w_ptr = (const float *)weights->ptr; + const float *x_ptr = (const float *)x->ptr; + float *mid_ptr = (float *)mid->ptr; + float *out_ptr = (float *)out->ptr; + + switch (gate_type) { + case GLM_CUDA_TENSOR_IQ2_XXS: + glm_moe_pair_swiglu_kernel<<>>( + mid_ptr, gate_w, up_w, x_ptr, sel_ptr, w_ptr, + expert_in_dim, expert_mid_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes); + break; + case GLM_CUDA_TENSOR_Q2_K: + glm_moe_pair_swiglu_kernel<<>>( + mid_ptr, gate_w, up_w, x_ptr, sel_ptr, w_ptr, + expert_in_dim, expert_mid_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes); + break; + case GLM_CUDA_TENSOR_Q4_K: + glm_moe_pair_swiglu_kernel<<>>( + mid_ptr, gate_w, up_w, x_ptr, sel_ptr, w_ptr, + expert_in_dim, expert_mid_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes); + break; + default: /* Q5_K, validated above */ + glm_moe_pair_swiglu_kernel<<>>( + mid_ptr, gate_w, up_w, x_ptr, sel_ptr, w_ptr, + expert_in_dim, expert_mid_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, + gate_expert_bytes, gate_row_bytes, up_expert_bytes, up_row_bytes); + break; + } + if (!cuda_ok(cudaGetLastError(), "glm routed moe pair launch")) return 0; + + switch (down_type) { + case GLM_CUDA_TENSOR_IQ2_XXS: + glm_moe_down_kernel<<>>( + out_ptr, down_w, sel_ptr, mid_ptr, + expert_mid_dim, out_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, down_expert_bytes, down_row_bytes); + break; + case GLM_CUDA_TENSOR_Q2_K: + glm_moe_down_kernel<<>>( + out_ptr, down_w, sel_ptr, mid_ptr, + expert_mid_dim, out_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, down_expert_bytes, down_row_bytes); + break; + case GLM_CUDA_TENSOR_Q4_K: + glm_moe_down_kernel<<>>( + out_ptr, down_w, sel_ptr, mid_ptr, + expert_mid_dim, out_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, down_expert_bytes, down_row_bytes); + break; + case GLM_CUDA_TENSOR_Q5_K: + glm_moe_down_kernel<<>>( + out_ptr, down_w, sel_ptr, mid_ptr, + expert_mid_dim, out_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, down_expert_bytes, down_row_bytes); + break; + default: /* Q6_K, validated above */ + glm_moe_down_kernel<<>>( + out_ptr, down_w, sel_ptr, mid_ptr, + expert_mid_dim, out_dim, n_total_expert, n_expert, + n_tokens, mid_token_stride, down_expert_bytes, down_row_bytes); + break; + } + return cuda_ok(cudaGetLastError(), "glm routed moe down launch"); +} + +/* ========================================================================= + * extern "C" backend entry points. + * ========================================================================= */ + +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 || !dst->ptr || !src->ptr) return 0; + if (count == 0) return 1; + if (count > UINT64_MAX / sizeof(float) || + count > UINT64_MAX / sizeof(uint16_t)) { + return 0; + } + const uint64_t src_bytes = count * sizeof(float); + const uint64_t dst_bytes = count * sizeof(uint16_t); + if (src_offset > src->bytes || src_bytes > src->bytes - src_offset || + dst_offset > dst->bytes || dst_bytes > dst->bytes - dst_offset) { + return 0; + } + f32_to_f16_kernel<<<(unsigned)((count + 255u) / 256u), 256>>>( + (__half *)((char *)dst->ptr + dst_offset), + (const float *)((const char *)src->ptr + src_offset), + count); + return cuda_ok(cudaGetLastError(), "tensor copy f32 to f16 launch"); +} + +/* On Metal this closes the current compute encoder so later blit work can be + * encoded. The CUDA stream is in-order with no encoder concept: no-op. */ +extern "C" int ds4_gpu_flush_encoder(void) { + return 1; +} + +extern "C" void ds4_gpu_release_q8_f16_cache(void) { + cuda_q8_f16_cache_release_all(); +} + +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; + } + glm_add3_kernel<<<(n + 255u) / 256u, 256>>>( + (float *)out->ptr, (const float *)a->ptr, + (const float *)b->ptr, (const float *)c->ptr, n); + return cuda_ok(cudaGetLastError(), "add3 launch"); +} + +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) return 0; + const uint64_t row_bytes = (uint64_t)n * sizeof(float); + if (norm_out->bytes < row_bytes || sum_out->bytes < row_bytes || + a->bytes < row_bytes || b->bytes < row_bytes || + weight_offset > model_size || row_bytes > model_size - weight_offset) { + return 0; + } + const float *w = (const float *)cuda_model_range_ptr(model_map, + weight_offset, row_bytes, "add_rms_weight"); + if (!w) return 0; + glm_add_rms_norm_weight_kernel<<<1, 256>>>( + (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 launch"); +} + +extern "C" int ds4_gpu_sort_i32_rows_asc_tensor( + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *src, + uint32_t row_width, + uint32_t n_rows) { + if (!dst || !src || row_width == 0 || n_rows == 0 || + (row_width & (row_width - 1u)) != 0) { + return 0; + } + const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t); + if (src->bytes < bytes || dst->bytes < bytes) { + fprintf(stderr, "ds4: CUDA row sort received undersized buffers\n"); + return 0; + } + const uint32_t shared_bytes = row_width * (uint32_t)sizeof(int32_t); + if (shared_bytes > 48u * 1024u) { + fprintf(stderr, "ds4: CUDA row sort scratch exceeds shared memory limit\n"); + return 0; + } + uint32_t threads = row_width < 256u ? row_width : 256u; + glm_sort_i32_rows_asc_kernel<<>>( + (int32_t *)dst->ptr, (const int32_t *)src->ptr, row_width, n_rows); + return cuda_ok(cudaGetLastError(), "sort i32 rows asc launch"); +} + +/* The Metal "decode mpp" and "model view" variants only change dispatch + * strategy / buffer wrapping; the math is the plain Q8_0 matmul. */ +extern "C" int ds4_gpu_matmul_q8_0_decode_mpp_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 ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, weight_offset, + in_dim, out_dim, x, n_tok); +} + +extern "C" int ds4_gpu_matmul_q8_0_decode_mpp_model_view_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 ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size, weight_offset, + in_dim, out_dim, x, n_tok); +} + +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_tensor(gate, up, mid, + model_map, model_size, + gate_offset, up_offset, + in_dim, out_dim, x, clamp); +} + +/* Mid-only fused shared expert: the Metal impl runs a fused kernel with no + * gate/up outputs. Here we keep two small persistent scratch buffers and + * delegate to the existing pair-matmul + swiglu path. */ +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 scratch_gate = {NULL, 0, 0}; + static ds4_gpu_tensor scratch_up = {NULL, 0, 0}; + if (!mid || out_dim == 0 || out_dim > UINT64_MAX / sizeof(float)) return 0; + const uint64_t bytes = out_dim * sizeof(float); + if (scratch_gate.bytes < bytes) { + if (scratch_gate.ptr) (void)cudaFree(scratch_gate.ptr); + if (scratch_up.ptr) (void)cudaFree(scratch_up.ptr); + scratch_gate.ptr = NULL; + scratch_up.ptr = NULL; + scratch_gate.bytes = 0; + scratch_up.bytes = 0; + void *g = NULL; + void *u = NULL; + if (!cuda_ok(cudaMalloc(&g, (size_t)bytes), "shared mid gate scratch") || + !cuda_ok(cudaMalloc(&u, (size_t)bytes), "shared mid up scratch")) { + if (g) (void)cudaFree(g); + return 0; + } + scratch_gate.ptr = g; + scratch_gate.bytes = bytes; + scratch_up.ptr = u; + scratch_up.bytes = bytes; + } + return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(&scratch_gate, &scratch_up, + mid, + model_map, model_size, + gate_offset, up_offset, + in_dim, out_dim, x, clamp); +} + +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 || !model_map || n_tok == 0 || + out_dim == 0 || n_tok > UINT32_MAX || + out_dim > UINT32_MAX || n_tok * out_dim > UINT32_MAX) { + return 0; + } + const uint64_t out_bytes = n_tok * out_dim * sizeof(float); + if (gate->bytes < out_bytes || up->bytes < out_bytes || mid->bytes < out_bytes) { + return 0; + } + 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, n_tok) && + ds4_gpu_swiglu_tensor(mid, gate, up, + (uint32_t)(n_tok * out_dim), clamp, 1.0f); +} + +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 glm_router_select_launch(selected, weights, probs, + model_map, model_size, bias_offset, + logits, n_expert, n_expert_used, + expert_weight_scale, 1); +} + +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) { + return glm_router_select_launch(selected, weights, probs, + model_map, model_size, bias_offset, + logits, n_expert, n_expert_used, + expert_weight_scale, n_tokens); +} + +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) { + if (n_expert == 0 || (uint64_t)n_expert * expert_mid_dim > UINT32_MAX) return 0; + return glm_routed_moe_launch(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, + 1u, + n_expert * expert_mid_dim, + force_resident); +} + +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) { + return glm_routed_moe_launch(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, + n_tokens, mid_token_stride, false); +} + +/* On Metal "direct scalar q4" only picks different pipelines for the same + * math (ds4_metal.m routes both through the shared batch impl); on CUDA the + * batch launcher is already the scalar path. */ +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) { + return glm_routed_moe_launch(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, + n_tokens, mid_token_stride, false); +} diff --git a/ds4_cuda_glm_stubs.inc b/ds4_cuda_glm_stubs.inc new file mode 100644 index 000000000..5e887b238 --- /dev/null +++ b/ds4_cuda_glm_stubs.inc @@ -0,0 +1,4 @@ +/* Auto-generated (tools/glm_stubgen.py). Stubs for backend functions the + * CUDA GLM port has not implemented yet. */ + + diff --git a/ds4_cuda_gqa.inc b/ds4_cuda_gqa.inc new file mode 100644 index 000000000..4f2cba010 --- /dev/null +++ b/ds4_cuda_gqa.inc @@ -0,0 +1,393 @@ +/* GQA full attention for Hy3-class models (hy_v3: 64 q heads / 8 kv heads, + * head_dim 128, per-head QK-RMSNorm with weights, rotate-half RoPE). + * + * ds4's existing attention paths are all MLA-family (shared latent KV, no + * per-kv-head cache); this file adds the grouped-query path: a standard + * per-head KV cache [n_kv_head][cap][head_dim] and a decode/prefill + * attention kernel where q head h attends kv head h / (n_head / n_kv_head). + * + * Everything here is validated by ds4_gpu_gqa_selftest(): a CPU reference + * implementation runs the identical pipeline on random data and the kernels + * must match it numerically. Run with `ds4 --gqa-selftest` (no model file + * needed), including a multi-chunk phase for the online-softmax rescale. + * + * RoPE convention: HF Llama rotate-half. Pair (i, i + d/2), angle + * pos * theta^(-2i/d). Hy3 rope_type is "default" (no yarn/ntk scaling). */ + +/* ---- kernels ---- */ + +/* Per-head RMSNorm with a [head_dim] weight vector, in place. + * x: [rows][head_dim] where rows = n_tok * n_head (or n_kv_head). */ +__global__ static void gqa_head_rms_norm_weight_kernel( + float *x, + const float *w, + uint32_t rows, + uint32_t head_dim, + float eps) { + uint32_t row = blockIdx.x; + if (row >= rows) return; + float *v = x + (uint64_t)row * head_dim; + __shared__ float red[256]; + float acc = 0.0f; + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + acc += v[d] * v[d]; + } + red[threadIdx.x] = acc; + __syncthreads(); + for (uint32_t s = blockDim.x / 2u; s != 0; s >>= 1u) { + if (threadIdx.x < s) red[threadIdx.x] += red[threadIdx.x + s]; + __syncthreads(); + } + float inv = rsqrtf(red[0] / (float)head_dim + eps); + for (uint32_t d = threadIdx.x; d < head_dim; d += blockDim.x) { + v[d] = v[d] * inv * w[d]; + } +} + +/* Rotate-half RoPE over the full head, in place. + * x: [n_tok][n_head][head_dim], token t is at position pos0 + t. */ +__global__ static void gqa_rope_neox_kernel( + float *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t pos0, + float theta) { + uint32_t half = head_dim / 2u; + uint64_t total = (uint64_t)n_tok * n_head * half; + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + uint32_t d = (uint32_t)(i % half); + uint32_t rem = (uint32_t)(i / half); + uint32_t h = rem % n_head; + uint32_t t = rem / n_head; + float *v = x + ((uint64_t)t * n_head + h) * head_dim; + float freq = powf(theta, -2.0f * (float)d / (float)head_dim); + float ang = (float)(pos0 + t) * freq; + float c = cosf(ang), s = sinf(ang); + float a = v[d], b = v[d + half]; + v[d] = a * c - b * s; + v[d + half] = a * s + b * c; +} + +/* Append n_tok tokens of K and V into the per-head cache. + * kv_in: [n_tok][n_kv_head][head_dim]; cache: [n_kv_head][cap][head_dim]. */ +__global__ static void gqa_kv_cache_append_kernel( + float *cache, + const float *kv_in, + uint32_t n_tok, + uint32_t n_kv_head, + uint32_t head_dim, + uint32_t cap, + uint32_t pos0) { + uint64_t total = (uint64_t)n_tok * n_kv_head * head_dim; + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + uint32_t d = (uint32_t)(i % head_dim); + uint32_t rem = (uint32_t)(i / head_dim); + uint32_t h = rem % n_kv_head; + uint32_t t = rem / n_kv_head; + cache[((uint64_t)h * cap + pos0 + t) * head_dim + d] = + kv_in[((uint64_t)t * n_kv_head + h) * head_dim + d]; +} + +/* Causal GQA attention over the cache, single-pass online softmax. + * q: [n_tok][n_head][head_dim] (post norm+rope); token t sees cache rows + * [0, pos0 + t]. out: [n_tok][n_head][head_dim]. One block per (t, h). */ +__global__ static void gqa_attention_kernel( + float *out, + const float *q, + const float *k_cache, + const float *v_cache, + uint32_t n_tok, + uint32_t n_head, + uint32_t n_kv_head, + uint32_t head_dim, + uint32_t cap, + uint32_t pos0) { + uint32_t t = blockIdx.x; + uint32_t h = blockIdx.y; + if (t >= n_tok || h >= n_head) return; + uint32_t kvh = h / (n_head / n_kv_head); + uint32_t n_rows = pos0 + t + 1u; + const float *qh = q + ((uint64_t)t * n_head + h) * head_dim; + const float *kh = k_cache + (uint64_t)kvh * cap * head_dim; + const float *vh = v_cache + (uint64_t)kvh * cap * head_dim; + float scale = rsqrtf((float)head_dim); + /* Online softmax, single pass over K and V. 128 threads: thread d owns + * output dim d; rows are processed in chunks of 128 with each thread + * computing one row's score, then everyone accumulates V against the + * shared per-chunk weights. Scores are computed exactly once (the old + * two-pass version recomputed them per output dim: ~130x the FLOPs and, + * past L1, ~head_dim re-reads of K per row at long context). */ + const uint32_t tid = threadIdx.x; + __shared__ float sq[128]; + __shared__ float s_w[128]; + __shared__ float red[128]; + if (tid < head_dim) sq[tid] = qh[tid]; + __syncthreads(); + + float acc = 0.0f; /* V accumulator for dim = tid */ + float den = 0.0f; /* running softmax denominator (same in all threads) */ + float m = -INFINITY;/* running max */ + for (uint32_t base = 0; base < n_rows; base += 128u) { + uint32_t chunk = n_rows - base < 128u ? n_rows - base : 128u; + /* one score per thread */ + float s = -INFINITY; + if (tid < chunk) { + const float *krow = kh + (uint64_t)(base + tid) * head_dim; + float dot = 0.0f; + for (uint32_t d = 0; d < head_dim; d++) dot += sq[d] * krow[d]; + s = dot * scale; + } + red[tid] = s; + __syncthreads(); + for (uint32_t st = 64u; st != 0; st >>= 1u) { + if (tid < st) red[tid] = fmaxf(red[tid], red[tid + st]); + __syncthreads(); + } + float m_new = fmaxf(m, red[0]); + float corr = expf(m - m_new); /* expf(-inf) == 0, first chunk safe */ + acc *= corr; + den *= corr; + if (tid < chunk) s_w[tid] = expf(s - m_new); + __syncthreads(); + for (uint32_t r = 0; r < chunk; r++) { + float w = s_w[r]; + den += w; + if (tid < head_dim) { + acc += w * vh[(uint64_t)(base + r) * head_dim + tid]; + } + } + m = m_new; + __syncthreads(); + } + if (tid < head_dim) { + out[((uint64_t)t * n_head + h) * head_dim + tid] = acc / den; + } +} + +/* ---- host wrappers ---- */ + +extern "C" int ds4_gpu_gqa_head_rms_norm_weight( + ds4_gpu_tensor *x, const ds4_gpu_tensor *w, + uint32_t rows, uint32_t head_dim, float eps) { + if (!x || !w || rows == 0 || head_dim == 0 || head_dim > 128u) return 0; + gqa_head_rms_norm_weight_kernel<<>>( + (float *)x->ptr, (const float *)w->ptr, rows, head_dim, eps); + return cuda_ok(cudaGetLastError(), "gqa head rms norm launch"); +} + +extern "C" int ds4_gpu_gqa_rope( + ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, + uint32_t head_dim, uint32_t pos0, float theta) { + if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || head_dim % 2u) return 0; + uint64_t total = (uint64_t)n_tok * n_head * (head_dim / 2u); + gqa_rope_neox_kernel<<<(uint32_t)((total + 255u) / 256u), 256>>>( + (float *)x->ptr, n_tok, n_head, head_dim, pos0, theta); + return cuda_ok(cudaGetLastError(), "gqa rope launch"); +} + +extern "C" int ds4_gpu_gqa_kv_cache_append( + ds4_gpu_tensor *cache, const ds4_gpu_tensor *kv_in, + uint32_t n_tok, uint32_t n_kv_head, uint32_t head_dim, + uint32_t cap, uint32_t pos0) { + if (!cache || !kv_in || n_tok == 0 || pos0 + n_tok > cap) return 0; + uint64_t total = (uint64_t)n_tok * n_kv_head * head_dim; + gqa_kv_cache_append_kernel<<<(uint32_t)((total + 255u) / 256u), 256>>>( + (float *)cache->ptr, (const float *)kv_in->ptr, + n_tok, n_kv_head, head_dim, cap, pos0); + return cuda_ok(cudaGetLastError(), "gqa kv append launch"); +} + +extern "C" int ds4_gpu_gqa_attention( + ds4_gpu_tensor *out, const ds4_gpu_tensor *q, + const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache, + uint32_t n_tok, uint32_t n_head, uint32_t n_kv_head, + uint32_t head_dim, uint32_t cap, uint32_t pos0) { + if (!out || !q || !k_cache || !v_cache || + n_head == 0 || n_kv_head == 0 || n_head % n_kv_head != 0 || + head_dim == 0 || head_dim > 128u || pos0 + n_tok > cap) { + return 0; + } + dim3 grid(n_tok, n_head, 1); + gqa_attention_kernel<<>>( + (float *)out->ptr, (const float *)q->ptr, + (const float *)k_cache->ptr, (const float *)v_cache->ptr, + n_tok, n_head, n_kv_head, head_dim, cap, pos0); + return cuda_ok(cudaGetLastError(), "gqa attention launch"); +} + +/* ---- CPU reference + self-test ---- */ + +static void gqa_ref_head_rms_norm(float *x, const float *w, + uint32_t rows, uint32_t hd, float eps) { + for (uint32_t r = 0; r < rows; r++) { + float *v = x + (uint64_t)r * hd; + double ss = 0.0; + for (uint32_t d = 0; d < hd; d++) ss += (double)v[d] * v[d]; + float inv = 1.0f / sqrtf((float)(ss / hd) + eps); + for (uint32_t d = 0; d < hd; d++) v[d] = v[d] * inv * w[d]; + } +} + +static void gqa_ref_rope(float *x, uint32_t n_tok, uint32_t n_head, + uint32_t hd, uint32_t pos0, float theta) { + uint32_t half = hd / 2u; + for (uint32_t t = 0; t < n_tok; t++) + for (uint32_t h = 0; h < n_head; h++) { + float *v = x + ((uint64_t)t * n_head + h) * hd; + for (uint32_t d = 0; d < half; d++) { + float freq = powf(theta, -2.0f * (float)d / (float)hd); + float ang = (float)(pos0 + t) * freq; + float c = cosf(ang), s = sinf(ang); + float a = v[d], b = v[d + half]; + v[d] = a * c - b * s; + v[d + half] = a * s + b * c; + } + } +} + +static void gqa_ref_attention(float *out, const float *q, + const float *kc, const float *vc, + uint32_t n_tok, uint32_t n_head, + uint32_t n_kv_head, uint32_t hd, + uint32_t cap, uint32_t pos0) { + float scale = 1.0f / sqrtf((float)hd); + uint32_t group = n_head / n_kv_head; + float *scores = (float *)malloc((size_t)(pos0 + n_tok) * sizeof(float)); + for (uint32_t t = 0; t < n_tok; t++) + for (uint32_t h = 0; h < n_head; h++) { + uint32_t kvh = h / group; + uint32_t n_rows = pos0 + t + 1u; + const float *qh = q + ((uint64_t)t * n_head + h) * hd; + const float *kh = kc + (uint64_t)kvh * cap * hd; + const float *vh = vc + (uint64_t)kvh * cap * hd; + float m = -INFINITY; + for (uint32_t r = 0; r < n_rows; r++) { + float dot = 0.0f; + for (uint32_t d = 0; d < hd; d++) dot += qh[d] * kh[(uint64_t)r * hd + d]; + scores[r] = dot * scale; + if (scores[r] > m) m = scores[r]; + } + float den = 0.0f; + for (uint32_t r = 0; r < n_rows; r++) den += expf(scores[r] - m); + float *oh = out + ((uint64_t)t * n_head + h) * hd; + for (uint32_t d = 0; d < hd; d++) { + float acc = 0.0f; + for (uint32_t r = 0; r < n_rows; r++) { + acc += expf(scores[r] - m) / den * vh[(uint64_t)r * hd + d]; + } + oh[d] = acc; + } + } + free(scores); +} + +static float gqa_max_abs_diff(const float *a, const float *b, uint64_t n) { + float m = 0.0f; + for (uint64_t i = 0; i < n; i++) { + float d = fabsf(a[i] - b[i]); + if (d > m) m = d; + } + return m; +} + +static uint32_t gqa_test_rng_state = 0x1234567u; +static float gqa_test_randf(void) { + gqa_test_rng_state = gqa_test_rng_state * 1664525u + 1013904223u; + return ((float)(gqa_test_rng_state >> 8) / (float)(1u << 24)) * 2.0f - 1.0f; +} + +/* Full-pipeline self-test on random data: qk-norm -> rope -> cache append + * -> causal GQA attention, CPU reference vs kernels. Hy3 dims: 64/8 heads, + * head_dim 128. Covers a 13-token prefill at pos0=0 and a 1-token decode at + * pos0=13 against the populated cache. Returns 1 on pass. */ +extern "C" int ds4_gpu_gqa_selftest(void) { + const uint32_t n_head = 64, n_kv_head = 8, hd = 128, cap = 256; + const float eps = 1e-5f, theta = 11158840.0f; + /* {n_tok, pos0}: short prefill, decode, then a 200-token prefill and a + * decode at row 215 to exercise the multi-chunk online-softmax path + * (rows > 128 crosses a chunk boundary and rescales acc/den). */ + const uint32_t phases[4][2] = { {13u, 0u}, {1u, 13u}, {200u, 14u}, {1u, 214u} }; + int all_ok = 1; + + float *qn_w = (float *)malloc(hd * sizeof(float)); + float *kn_w = (float *)malloc(hd * sizeof(float)); + for (uint32_t d = 0; d < hd; d++) { qn_w[d] = 0.5f + 0.5f * fabsf(gqa_test_randf()); kn_w[d] = 0.5f + 0.5f * fabsf(gqa_test_randf()); } + + uint64_t kcache_elems = (uint64_t)n_kv_head * cap * hd; + float *ref_kc = (float *)calloc(kcache_elems, sizeof(float)); + float *ref_vc = (float *)calloc(kcache_elems, sizeof(float)); + + ds4_gpu_tensor *g_qnw = ds4_gpu_tensor_alloc(hd * sizeof(float)); + ds4_gpu_tensor *g_knw = ds4_gpu_tensor_alloc(hd * sizeof(float)); + ds4_gpu_tensor *g_kc = ds4_gpu_tensor_alloc(kcache_elems * sizeof(float)); + ds4_gpu_tensor *g_vc = ds4_gpu_tensor_alloc(kcache_elems * sizeof(float)); + if (!g_qnw || !g_knw || !g_kc || !g_vc) { fprintf(stderr, "gqa-selftest: alloc failed\n"); return 0; } + ds4_gpu_tensor_write(g_qnw, 0, qn_w, hd * sizeof(float)); + ds4_gpu_tensor_write(g_knw, 0, kn_w, hd * sizeof(float)); + { float *z = (float *)calloc(kcache_elems, sizeof(float)); + ds4_gpu_tensor_write(g_kc, 0, z, kcache_elems * sizeof(float)); + ds4_gpu_tensor_write(g_vc, 0, z, kcache_elems * sizeof(float)); + free(z); } + + for (int phase = 0; phase < 4; phase++) { + uint32_t n_tok = phases[phase][0], pos0 = phases[phase][1]; + uint64_t q_elems = (uint64_t)n_tok * n_head * hd; + uint64_t kv_elems = (uint64_t)n_tok * n_kv_head * hd; + float *q = (float *)malloc(q_elems * sizeof(float)); + float *k = (float *)malloc(kv_elems * sizeof(float)); + float *v = (float *)malloc(kv_elems * sizeof(float)); + for (uint64_t i = 0; i < q_elems; i++) q[i] = gqa_test_randf(); + for (uint64_t i = 0; i < kv_elems; i++) { k[i] = gqa_test_randf(); v[i] = gqa_test_randf(); } + + ds4_gpu_tensor *g_q = ds4_gpu_tensor_alloc(q_elems * sizeof(float)); + ds4_gpu_tensor *g_k = ds4_gpu_tensor_alloc(kv_elems * sizeof(float)); + ds4_gpu_tensor *g_v = ds4_gpu_tensor_alloc(kv_elems * sizeof(float)); + ds4_gpu_tensor *g_out = ds4_gpu_tensor_alloc(q_elems * sizeof(float)); + ds4_gpu_tensor_write(g_q, 0, q, q_elems * sizeof(float)); + ds4_gpu_tensor_write(g_k, 0, k, kv_elems * sizeof(float)); + ds4_gpu_tensor_write(g_v, 0, v, kv_elems * sizeof(float)); + + /* reference pipeline */ + gqa_ref_head_rms_norm(q, qn_w, n_tok * n_head, hd, eps); + gqa_ref_head_rms_norm(k, kn_w, n_tok * n_kv_head, hd, eps); + gqa_ref_rope(q, n_tok, n_head, hd, pos0, theta); + gqa_ref_rope(k, n_tok, n_kv_head, hd, pos0, theta); + for (uint32_t t = 0; t < n_tok; t++) + for (uint32_t h = 0; h < n_kv_head; h++) + for (uint32_t d = 0; d < hd; d++) { + ref_kc[((uint64_t)h * cap + pos0 + t) * hd + d] = k[((uint64_t)t * n_kv_head + h) * hd + d]; + ref_vc[((uint64_t)h * cap + pos0 + t) * hd + d] = v[((uint64_t)t * n_kv_head + h) * hd + d]; + } + float *ref_out = (float *)malloc(q_elems * sizeof(float)); + gqa_ref_attention(ref_out, q, ref_kc, ref_vc, n_tok, n_head, n_kv_head, hd, cap, pos0); + + /* kernel pipeline */ + int ok = ds4_gpu_gqa_head_rms_norm_weight(g_q, g_qnw, n_tok * n_head, hd, eps) && + ds4_gpu_gqa_head_rms_norm_weight(g_k, g_knw, n_tok * n_kv_head, hd, eps) && + ds4_gpu_gqa_rope(g_q, n_tok, n_head, hd, pos0, theta) && + ds4_gpu_gqa_rope(g_k, n_tok, n_kv_head, hd, pos0, theta) && + ds4_gpu_gqa_kv_cache_append(g_kc, g_k, n_tok, n_kv_head, hd, cap, pos0) && + ds4_gpu_gqa_kv_cache_append(g_vc, g_v, n_tok, n_kv_head, hd, cap, pos0) && + ds4_gpu_gqa_attention(g_out, g_q, g_kc, g_vc, n_tok, n_head, n_kv_head, hd, cap, pos0) && + cuda_ok(cudaDeviceSynchronize(), "gqa selftest sync"); + float *gpu_out = (float *)malloc(q_elems * sizeof(float)); + if (ok) ok = ds4_gpu_tensor_read(g_out, 0, gpu_out, q_elems * sizeof(float)) != 0; + float diff = ok ? gqa_max_abs_diff(ref_out, gpu_out, q_elems) : INFINITY; + int pass = ok && diff < 2e-4f; + fprintf(stderr, "gqa-selftest phase %d (n_tok=%u pos0=%u): %s (max diff %.2e)\n", + phase, n_tok, pos0, pass ? "PASS" : "FAIL", diff); + if (!pass) all_ok = 0; + + free(q); free(k); free(v); free(ref_out); free(gpu_out); + ds4_gpu_tensor_free(g_q); ds4_gpu_tensor_free(g_k); + ds4_gpu_tensor_free(g_v); ds4_gpu_tensor_free(g_out); + } + + free(qn_w); free(kn_w); free(ref_kc); free(ref_vc); + ds4_gpu_tensor_free(g_qnw); ds4_gpu_tensor_free(g_knw); + ds4_gpu_tensor_free(g_kc); ds4_gpu_tensor_free(g_vc); + return all_ok; +} diff --git a/ds4_gpu.h b/ds4_gpu.h index fbc06fb14..bbcc83ba3 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -454,6 +454,40 @@ int ds4_gpu_rms_norm_weight_rows_tensor( uint32_t rows, float eps); +/* GQA attention (Hy3-class models): ds4_cuda_gqa.inc */ +int ds4_gpu_gqa_head_rms_norm_weight( + ds4_gpu_tensor *x, + const ds4_gpu_tensor *w, + uint32_t rows, + uint32_t head_dim, + float eps); +int ds4_gpu_gqa_rope( + ds4_gpu_tensor *x, + uint32_t n_tok, + uint32_t n_head, + uint32_t head_dim, + uint32_t pos0, + float theta); +int ds4_gpu_gqa_kv_cache_append( + ds4_gpu_tensor *cache, + const ds4_gpu_tensor *kv_in, + uint32_t n_tok, + uint32_t n_kv_head, + uint32_t head_dim, + uint32_t cap, + uint32_t pos0); +int ds4_gpu_gqa_attention( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *q, + const ds4_gpu_tensor *k_cache, + const ds4_gpu_tensor *v_cache, + uint32_t n_tok, + uint32_t n_head, + uint32_t n_kv_head, + uint32_t head_dim, + uint32_t cap, + uint32_t pos0); + int ds4_gpu_add_rms_norm_weight_tensor( ds4_gpu_tensor *norm_out, ds4_gpu_tensor *sum_out,