diff --git a/ds4.c b/ds4.c index 582d5afb40..1ab223102b 100644 --- a/ds4.c +++ b/ds4.c @@ -39437,6 +39437,19 @@ static bool ds41_bf16(ds4_gpu_tensor *x, uint32_t width) { static bool ds41_matmul(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, bool round) { +#if defined(__APPLE__) + /* Round on the store instead of a rounding dispatch over the result. */ + static int fused = -1; + if (fused < 0) { + fused = getenv("DS4_METAL_DISABLE_V41_MATVEC_BF16") == NULL && + getenv("DS4_METAL_V41_SKIP_BF16") == NULL; + } + if (round && fused) { + const int rc = ds4_gpu_dsv41_matvec_bf16(out, m->map, m->size, weight->abs_offset, + weight->type, (uint32_t)weight->dim[0], (uint32_t)weight->dim[1], in); + if (rc) return rc > 0; + } +#endif return metal_graph_matmul_plain_tensor(out, m, weight, weight->dim[0], weight->dim[1], in, 1) && (!round || ds41_bf16(out, (uint32_t)weight->dim[1])); } @@ -39444,6 +39457,7 @@ static bool ds41_matmul(ds4_gpu_tensor *out, const ds4_model *m, static bool ds41_matmul_batch(ds4_gpu_tensor *out, const ds4_model *m, const ds4_tensor *weight, const ds4_gpu_tensor *in, uint32_t count, bool round) { + if (count == 1) return ds41_matmul(out, m, weight, in, round); const uint32_t width = (uint32_t)weight->dim[0], outputs = (uint32_t)weight->dim[1]; bool ok; /* Small decode batches retain scalar reductions before BF16 and sparse @@ -39611,15 +39625,16 @@ static bool ds41_attention_low(ds41_gpu_graph *g, const ds4_model *m, 4096, 1024, groups, g->heads) && ds41_bf16(g->low, groups * DS4_N_LORA_O); } +/* round: the single-box block leaves rounded to BF16 (fused attention glue). */ static bool ds41_attention_output(ds41_gpu_graph *g, const ds4_model *m, - const ds4_layer_weights *l) { + const ds4_layer_weights *l, bool round) { const uint32_t groups = DS4_N_OUT_GROUP / g->tp_world; if (!ds41_attention_low(g, m, l)) return false; return g->tp_world == 2 ? metal_graph_matmul_dense_quant_kslice(g->block, m, l->attn_output_b, 8192, (uint64_t)g->tp_rank * groups * 1024u, (uint64_t)groups * 1024u, DS4_N_EMBD, g->low, 0) : - ds41_matmul(g->block, m, l->attn_output_b, g->low, false); + ds41_matmul(g->block, m, l->attn_output_b, g->low, round); } static bool ds41_attention_publish(ds41_gpu_graph *g, const ds4_model *m, @@ -39702,9 +39717,83 @@ static bool ds41_attention_select(ds41_gpu_graph *g, const ds4_model *m, ds41_attention_select_published(g, m, l, il); } +#if defined(__APPLE__) +/* DeepSeek's production decode runs each half-layer's hyper-connection work + * in one kernel ("Mega-mHC", tech report 3.2) and the router as "Mega-Gate". + * ds4's V4.1 graph mirrored the reference op by op: 19 HC glue dispatches + * per layer (norm, mix matvec, split, collapse, BF16, weighted norm, BF16, + * expand, BF16, ...) and 22 of MoE glue (router matvec + ten select + * dispatches, shared gate/up/down + their BF16 passes, SwiGLU, the routed + + * shared sum). The fused paths keep every reduction tree and rounding point + * of those sequences (tests/test_deepseek41_metal --hc-fuse / --moe-fuse) + * and issue 6 + 5. Single box only; the TP graph keeps the unfused + * sequences. */ +static bool ds41_hc_fused(const ds41_gpu_graph *g) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_HC_FUSE") == NULL && + ds4_gpu_hc_rms_norm_mix_f16_available() != 0; + } + return enabled && g->tp_world == 1; +} + +/* Router matvec + select as one (two on pre-M5) dispatch. */ +static bool ds41_router_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL && + getenv("DS4_METAL_DISABLE_V41_ROUTER_FUSE") == NULL; + } + return enabled && ds41_hc_fused(g) && l->ffn_gate_inp->type == DS4_TENSOR_F32; +} + +/* Shared expert gate/up/SwiGLU as one dispatch, its down projection rounded + * on the store, and the routed + shared sum folded into the FFN tail's + * expand (ds41_moe_finish / ds41_graph_after_moe agree). The down + * projection runs BEFORE the routed experts: a fused tail that read + * shared_mid after them saw clobbered values on the M3 Ultra. */ +static bool ds41_moe_fused(const ds41_gpu_graph *g, const ds4_layer_weights *l) { + static int enabled = -1; + if (enabled < 0) { + enabled = getenv("DS4_METAL_DISABLE_V41_MOE_FUSE") == NULL && + getenv("DS4_METAL_DISABLE_V41_SHARED_FUSE") == NULL; + } + return enabled && ds41_hc_fused(g) && + l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + l->ffn_down_shexp->type == DS4_TENSOR_Q8_0; +} + +/* Attention glue of the single-box decode graph as one dispatch each: the + * q/kv norms with the KV RoPE, FP8 quantization and raw-window store; the + * heads' rounding with their inverse RoPE; the output projection rounded on + * its store. Byte-identical to the standalone sequences + * (tests/test_deepseek41_metal --attn-fuse). */ +static bool ds41_attn_fused(const ds41_gpu_graph *g) { + static int enabled = -1; + if (enabled < 0) enabled = getenv("DS4_METAL_DISABLE_V41_ATTN_FUSE") == NULL; + return enabled && ds41_hc_fused(g); +} +#endif + +/* With the fused attention glue the KV row leaves here roped, quantized and + * stored in the raw window; ds41_attention skips those steps. */ static bool ds41_attention_project(ds41_gpu_graph *g, const ds4_model *m, - const ds4_layer_weights *l) { + const ds4_layer_weights *l, uint32_t il) { const uint32_t q_dim = DS4_N_HEAD / g->tp_world * DS4_N_HEAD_DIM; +#if defined(__APPLE__) + if (ds41_attn_fused(g)) + return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && + ds41_matmul(g->kv, m, l->attn_kv, g->norm, true) && + ds4_gpu_dsv41_qkv_norm_kv_tail(g->qr, g->kv, g->window[il], + (uint64_t)(g->pos % 128u) * 512u * 4u, m->map, m->size, + l->attn_q_a_norm->abs_offset, l->attn_kv_a_norm->abs_offset, + DS4_N_LORA_Q, DS4_N_HEAD_DIM, DS4_RMS_EPS, g->pos, + ds4_layer_compress_ratio(il) != 0) && + ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, 0, q_dim); +#else + (void)il; +#endif return ds41_matmul(g->qr, m, l->attn_q_a, g->norm, true) && ds41_norm(g->qr, g->qr, m, l->attn_q_a_norm) && ds41_matmul_rows(g->q, m, l->attn_q_b, g->qr, @@ -39720,12 +39809,20 @@ static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, const uint32_t n_comp = ratio ? (pos + 1u) / ratio : 0u; const uint32_t heads = DS4_N_HEAD / g->tp_world; const uint32_t head0 = g->tp_rank * heads; - if (!projected && !ds41_attention_project(g, m, l)) return false; + if (!projected && !ds41_attention_project(g, m, l, il)) return false; + /* Projected rows (the prefill's per-row path) had their projections done + * elsewhere: their KV tail runs here, unfused. */ +#if defined(__APPLE__) + const bool fused = !projected && ds41_attn_fused(g); +#else + const bool fused = false; +#endif if (!ds41_rope(g->q, heads, DS4_N_HEAD_DIM, il, pos, false) || - !ds41_rope(g->kv, 1, DS4_N_HEAD_DIM, il, pos, false) || - !ds4_gpu_dsv41_quantize(g->kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0) || - !ds4_gpu_tensor_copy(g->window[il], (uint64_t)(pos % 128u) * 512u * 4u, - g->kv, 0, 512u * 4u) || + (!fused && + (!ds41_rope(g->kv, 1, DS4_N_HEAD_DIM, il, pos, false) || + !ds4_gpu_dsv41_quantize(g->kv, DS4_N_HEAD_DIM, 1, DS4_V41_FP8_E8M0) || + !ds4_gpu_tensor_copy(g->window[il], (uint64_t)(pos % 128u) * 512u * 4u, + g->kv, 0, 512u * 4u))) || !ds41_attention_select(g, m, l, il)) return false; const uint32_t attended = n_comp < DS4_N_INDEXER_TOP_K ? n_comp : DS4_N_INDEXER_TOP_K; if (n_comp && !ds4_gpu_dsv41_gather_kv(g->selected_kv, g->compressed[owner], @@ -39736,12 +39833,15 @@ static bool ds41_attention(ds41_gpu_graph *g, const ds4_model *m, g->q, g->window[il], n_raw, 128, (pos + 1u - n_raw) % 128u, g->selected_kv, 0, attended, NULL, 0, heads, DS4_N_HEAD_DIM) || - !ds41_bf16(g->heads, heads * DS4_N_HEAD_DIM) || - !ds41_rope(g->heads, heads, DS4_N_HEAD_DIM, il, pos, true)) return false; + (fused ? + !ds4_gpu_dsv41_bf16_rope(g->heads, DS4_N_HEAD_DIM, heads, 1, pos, + ds4_layer_compress_ratio(il) != 0, true) : + (!ds41_bf16(g->heads, heads * DS4_N_HEAD_DIM) || + !ds41_rope(g->heads, heads, DS4_N_HEAD_DIM, il, pos, true)))) return false; if (projected) return true; - return ds41_attention_output(g, m, l) && + return ds41_attention_output(g, m, l, fused) && ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && - ds41_bf16(g->block, DS4_N_EMBD); + (fused || ds41_bf16(g->block, DS4_N_EMBD)); } static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, @@ -39754,6 +39854,15 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; const ds4_tensor *bias = ds41_image_at(g, g->pos) ? l->ffn_exp_probs_vl : l->ffn_exp_probs_b; if (!bias) return false; +#if defined(__APPLE__) + const bool fused = ds41_moe_fused(g, l); + if (ds41_router_fused(g, l)) { + if (!ds4_gpu_dsv41_router_select(g->selected, g->route_weights, g->route_probs, + g->route_logits, g->norm, m->map, m->size, l->ffn_gate_inp->abs_offset, + bias->abs_offset, true, DS4_N_EMBD, DS4_N_EXPERT, DS4_N_EXPERT_USED, + DS4_EXPERT_WEIGHT_SCALE)) return false; + } else +#endif if (!ds41_matmul(g->route_logits, m, l->ffn_gate_inp, g->norm, false) || !ds4_gpu_router_select_tensor(g->selected, g->route_weights, g->route_probs, m->map, m->size, bias->abs_offset, 0, 0, token, @@ -39761,6 +39870,16 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, g->route_logits)) return false; const bool shared_here = !shared_owner || g->tp_rank == (il & 1u); bool shared_queued = false; +#if defined(__APPLE__) + /* Fused: gate/up/SwiGLU here; the down projection, the routed + shared + * sum and the HC expand run as one dispatch in ds41_graph_after_moe. */ + if (fused && shared_here && + (!ds4_gpu_dsv41_shared_gate_up_swiglu(g->shared_mid, g->norm, m->map, m->size, + l->ffn_gate_shexp->abs_offset, l->ffn_up_shexp->abs_offset, + DS4_N_EMBD, DS4_N_FF_EXP, DS4_SWIGLU_CLAMP_EXP) || + !ds41_matmul(g->shared, m, l->ffn_down_shexp, g->shared_mid, true))) return false; + if (fused) shared_queued = true; +#endif #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) if (shared_owner && shared_here && l->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && @@ -39812,32 +39931,60 @@ static bool ds41_moe_partial(ds41_gpu_graph *g, const ds4_model *m, return true; } -static bool ds41_moe_finish(ds41_gpu_graph *g, uint32_t il) { +static bool ds41_moe_finish(ds41_gpu_graph *g, const ds4_layer_weights *l, uint32_t il) { const bool shared_owner = g->tp_world == 2 && !getenv("DS4_METAL_DISABLE_V41_TP_SHARED_OWNER"); ds4_gpu_tensor *routed = shared_owner ? g->block : g->routed; if (!ds41_sum_partial(g, routed, il, DS4_TP_GATE_FFN)) return false; +#if defined(__APPLE__) + if (ds41_moe_fused(g, l)) return true; +#endif return (shared_owner || ds4_gpu_add_tensor(g->block, routed, g->shared, DS4_N_EMBD)) && ds41_bf16(g->block, DS4_N_EMBD); } static bool ds41_moe(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, uint32_t token) { - return ds41_moe_partial(g, m, l, il, token) && ds41_moe_finish(g, il); + return ds41_moe_partial(g, m, l, il, token) && ds41_moe_finish(g, l, il); } static bool ds41_graph_logits(ds41_gpu_graph *g, const ds4_model *m, const ds4_weights *w, float *logits) { if (!g->valid || !logits || !ds4_gpu_begin_commands()) return false; - bool ok = ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && - ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, w->output_norm) && - ds41_output_projection(g, g->tp_logits_half ? g->tp_logits_half : g->logits, + bool ok = +#if defined(__APPLE__) + ds41_hc_fused(g) ? + ds4_gpu_dsv41_hc_collapse_norm(NULL, g->x, g->norm, NULL, g->pre, g->residual, m->map, m->size, + 0, 0, w->output_norm->abs_offset, DS4_N_EMBD, DS4_N_HC, 0, 0.0f, DS4_RMS_EPS) != 0 : +#endif + ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && + ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, w->output_norm); + ok = ok && ds41_output_projection(g, g->tp_logits_half ? g->tp_logits_half : g->logits, m, w, g->norm, 1); if (!ds4_gpu_end_commands()) ok = false; return ok && ds4_gpu_tensor_read(g->logits, 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)); } +#if defined(__APPLE__) +/* Plain RMSNorm over the flattened HC row and the F16 mix matvec, one dispatch. */ +static bool ds41_hc_mix_fused(ds41_gpu_graph *g, const ds4_model *m, + const ds4_tensor *fn, const ds4_gpu_tensor *residual) { + return ds4_gpu_hc_rms_norm_mix_f16_tensor(g->mix, residual, m->map, m->size, fn->abs_offset, + DS4_N_HC * DS4_N_EMBD, (uint32_t)fn->dim[1], DS4_RMS_EPS) != 0; +} + +static bool ds41_hc_collapse_norm_fused(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l, bool ffn) { + return ds4_gpu_dsv41_hc_collapse_norm(ffn ? g->ffn_split : g->attn_split, g->x, g->norm, g->mix, + ffn ? g->attn_split : g->pre, ffn ? g->after_attn : g->residual, m->map, m->size, + (ffn ? l->hc_ffn_scale : l->hc_attn_scale)->abs_offset, + (ffn ? l->hc_ffn_base : l->hc_attn_base)->abs_offset, + (ffn ? l->ffn_norm : l->attn_norm)->abs_offset, + DS4_N_EMBD, DS4_N_HC, DS4_N_HC_SINKHORN_ITER, DS4_HC_EPS, DS4_RMS_EPS) != 0; +} +#endif + static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il) { if (ds41_engram_layer(il) && !ds41_image_at(g, g->pos)) { @@ -39847,6 +39994,11 @@ static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, g->engram_q_norm[i], g->engram_k_norm[i], NULL, DS4_N_EMBD, 1, DS4_RMS_EPS)) return false; } +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds41_hc_mix_fused(g, m, l->hc_attn_fn, g->residual) && + ds41_hc_collapse_norm_fused(g, m, l, false); +#endif return ds41_hc_mix(g, m, l, false) && ds4_gpu_hc_weighted_sum_tensor(g->x, g->residual, g->pre, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->x, DS4_N_EMBD) && ds41_norm(g->norm, g->x, m, l->attn_norm); @@ -39854,6 +40006,12 @@ static bool ds41_graph_before_attention(ds41_gpu_graph *g, const ds4_model *m, static bool ds41_graph_after_attention(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l) { +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds4_gpu_dsv41_hc_expand4(g->after_attn, g->block, g->residual, g->attn_split, NULL, NULL, NULL, DS4_N_EMBD) && + ds41_hc_mix_fused(g, m, l->hc_ffn_fn, g->after_attn) && + ds41_hc_collapse_norm_fused(g, m, l, true); +#endif return ds4_gpu_hc_expand_split_tensor(g->after_attn, g->block, g->residual, g->attn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->after_attn, DS4_N_EMBD * DS4_N_HC) && ds41_hc_mix(g, m, l, true) && @@ -40137,16 +40295,37 @@ static bool ds41_attention_batch(ds41_gpu_graph *g, const ds4_model *m, (n_raw - kept + part) * row_bytes, (kept - part) * row_bytes)); } -static bool ds41_graph_after_moe(ds41_gpu_graph *g) { +/* HC expand of an already summed and rounded FFN block (g->block). */ +static bool ds41_hc_expand_after_moe(ds41_gpu_graph *g) { +#if defined(__APPLE__) + if (ds41_hc_fused(g)) + return ds4_gpu_dsv41_hc_expand4(g->residual, g->block, g->after_attn, g->ffn_split, g->pre, NULL, NULL, DS4_N_EMBD); +#endif return ds4_gpu_hc_expand_split_tensor(g->residual, g->block, g->after_attn, g->ffn_split, DS4_N_EMBD, DS4_N_HC) && ds41_bf16(g->residual, DS4_N_EMBD * DS4_N_HC) && ds4_gpu_tensor_copy(g->pre, 0, g->ffn_split, 0, DS4_N_HC * sizeof(float)); } +/* The layer's FFN tail after ds41_moe: with the fused MoE glue this is where + * the routed + shared sum and its rounding run, inside the expand dispatch; + * otherwise ds41_moe_finish did them. */ +static bool ds41_graph_after_moe(ds41_gpu_graph *g, const ds4_model *m, + const ds4_layer_weights *l) { +#if defined(__APPLE__) + if (ds41_moe_fused(g, l)) + return ds4_gpu_dsv41_hc_expand4(g->residual, g->routed, g->after_attn, g->ffn_split, + g->pre, g->shared, g->block, DS4_N_EMBD) != 0; + (void)m; +#else + (void)m; (void)l; +#endif + return ds41_hc_expand_after_moe(g); +} + static bool ds41_graph_layer(ds41_gpu_graph *g, const ds4_model *m, const ds4_layer_weights *l, uint32_t il, int token) { return ds41_graph_before_moe(g, m, l, il) && ds41_moe(g, m, l, il, (uint32_t)token) && - ds41_graph_after_moe(g); + ds41_graph_after_moe(g, m, l); } #if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) @@ -40164,7 +40343,7 @@ static bool ds41_decode_island(ds41_gpu_graph *g, const ds4_model *m, const int state = ds4_gpu_decode_graph_begin(&key); if (state == 1) return true; const bool ok = island == 0 ? - ds41_graph_before_attention(g, m, l, il) && ds41_attention_project(g, m, l) : + ds41_graph_before_attention(g, m, l, il) && ds41_attention_project(g, m, l, il) : island == 2 ? ds41_attention_output(g, m, l) : ds41_graph_after_attention(g, m, l) && ds41_moe_partial(g, m, l, il, 0); if (state != 0) return ok; @@ -40187,7 +40366,7 @@ static bool ds41_graph_decode_layer(ds41_gpu_graph *g, const ds4_model *m, ds41_sum_partial(g, g->block, il, DS4_TP_GATE_ATTN) && ds41_bf16(g->block, DS4_N_EMBD) && ds41_decode_island(g, m, l, il, 1) && - ds41_moe_finish(g, il) && ds41_graph_after_moe(g); + ds41_moe_finish(g, l, il) && ds41_graph_after_moe(g, m, l); } #endif @@ -40901,7 +41080,7 @@ static bool ds41_graph_prefill_sweep(ds41_gpu_graph *g, const ds4_model *m, #undef DS41_USE_FFN_ROW ok = ds41_graph_after_attention(&row, m, l); if (ok && !batch_moe) ok = ds41_moe(&row, m, l, il, (uint32_t)tokens[off + t]) && - ds41_graph_after_moe(&row); + ds41_graph_after_moe(&row, m, l); } } if (ok && batch_moe) { @@ -40919,7 +41098,7 @@ static bool ds41_graph_prefill_sweep(ds41_gpu_graph *g, const ds4_model *m, DS41_PREFILL_ROWS(DS41_USE_MOE_ROW) #undef DS41_USE_MOE_ROW ok = ds4_gpu_add_tensor(row.block, row.routed, row.shared, DS4_N_EMBD) && - ds41_bf16(row.block, DS4_N_EMBD) && ds41_graph_after_moe(&row); + ds41_bf16(row.block, DS4_N_EMBD) && ds41_hc_expand_after_moe(&row); } } DS41_STAGE("hc expand"); @@ -41098,7 +41277,7 @@ static DS4_MAYBE_UNUSED bool ds41_graph_step_batch(ds41_gpu_graph *const *graphs row.q = queries[i]; row.heads = heads[i]; ok = ds41_attention(&row, model, l, il, true) && - ds41_attention_output(&row, model, l); + ds41_attention_output(&row, model, l, false); } if (ok) ok = ds41_sum_partial_batch(g, active.block, il, rows); if (ok) ok = ds4_gpu_dsv41_quantize(active.block, DS4_N_EMBD, rows, DS4_V41_BF16) && diff --git a/ds4_deepseek41_gpu.h b/ds4_deepseek41_gpu.h index 39f11699f8..2a7a5b0d87 100644 --- a/ds4_deepseek41_gpu.h +++ b/ds4_deepseek41_gpu.h @@ -123,6 +123,59 @@ int ds4_gpu_dsv41_projection_rows(ds4_gpu_tensor *out, int ds4_gpu_dsv41_gather_kv(ds4_gpu_tensor *out, const ds4_gpu_tensor *source, const ds4_gpu_tensor *ids, uint32_t source_rows, uint32_t selected_rows); +/* Decode HC glue for one token row, HC=4, byte-identical to the standalone + * split/collapse/BF16/norm/BF16 and expand/BF16 dispatch sequences. + * collapse_norm: split(mix) -> split; x = bf16(pre-weighted collapse of + * residual); norm = bf16(rmsnorm(x) * weight). `pre` is the coefficient row + * of the PREVIOUS sublayer's split (the first four floats of that tensor). + * expand4: out = bf16(post/comb expand of block into residual); when `pre` + * is given, split[0..3] is also copied into it; when `add` is given, block + * is first bf16(block + add) and that sum is written to block_sum. */ +int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, + const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, + const ds4_gpu_tensor *residual, + const void *model_map, uint64_t model_size, + uint64_t scale_offset, uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters, + float hc_eps, float norm_eps); +int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, + const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, + ds4_gpu_tensor *pre, const ds4_gpu_tensor *add, + ds4_gpu_tensor *block_sum, uint32_t n_embd); +/* Decode MoE glue for one token row, byte-identical to the standalone + * sequences: router = F32 logits matvec + softplus/sqrt + bias + canonical + * top-k + normalized weights in one dispatch; shared gate/up = two Q8_0 + * matvecs + BF16 + SwiGLU + BF16; shared down + HC tail = Q8_0 matvec + + * BF16 + (routed + shared) + BF16 + post/comb expand + BF16 (+ pre carry). */ +int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, ds4_gpu_tensor *logits, + const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint64_t bias_offset, bool has_bias, + uint32_t n_embd, uint32_t n_expert, uint32_t n_used, float scale); +int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, + uint32_t n_embd, uint32_t n_ff, float clamp); +/* Decode attention glue for one token row, byte-identical to the standalone + * sequences. matvec_bf16: a Q8_0 or F16 single-row matvec whose store is + * the BF16 rounding (1 done, 0 not covered, -1 error). qkv_norm_kv_tail: + * the q LoRA and KV weighted norms with their roundings, then the KV RoPE, + * FP8 (E8M0) block quantization and the copy into the raw window row at + * window_offset. bf16_rope: the rounding pass over whole rows plus the + * (inverse) RoPE on their last 64 values. */ +enum { DS4_V41_WEIGHT_F16 = 1, DS4_V41_WEIGHT_Q8_0 = 8 }; /* GGUF tensor type ids */ +int ds4_gpu_dsv41_matvec_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint32_t type, uint32_t in_dim, uint32_t out_dim, + const ds4_gpu_tensor *x); +int ds4_gpu_dsv41_qkv_norm_kv_tail(ds4_gpu_tensor *q, ds4_gpu_tensor *kv, ds4_gpu_tensor *window, + uint64_t window_offset, const void *model_map, uint64_t model_size, + uint64_t q_weight_offset, uint64_t kv_weight_offset, + uint32_t q_n, uint32_t kv_n, float eps, uint32_t pos, + bool compressed); +int ds4_gpu_dsv41_bf16_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, bool compressed, bool inverse); #ifdef __cplusplus } diff --git a/ds4_metal.m b/ds4_metal.m index b61cbee8d7..0bb0899500 100644 --- a/ds4_metal.m +++ b/ds4_metal.m @@ -44814,7 +44814,10 @@ int ds4_gpu_hc_rms_norm_mix_f16_tensor( uint32_t out_dim, float eps) { if (!g_initialized && !ds4_gpu_init()) return 0; - if (!out || !x || !model_map || n != 16384u || out_dim != 24u) return 0; + /* 16384 = V4's 4x4096 HC row, 20480 = V4.1's 4x5120: both norm with 1024 + * threads and both are whole multiples of the matvec's 4096-value stride, + * so the kernel's virtual-thread replica and empty tail hold for either. */ + if (!out || !x || !model_map || (n != 16384u && n != 20480u) || out_dim != 24u) return 0; @autoreleasepool { const uint64_t row_bytes = (uint64_t)n * sizeof(uint16_t); @@ -47578,14 +47581,9 @@ static bool dsv41_tensor_has_floats(const ds4_gpu_tensor *tensor, uint64_t count /* Frequency rounding errors accumulate into phase errors at long contexts. * Preserve the reference's pow/reciprocal and YaRN operation order here. */ #pragma float_control(precise, on, push) -int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, - uint32_t rows, uint32_t start, uint32_t stride, - bool compressed, bool inverse) { - if (width < 64 || !heads || !rows || rows > 1048576 || !stride || - (uint64_t)start + (uint64_t)(rows - 1u) * stride >= 1048576u || - (uint64_t)heads * rows > UINT64_MAX / width || - !dsv41_tensor_has_floats(x, (uint64_t)width * heads * rows)) return 0; - if (!g_initialized && !ds4_gpu_init()) return 0; +/* The released V4.1 RoPE frequencies: plain (base 10000) and the YaRN-style + * compressed table (base 160000 with the 1/16 interpolation ramp). */ +static const float *dsv41_rope_frequencies(bool compressed) { static float frequencies[2][32]; static dispatch_once_t once; dispatch_once(&once, ^{ @@ -47607,14 +47605,25 @@ int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, } } }); + return frequencies[compressed ? 1 : 0]; +} + +static int dsv41_rope_dispatch(const char *kernel, ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, uint32_t stride, + bool compressed, bool inverse) { + if (width < 64 || !heads || !rows || rows > 1048576 || !stride || + (uint64_t)start + (uint64_t)(rows - 1u) * stride >= 1048576u || + (uint64_t)heads * rows > UINT64_MAX / width || + !dsv41_tensor_has_floats(x, (uint64_t)width * heads * rows)) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; @autoreleasepool { - id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_rope"); + id pipeline = ds4_gpu_get_pipeline(kernel); if (!pipeline) return 0; struct { uint32_t width, heads, rows, start, inverse, stride; float frequencies[32]; } args = {width, heads, rows, start, inverse, stride, {0}}; - memcpy(args.frequencies, frequencies[compressed ? 1 : 0], sizeof(args.frequencies)); + memcpy(args.frequencies, dsv41_rope_frequencies(compressed), sizeof(args.frequencies)); int owned = 0; id cb = ds4_gpu_command_buffer(&owned); id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; @@ -47630,6 +47639,126 @@ int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, } #pragma float_control(pop) +int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, uint32_t stride, + bool compressed, bool inverse) { + return dsv41_rope_dispatch("kernel_dsv41_rope", x, width, heads, rows, start, stride, + compressed, inverse); +} + +int ds4_gpu_dsv41_bf16_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, + uint32_t rows, uint32_t start, bool compressed, bool inverse) { + if (width % 32u) return 0; + return dsv41_rope_dispatch("kernel_dsv41_bf16_rope", x, width, heads, rows, start, 1, + compressed, inverse); +} + +int ds4_gpu_dsv41_qkv_norm_kv_tail(ds4_gpu_tensor *q, ds4_gpu_tensor *kv, ds4_gpu_tensor *window, + uint64_t window_offset, const void *model_map, uint64_t model_size, + uint64_t q_weight_offset, uint64_t kv_weight_offset, + uint32_t q_n, uint32_t kv_n, float eps, uint32_t pos, + bool compressed) { + const uint64_t q_bytes = (uint64_t)q_n * sizeof(float), kv_bytes = (uint64_t)kv_n * sizeof(float); + if (!q_n || q_n % 4u || kv_n < 64u || kv_n % 32u || kv_n > q_n || !model_map || + !dsv41_tensor_has_floats(q, q_n) || !dsv41_tensor_has_floats(kv, kv_n) || + !window || window_offset > ds4_gpu_tensor_bytes(window) || + kv_bytes > ds4_gpu_tensor_bytes(window) - window_offset || + q_weight_offset > model_size || q_bytes > model_size - q_weight_offset || + kv_weight_offset > model_size || kv_bytes > model_size - kv_weight_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_qkv_norm_kv_tail"); + if (!pipeline) return 0; + /* The q row's standalone norm thread count (its reduction tree); the + * KV row has fewer elements than that and only its own lanes hold any. */ + const NSUInteger threads = ds4_gpu_rms_norm_threads(q_n); + const NSUInteger shared_bytes = ((NSUInteger)q_n + 32u) * sizeof(float); + if (threads < ds4_gpu_rms_norm_threads(kv_n) || threads > pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0; + uint64_t q_inner = 0, kv_inner = 0; + id qwbuf = ds4_gpu_wrap_model_range(model_map, model_size, q_weight_offset, q_bytes, &q_inner); + id kvwbuf = ds4_gpu_wrap_model_range(model_map, model_size, kv_weight_offset, kv_bytes, &kv_inner); + if (!qwbuf || !kvwbuf) return 0; + struct { uint32_t q_n, kv_n; float eps; uint32_t pos; float frequencies[32]; } args = + {q_n, kv_n, eps, pos, {0}}; + memcpy(args.frequencies, dsv41_rope_frequencies(compressed), sizeof(args.frequencies)); + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(q) offset:ds4_gpu_tensor_offset(q) atIndex:1]; + [enc setBuffer:qwbuf offset:(NSUInteger)q_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(kv) offset:ds4_gpu_tensor_offset(kv) atIndex:3]; + [enc setBuffer:kvwbuf offset:(NSUInteger)kv_inner atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(window) + offset:ds4_gpu_tensor_offset(window) + (NSUInteger)window_offset atIndex:5]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(2, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 qkv norm + KV tail"); + } +} + +/* Single-row matvec with the BF16 rounding on the store. 1: done; 0: shape or + * type not covered (caller runs the standalone matvec + rounding); -1: error. */ +int ds4_gpu_dsv41_matvec_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint32_t type, uint32_t in_dim, uint32_t out_dim, + const ds4_gpu_tensor *x) { + if (!in_dim || !out_dim || !model_map || !dsv41_tensor_has_floats(x, in_dim) || + !dsv41_tensor_has_floats(out, out_dim)) return -1; + if (!g_initialized && !ds4_gpu_init()) return -1; + @autoreleasepool { + ds4_gpu_q8_0_matvec_args args; + ds4_gpu_mv_dispatch mv_dispatch; + uint64_t weight_bytes; + const char *kernel; + if (type == DS4_V41_WEIGHT_Q8_0) { + if (in_dim % 32u) return 0; + weight_bytes = ((uint64_t)in_dim / 32u) * 34u * out_dim; + args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim); + mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + if (out_dim > 65536u) mv_dispatch.nsg = 8; + kernel = "kernel_dsv41_mul_mv_q8_0_f32_bf16"; + } else if (type == DS4_V41_WEIGHT_F16) { + if (in_dim < 32u || in_dim % 4u) return 0; + weight_bytes = (uint64_t)in_dim * sizeof(uint16_t) * out_dim; + ds4_gpu_f16_matvec_args f16_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim); + mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0); + if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) { + mv_dispatch.nr0 = 4; + mv_dispatch.smem = 32u * 4u * sizeof(float); + } + memcpy(&args, &f16_args, sizeof(args)); + kernel = "kernel_dsv41_mul_mv_f16_f32_4_bf16"; + } else { + return 0; + } + if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return -1; + args.nr0 = mv_dispatch.nr0; + uint64_t inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner); + id pipeline = ds4_gpu_get_mul_mv_pipeline(kernel, mv_dispatch.nsg); + if (!wbuf || !pipeline) return -1; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return -1; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:3]; + [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 BF16-rounded matvec") ? 1 : -1; + } +} + int ds4_gpu_dsv41_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads, uint32_t rows, uint32_t start, bool compressed, bool inverse) { return ds4_gpu_dsv41_rope_stride(x, width, heads, rows, start, 1, compressed, inverse); @@ -47707,6 +47836,211 @@ int ds4_gpu_dsv41_engram_add(ds4_gpu_tensor *residual, } } +/* V4.1 decode HC glue (metal/dsv41.metal): one token row, HC=4. */ +typedef struct { uint32_t n_embd, sinkhorn_iters; float hc_eps, norm_eps; uint32_t copy_pre, has_add; } ds4_gpu_dsv41_hc_args; + +int ds4_gpu_dsv41_hc_collapse_norm(ds4_gpu_tensor *split, ds4_gpu_tensor *x, ds4_gpu_tensor *norm, + const ds4_gpu_tensor *mix, const ds4_gpu_tensor *pre, + const ds4_gpu_tensor *residual, + const void *model_map, uint64_t model_size, + uint64_t scale_offset, uint64_t base_offset, + uint64_t norm_weight_offset, + uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters, + float hc_eps, float norm_eps) { + const uint64_t weight_bytes = (uint64_t)n_embd * sizeof(float); + /* sinkhorn_iters == 0: no split (the logits' collapse), mix/split/scale/base unused. */ + const bool with_split = sinkhorn_iters != 0; + if (n_hc != 4u || !n_embd || n_embd % 4u || !model_map || + (with_split && (!dsv41_tensor_has_floats(split, 24) || !dsv41_tensor_has_floats(mix, 24) || + scale_offset > model_size || 12u > model_size - scale_offset || + base_offset > model_size || 96u > model_size - base_offset)) || + !dsv41_tensor_has_floats(pre, 4) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(norm, n_embd) || + norm_weight_offset > model_size || weight_bytes > model_size - norm_weight_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_collapse_norm4"); + if (!pipeline) return 0; + /* The norm's reduction tree is the standalone kernel's only at its thread count. */ + const NSUInteger threads = ds4_gpu_rms_norm_threads(n_embd); + const NSUInteger shared_bytes = ((NSUInteger)n_embd + 32u) * sizeof(float); + if (threads > pipeline.maxTotalThreadsPerThreadgroup || + shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0; + uint64_t scale_inner = 0, base_inner = 0, norm_inner = 0; + id normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, + weight_bytes, &norm_inner); + id scalebuf = with_split ? + ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, 12u, &scale_inner) : normwbuf; + id basebuf = with_split ? + ds4_gpu_wrap_model_range(model_map, model_size, base_offset, 96u, &base_inner) : normwbuf; + if (!scalebuf || !basebuf || !normwbuf) return 0; + if (!with_split) { mix = x; split = x; scale_inner = base_inner = norm_inner; } + const ds4_gpu_dsv41_hc_args args = {n_embd, sinkhorn_iters, hc_eps, norm_eps, 0, 0}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:ds4_gpu_tensor_buffer(mix) offset:ds4_gpu_tensor_offset(mix) atIndex:1]; + [enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2]; + [enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(pre) offset:ds4_gpu_tensor_offset(pre) atIndex:4]; + [enc setBuffer:ds4_gpu_tensor_buffer(residual) offset:ds4_gpu_tensor_offset(residual) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(split) offset:ds4_gpu_tensor_offset(split) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:7]; + [enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:8]; + [enc setBuffer:ds4_gpu_tensor_buffer(norm) offset:ds4_gpu_tensor_offset(norm) atIndex:9]; + [enc setThreadgroupMemoryLength:shared_bytes atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 HC collapse/norm"); + } +} + +int ds4_gpu_dsv41_hc_expand4(ds4_gpu_tensor *out, const ds4_gpu_tensor *block, + const ds4_gpu_tensor *residual, const ds4_gpu_tensor *split, + ds4_gpu_tensor *pre, const ds4_gpu_tensor *add, + ds4_gpu_tensor *block_sum, uint32_t n_embd) { + if (!n_embd || !dsv41_tensor_has_floats(out, 4ull * n_embd) || + !dsv41_tensor_has_floats(block, n_embd) || !dsv41_tensor_has_floats(residual, 4ull * n_embd) || + !dsv41_tensor_has_floats(split, 24) || (pre && !dsv41_tensor_has_floats(pre, 4)) || + (add != NULL) != (block_sum != NULL) || (add && !dsv41_tensor_has_floats(add, n_embd)) || + (block_sum && !dsv41_tensor_has_floats(block_sum, n_embd))) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_expand4_bf16"); + if (!pipeline) return 0; + const ds4_gpu_dsv41_hc_args args = {n_embd, 0, 0.0f, 0.0f, pre != NULL, add != NULL}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + const ds4_gpu_tensor *buffers[] = {block, residual, split, out, pre ? pre : split, + add ? add : block, block_sum ? block_sum : out}; + for (NSUInteger i = 0; i < 7; i++) + [enc setBuffer:ds4_gpu_tensor_buffer(buffers[i]) + offset:ds4_gpu_tensor_offset(buffers[i]) atIndex:i + 1]; + [enc dispatchThreadgroups:MTLSizeMake(((uint64_t)n_embd + 255u) / 256u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 HC expand"); + } +} + +/* V4.1 decode MoE glue (metal/dsv41.metal): one token row. */ +int ds4_gpu_dsv41_router_select(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, + ds4_gpu_tensor *probs, ds4_gpu_tensor *logits, + const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t weight_offset, uint64_t bias_offset, bool has_bias, + uint32_t n_embd, uint32_t n_expert, uint32_t n_used, float scale) { + const uint64_t weight_bytes = (uint64_t)n_embd * n_expert * sizeof(float); + const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float); + if (!n_embd || n_embd % 32u || !n_expert || n_expert % 2u || n_expert > 512u || + !n_used || n_used > 8u || n_used > n_expert || !model_map || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(logits, n_expert) || + !dsv41_tensor_has_floats(probs, n_expert) || !dsv41_tensor_has_floats(weights, n_used) || + !selected || ds4_gpu_tensor_bytes(selected) < (uint64_t)n_used * sizeof(int32_t) || + weight_offset > model_size || weight_bytes > model_size - weight_offset || + (has_bias && (bias_offset > model_size || bias_bytes > model_size - bias_offset))) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + id pipeline = ds4_gpu_get_pipeline("kernel_dsv41_router_select"); + if (!pipeline || pipeline.maxTotalThreadsPerThreadgroup < 256u || !g_dsv4_completion_cache) return 0; + uint64_t weight_inner = 0, bias_inner = 0; + id wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &weight_inner); + id biasbuf = has_bias ? + ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner) : wbuf; + if (!wbuf || !biasbuf) return 0; + id logitsbuf = ds4_gpu_tensor_buffer(logits); + NSString *completion_key = [NSString stringWithFormat:@"v41router:%p:%llu", + (void *)logitsbuf, (unsigned long long)ds4_gpu_tensor_offset(logits)]; + id completion = [g_dsv4_completion_cache objectForKey:completion_key]; + if (!completion) { + completion = [g_device newBufferWithLength:sizeof(uint32_t) options:MTLResourceStorageModeShared]; + if (!completion) return 0; + *((uint32_t *)[completion contents]) = 0u; + [g_dsv4_completion_cache setObject:completion forKey:completion_key]; + } + [g_transient_buffers addObject:completion]; + /* One dispatch (last-arriving threadgroup selects) on M5 only, like + * V4's fused router: on the M3 Ultra the other groups' logits were + * not reliably visible to the last one. Elsewhere the standalone + * F32 matvec runs first and one threadgroup selects. */ + const bool fused_matvec = ds4_gpu_device_is_m5_apple_silicon() && + getenv("DS4_METAL_DISABLE_V41_ROUTER_SINGLE_DISPATCH") == NULL; + if (!fused_matvec && + !ds4_gpu_matmul_f32_tensor(logits, model_map, model_size, weight_offset, + n_embd, n_expert, x, 1)) return 0; + const struct { uint32_t n_embd, n_expert, n_used, has_bias; float scale; uint32_t fused_matvec; } args = + {n_embd, n_expert, n_used, has_bias, scale, fused_matvec}; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:wbuf offset:(NSUInteger)weight_inner atIndex:1]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2]; + [enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(probs) offset:ds4_gpu_tensor_offset(probs) atIndex:4]; + [enc setBuffer:biasbuf offset:(NSUInteger)(has_bias ? bias_inner : weight_inner) atIndex:5]; + [enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:6]; + [enc setBuffer:ds4_gpu_tensor_buffer(weights) offset:ds4_gpu_tensor_offset(weights) atIndex:7]; + [enc setBuffer:completion offset:0 atIndex:8]; + [enc setThreadgroupMemoryLength:(512u + 32u) * sizeof(float) atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(fused_matvec ? (n_expert + 1u) / 2u : 1u, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 router select"); + } +} + +int ds4_gpu_dsv41_shared_gate_up_swiglu(ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, + const void *model_map, uint64_t model_size, + uint64_t gate_offset, uint64_t up_offset, + uint32_t n_embd, uint32_t n_ff, float clamp) { + const uint64_t row_bytes = ((uint64_t)n_embd / 32u) * 34u; + const uint64_t weight_bytes = row_bytes * n_ff; + if (!n_embd || n_embd % 32u || !n_ff || !model_map || !isfinite(clamp) || clamp < 0.0f || + !dsv41_tensor_has_floats(x, n_embd) || !dsv41_tensor_has_floats(mid, n_ff) || + gate_offset > model_size || weight_bytes > model_size - gate_offset || + up_offset > model_size || weight_bytes > model_size - up_offset) return 0; + if (!g_initialized && !ds4_gpu_init()) return 0; + @autoreleasepool { + uint64_t gate_inner = 0, up_inner = 0; + id gate_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner); + id up_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner); + if (!gate_wbuf || !up_wbuf) return 0; + ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(n_embd, n_ff); + ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch(); + args.nr0 = mv_dispatch.nr0; + id pipeline = + ds4_gpu_get_mul_mv_pipeline("kernel_dsv41_shared_gate_up_swiglu_q8_0", mv_dispatch.nsg); + if (!pipeline) return 0; + int owned = 0; + id cb = ds4_gpu_command_buffer(&owned); + id enc = cb ? ds4_gpu_compute_encoder(cb) : nil; + if (!enc) return 0; + [enc setComputePipelineState:pipeline]; + [enc setBytes:&args length:sizeof(args) atIndex:0]; + [enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1]; + [enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2]; + [enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:3]; + [enc setBuffer:ds4_gpu_tensor_buffer(mid) offset:ds4_gpu_tensor_offset(mid) atIndex:4]; + [enc setBytes:&clamp length:sizeof(clamp) atIndex:5]; + [enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0]; + [enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_ff + (NSUInteger)mv_dispatch.nr0 - 1u) / + (NSUInteger)mv_dispatch.nr0, 1, 1) + threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)]; + ds4_gpu_end_compute_encoder(cb, enc); + return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 shared gate/up"); + } +} + int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset, ds4_gpu_tensor *plain, uint32_t width, uint32_t rows, uint32_t format, bool pack) { diff --git a/metal/dsv41.metal b/metal/dsv41.metal index 03afd310dd..e2ce1e66c2 100644 --- a/metal/dsv41.metal +++ b/metal/dsv41.metal @@ -291,3 +291,691 @@ kernel void kernel_dsv41_indexer_scores_packed( } } #endif + +/* Decode-time V4.1 hyper-connection glue, one token row, HC=4. + * + * The released graph runs each half-layer's HC work as separate dispatches: + * split/sinkhorn, the pre-weighted collapse of the four streams, a BF16 + * rounding pass, the weighted RMSNorm and another BF16 pass (and, after the + * sublayer, the post/comb expand plus its BF16 pass). DeepSeek's production + * decode does the whole thing in one "Mega-mHC" kernel. These two kernels + * are that fusion for ds4's graph: every reduction keeps the standalone + * kernel's thread mapping and accumulation order, and every rounding point + * is the same dsv41_bf16, so the outputs are byte-identical (checked by + * tests/test_deepseek41_metal --hc-fuse). */ +struct ds4_metal_args_dsv41_hc { + uint n_embd; + uint sinkhorn_iters; + float hc_eps; + float norm_eps; + uint copy_pre; + uint has_add; +}; + +static inline float4 dsv41_bf16x4(float4 v) { + return float4(dsv41_bf16(v.x), dsv41_bf16(v.y), dsv41_bf16(v.z), dsv41_bf16(v.w)); +} + +/* kernel_dsv4_hc_split_sinkhorn's HC == 4 body, verbatim, on one lane. */ +static inline void dsv41_hc_split4(device const float *mix, device const float *scale, + device const float *base, device float *out, + uint sinkhorn_iters, float epsv) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + + const float4 pre_z = *((device const float4 *)mix) * pre_scale + *((device const float4 *)base); + *((device float4 *)out) = ds4_hc_sigmoid(pre_z) + epsv; + + const float4 post_z = *((device const float4 *)(mix + 4)) * post_scale + *((device const float4 *)(base + 4)); + *((device float4 *)(out + 4)) = ds4_hc_twice_sigmoid(post_z); + + float4 r0 = *((device const float4 *)(mix + 8)) * comb_scale + *((device const float4 *)(base + 8)); + float4 r1 = *((device const float4 *)(mix + 12)) * comb_scale + *((device const float4 *)(base + 12)); + float4 r2 = *((device const float4 *)(mix + 16)) * comb_scale + *((device const float4 *)(base + 16)); + float4 r3 = *((device const float4 *)(mix + 20)) * comb_scale + *((device const float4 *)(base + 20)); + + const float m0 = max(max(r0.x, r0.y), max(r0.z, r0.w)); + const float m1 = max(max(r1.x, r1.y), max(r1.z, r1.w)); + const float m2 = max(max(r2.x, r2.y), max(r2.z, r2.w)); + const float m3 = max(max(r3.x, r3.y), max(r3.z, r3.w)); + + r0 = exp(r0 - m0); + r1 = exp(r1 - m1); + r2 = exp(r2 - m2); + r3 = exp(r3 - m3); + + r0 = r0 * (1.0f / (r0.x + r0.y + r0.z + r0.w)) + epsv; + r1 = r1 * (1.0f / (r1.x + r1.y + r1.z + r1.w)) + epsv; + r2 = r2 * (1.0f / (r2.x + r2.y + r2.z + r2.w)) + epsv; + r3 = r3 * (1.0f / (r3.x + r3.y + r3.z + r3.w)) + epsv; + + float4 col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + + for (uint iter = 1; iter < sinkhorn_iters; ++iter) { + r0 *= 1.0f / (r0.x + r0.y + r0.z + r0.w + epsv); + r1 *= 1.0f / (r1.x + r1.y + r1.z + r1.w + epsv); + r2 *= 1.0f / (r2.x + r2.y + r2.z + r2.w + epsv); + r3 *= 1.0f / (r3.x + r3.y + r3.z + r3.w + epsv); + + col_inv = 1.0f / (r0 + r1 + r2 + r3 + epsv); + r0 *= col_inv; + r1 *= col_inv; + r2 *= col_inv; + r3 *= col_inv; + } + + *((device float4 *)(out + 8)) = r0; + *((device float4 *)(out + 12)) = r1; + *((device float4 *)(out + 16)) = r2; + *((device float4 *)(out + 20)) = r3; +} + +/* split(mix) -> split_out; x = bf16(sum_h pre[h] * residual[h]); norm = + * bf16(rmsnorm(x) * weight). `pre` is the coefficient row the previous + * sublayer produced (V4.1's single-pass mHC), not this split's. One + * threadgroup of ds4_gpu_rms_norm_threads(n_embd) threads: the collapse + * keeps kernel_dsv4_hc_weighted_sum's per-stream order and the norm keeps + * kernel_rms_norm_mul_f32_4's loop, simd tree and 1/sqrt scale. */ +kernel void kernel_dsv41_hc_collapse_norm4( + constant ds4_metal_args_dsv41_hc &args, + device const float *mix, + device const float *scale, + device const float *base, + device const float *pre, + device const float4 *residual, + device float *split_out, + device float4 *x_out, + device const float4 *norm_weight, + device float4 *norm_out, + threadgroup float *shared [[threadgroup(0)]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + const uint n4 = args.n_embd >> 2; + threadgroup float4 *row = (threadgroup float4 *)shared; + threadgroup float *sums = shared + args.n_embd; + + if (tid == 0 && args.sinkhorn_iters) dsv41_hc_split4(mix, scale, base, split_out, args.sinkhorn_iters, args.hc_eps); + if (sgitg == 0) sums[tiisg] = 0.0f; + + const float4 p = *((device const float4 *)pre); + device const float4 *x0 = residual; + device const float4 *x1 = residual + n4; + device const float4 *x2 = residual + 2u * n4; + device const float4 *x3 = residual + 3u * n4; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + float4 v = 0.0f; + v += x0[i] * p.x; + v += x1[i] * p.y; + v += x2[i] * p.z; + v += x3[i] * p.w; + v = dsv41_bf16x4(v); + row[i] = v; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) sums[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(sums[tiisg]); + + const float mean = sumf / (float)args.n_embd; + const float norm_scale = 1.0f / sqrt(mean + args.norm_eps); + for (uint i = tid; i < n4; i += ntg) { + const float4 v = row[i]; + x_out[i] = v; + norm_out[i] = dsv41_bf16x4((v * norm_scale) * norm_weight[i]); + } +} + +/* out[h] = bf16(post[h] * block + sum_s comb[h, s] * residual[s]) for the + * four streams, kernel_dsv4_hc_expand4's index arithmetic and accumulation + * order; copy_pre also carries split[0..3] into `pre` for the next layer. + * has_add: block = bf16(block_in + add) first (the routed + shared sum and + * its rounding), written back to block_sum. */ +kernel void kernel_dsv41_hc_expand4_bf16( + constant ds4_metal_args_dsv41_hc &args, + device const float *block_in, + device const float *residual, + device const float *split, + device float *out, + device float *pre_out, + device const float *add, + device float *block_sum, + uint d [[thread_position_in_grid]]) { + if (d >= args.n_embd) return; + device const float *post = split + 4; + device const float *comb = split + 8; + + float block_v = block_in[d]; + if (args.has_add) { + block_v += add[d]; + block_v = dsv41_bf16(block_v); + block_sum[d] = block_v; + } + const float r0 = residual[d]; + const float r1 = residual[d + args.n_embd]; + const float r2 = residual[d + 2u * args.n_embd]; + const float r3 = residual[d + 3u * args.n_embd]; + + for (uint dst_hc = 0; dst_hc < 4; ++dst_hc) { + float acc = block_v * post[dst_hc]; + acc += comb[dst_hc + 0u * 4u] * r0; + acc += comb[dst_hc + 1u * 4u] * r1; + acc += comb[dst_hc + 2u * 4u] * r2; + acc += comb[dst_hc + 3u * 4u] * r3; + out[d + dst_hc * args.n_embd] = dsv41_bf16(acc); + } + if (args.copy_pre && d < 4) pre_out[d] = split[d]; +} + +/* Single-row matvecs that round to BF16 on the store: the V4.1 graph rounds + * almost every projection output, as a separate dispatch over the result. + * Bodies are kernel_mul_mv_q8_0_f32_impl and kernel_mul_mv_t_t_4_impl + * with helper_mv_reduce_and_write's tree; only + * the store changes (tests/test_deepseek41_metal --attn-fuse). */ +template +static inline void dsv41_mv_reduce_write_bf16( + device float * dst_f32, float sumf[NR0], const int r0, const int ne01, + ushort tiisg, ushort sgitg, threadgroup char * shmem) { + constexpr short NW = N_SIMDWIDTH; + threadgroup float * shmem_f32[NR0]; + for (short row = 0; row < NR0; ++row) { + shmem_f32[row] = (threadgroup float *) shmem + NW*row; + if (sgitg == 0) shmem_f32[row][tiisg] = 0.0f; + sumf[row] = simd_sum(sumf[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (short row = 0; row < NR0; ++row) { + if (tiisg == 0) shmem_f32[row][sgitg] = sumf[row]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (short row = 0; row < NR0 && r0 + row < ne01; ++row) { + float tot = simd_sum(shmem_f32[row][tiisg]); + if (tiisg == 0 && sgitg == 0) dst_f32[r0 + row] = dsv41_bf16(tot); + } +} + +kernel void kernel_dsv41_mul_mv_q8_0_f32_bf16( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short NR0 = N_R0_Q8_0; + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + const int nb = args.ne00/QK8_0; + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + device const float * y = (device const float *) (src1 + offset1); + device const block_q8_0 * ax[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0); + } + float sumf[NR0] = { 0.f }; + const short ix = tiisg/(NW/NQ); + const short il = tiisg%(NW/NQ); + const int ib0 = sgitg*NQ + ix; + float yl[NQ]; + device const float * yb = y + ib0*QK8_0 + il*NQ; + for (int ib = ib0; ib < nb; ib += NSG*NQ) { + for (short i = 0; i < NQ; ++i) yl[i] = yb[i]; + for (short row = 0; row < NR0; row++) { + device const int8_t * qs = ax[row][ib].qs + il*NQ; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) sumq += qs[i] * yl[i]; + sumf[row] += sumq*ax[row][ib].d; + } + yb += NSG*NQ*QK8_0; + } + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + dsv41_mv_reduce_write_bf16(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +template +void dsv41_mul_mv_f16_f32_4_bf16_impl( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + const int nb = args.ne00/NB; + const int r0 = tgpig.x*NR0; + const int r1 = tgpig.y; + const int im = tgpig.z; + const uint i12 = im%args.ne12; + const uint i13 = im/args.ne12; + const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13; + device const float * y = (device const float *) (src1 + offset1); + device const float4 * y4 = (device const float4 *) (src1 + offset1); + device const half * ax [NR0]; + device const half4 * ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/args.r2)*args.nb02 + (i13/args.r3)*args.nb03; + ax [row] = (device const half *) ((device char *) src0 + offset0); + ax4[row] = (device const half4 *) ((device char *) src0 + offset0); + } + float sumf[NR0] = { 0.f }; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + float4 yl4[NF4]; + device const float4 * yb4 = y4 + (ib0*NB + il*NF)/4; + for (int ib = ib0; ib < nb; ib += NSG*NF) { + for (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + for (short row = 0; row < NR0; row++) { + device const half4 * xb4 = ax4[row] + (ib*NB + il*NF)/4; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; + } + for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) { + for (short row = 0; row < NR0; row++) sumf[row] += ax[row][i] * y[i]; + } + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + dsv41_mv_reduce_write_bf16(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem); +} + +kernel void kernel_dsv41_mul_mv_f16_f32_4_bf16( + constant ds4_metal_args_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + switch (args.nr0) { + case 2: dsv41_mul_mv_f16_f32_4_bf16_impl<2>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + case 4: dsv41_mul_mv_f16_f32_4_bf16_impl<4>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break; + } +} + +/* BF16 rounding of a whole row followed by the adjacent-pair RoPE on its + * last 64 values (kernel_dsv41_rope), one simdgroup per row: what the graph + * ran as a rounding pass over all heads plus the inverse RoPE. */ +kernel void kernel_dsv41_bf16_rope( + constant ds4_metal_args_dsv41_rope &args, + device float *x, + uint2 group [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const ulong row = ((ulong)group.y * args.heads + group.x) * args.width; + for (uint i = lane; i < args.width; i += 32u) x[row + i] = dsv41_bf16(x[row + i]); + simdgroup_barrier(mem_flags::mem_device); + const float theta = float(args.start + group.y * args.stride) * args.frequencies[lane]; + const float c = precise::cos(theta); + const float s = args.inverse ? -precise::sin(theta) : precise::sin(theta); + const ulong i = row + args.width - 64u + 2u * lane; + const float re = x[i], im = x[i + 1u]; + x[i] = dsv41_bf16(re * c - im * s); + x[i + 1u] = dsv41_bf16(re * s + im * c); +} + +/* The attention projections' normalization and the KV tail in one dispatch: + * threadgroup y=0 is kernel_rms_norm_mul_f32_4 over the q LoRA row plus its + * BF16 rounding; y=1 the same over the 512-wide KV row, then the RoPE on + * its last 64 values, the FP8 (E8M0-scaled) block quantization of + * kernel_dsv41_quantize mode 1 and the store into the raw window slot. + * The threadgroup is the q row's norm thread count; the KV lanes beyond its + * own count hold no elements, so their zero partials leave the standalone + * reduction tree unchanged. */ +struct ds4_metal_args_dsv41_qkv_tail { + uint q_n; + uint kv_n; + float eps; + uint pos; + float frequencies[32]; +}; + +kernel void kernel_dsv41_qkv_norm_kv_tail( + constant ds4_metal_args_dsv41_qkv_tail &args, + device float4 *q, + device const float4 *q_weight, + device float4 *kv, + device const float4 *kv_weight, + device float *window_row, + threadgroup float *shared [[threadgroup(0)]], + ushort task [[threadgroup_position_in_grid]], + ushort tid [[thread_position_in_threadgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort ntg [[threads_per_threadgroup]]) { + const bool kv_task = task != 0; + const uint n = kv_task ? args.kv_n : args.q_n; + const uint n4 = n >> 2; + device float4 *x = kv_task ? kv : q; + device const float4 *w = kv_task ? kv_weight : q_weight; + threadgroup float4 *row = (threadgroup float4 *)shared; + threadgroup float *sums = shared + 4u * n4; + + if (sgitg == 0) sums[tiisg] = 0.0f; + float sumf = 0.0f; + for (uint i = tid; i < n4; i += ntg) { + const float4 v = x[i]; + sumf += dot(v, v); + } + sumf = simd_sum(sumf); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) sums[sgitg] = sumf; + threadgroup_barrier(mem_flags::mem_threadgroup); + sumf = simd_sum(sums[tiisg]); + const float mean = sumf / (float)n; + const float scale = 1.0f / sqrt(mean + args.eps); + for (uint i = tid; i < n4; i += ntg) { + const float4 v = dsv41_bf16x4((x[i] * scale) * w[i]); + row[i] = v; + if (!kv_task) x[i] = v; + } + if (!kv_task) return; + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float *rowf = (threadgroup float *)row; + if (sgitg == 0) { + const float theta = float(args.pos) * args.frequencies[tiisg]; + const float c = precise::cos(theta); + const float s = precise::sin(theta); + const uint i = n - 64u + 2u * tiisg; + const float re = rowf[i], im = rowf[i + 1u]; + rowf[i] = dsv41_bf16(re * c - im * s); + rowf[i + 1u] = dsv41_bf16(re * s + im * c); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const uint nsg = (ntg + 31u) / 32u; + for (uint block = sgitg; block < n / 32u; block += nsg) { + const uint column = block * 32u + tiisg; + const float value = dsv41_bf16(rowf[column]); + const float amax = simd_max(abs(value)); + const float qscale = dsv41_pow2_ceil(max(amax, 1.0e-4f) * (1.0f / 448.0f)); + const float result = copysign(dsv4_e4m3fn_dequant(abs(value) / qscale), value) * qscale; + const float out = dsv41_bf16(result); + ((device float *)kv)[column] = out; + window_row[column] = out; + } +} + +/* Decode-time V4.1 MoE glue, one token row. Same discipline as the HC + * kernels above: every reduction keeps the standalone kernel's shape and + * every BF16 rounding point is kept (tests/test_deepseek41_metal --moe-fuse). */ + +struct ds4_metal_args_dsv41_router { + uint n_embd; + uint n_expert; + uint n_used; + uint has_bias; + float scale; + uint fused_matvec; +}; + +/* Router logits (F32 matvec, kernel_mul_mv_f32_f32_4 with nsg=8, nr0=2) and + * the generic select sequence the V4.1 graph ran as ten dispatches: + * probs = sqrt(softplus(logits)); score = probs + bias; top-k in the + * canonical (score desc, idx asc) order of kernel_argsort_f32_i32_desc_canon; + * weights = (probs[sel] / max(sum, 2^-14)) * scale, where sum is + * kernel_sum_rows_f32_f32's six-lane simd reduction. + * fused_matvec: the last-arriving threadgroup selects (DeepSeek's Mega-Gate + * shape); the host enables it on M5 only, like V4's fused router: on the + * M3 Ultra the other groups' logits were not reliably visible to the last + * one (nondeterministic probs). Otherwise one threadgroup selects from + * logits produced by the standalone matvec dispatch. */ +kernel void kernel_dsv41_router_select( + constant ds4_metal_args_dsv41_router &args, + device const float *weight, + device const float *x, + device float *logits, + device float *probs, + device const float *bias, + device int32_t *selected, + device float *weights, + device atomic_uint *completion, + threadgroup float *scratch [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tid [[thread_index_in_threadgroup]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + if (args.fused_matvec) { + constexpr short NSG = 8; + constexpr short NR0 = 2; + constexpr short NB = 32; + constexpr short NF = 16; + constexpr short NF4 = NF/4; + constexpr short NW = N_SIMDWIDTH; + const int nb = (int)args.n_embd/NB; + const int r0 = tgpig.x*NR0; + device const float4 *y4 = (device const float4 *)x; + device const float4 *ax4[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + ax4[row] = (device const float4 *)(weight + (uint64_t)(r0 + row)*args.n_embd); + } + float sumf[NR0] = {0.f}; + const short ix = tiisg/(NW/NF); + const short il = tiisg%(NW/NF); + const int ib0 = sgitg*NF + ix; + device const float4 *yb4 = y4 + (ib0*NB + il*NF)/4; + for (int ib = ib0; ib < nb; ib += NSG*NF) { + float4 yl4[NF4]; + FOR_UNROLL (short i = 0; i < NF4; ++i) yl4[i] = yb4[i]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const float4 *xb4 = ax4[row] + (ib*NB + il*NF)/4; + float sumq = 0.f; + FOR_UNROLL (short i = 0; i < NF4; ++i) sumq += dot(float4(xb4[i]), float4(yl4[i])); + sumf[row] += sumq; + } + yb4 += NSG*NF*NW/4; + } + for (int i = nb*NB + sgitg*NW + tiisg; i < (int)args.n_embd; i += NW*NSG) { + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sumf[row] += ((device const float *)ax4[row])[i] * x[i]; + } + } + helper_mv_reduce_and_write(logits, sumf, r0, (int)args.n_expert, + tiisg, sgitg, (threadgroup char *)scratch); + + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) { + const uint old = atomic_fetch_add_explicit(completion, 1u, memory_order_relaxed); + scratch[0] = old + 1u == (args.n_expert + 1u) / 2u ? 1.0f : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (scratch[0] == 0.0f) return; + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + } + + constexpr uint SLOTS = 512; + threadgroup float *score_tg = scratch; /* SLOTS */ + threadgroup float *red_s = scratch + SLOTS; /* 8 */ + threadgroup int32_t *red_i = (threadgroup int32_t *)(scratch + SLOTS + 8); /* 8 */ + threadgroup int32_t *sel_tg = (threadgroup int32_t *)(scratch + SLOTS + 16); /* 8 */ + threadgroup volatile float *stage = scratch + SLOTS + 24; /* 8 */ + + device volatile const float *lg = (device volatile const float *)logits; + for (uint i = tid; i < SLOTS; i += 256u) { + float s = -INFINITY; + if (i < args.n_expert) { + const float v = lg[i]; + const float sp = select(log(1.0f + exp(v)), v, v > 20.0f); + const float p = sqrt(sp); + probs[i] = p; + s = args.has_bias ? p + bias[i] : p; + } + score_tg[i] = s; + } + threadgroup_barrier(mem_flags::mem_device_and_threadgroup); + + float s0 = score_tg[tid], s1 = score_tg[tid + 256u]; + for (uint r = 0; r < args.n_used; r++) { + float bs = s0; + int32_t bi = (int32_t)tid; + if (s1 > bs) { bs = s1; bi = (int32_t)tid + 256; } + for (ushort o = 16; o > 0; o >>= 1) { + const float ps = simd_shuffle_xor(bs, o); + const int32_t pi = simd_shuffle_xor(bi, o); + if (ps > bs || (ps == bs && pi < bi)) { bs = ps; bi = pi; } + } + if (tiisg == 0) { red_s[sgitg] = bs; red_i[sgitg] = bi; } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid == 0) { + float best = red_s[0]; + int32_t best_i = red_i[0]; + for (uint g = 1; g < 8; g++) { + if (red_s[g] > best || (red_s[g] == best && red_i[g] < best_i)) { + best = red_s[g]; best_i = red_i[g]; + } + } + sel_tg[r] = best_i; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const int32_t sel = sel_tg[r]; + if ((int32_t)tid == sel) s0 = -INFINITY; + if ((int32_t)tid + 256 == sel) s1 = -INFINITY; + } + + /* kernel_sum_rows_f32_f32 runs six threads: lanes 0..5 hold one weight + * each and the simdgroup reduces them; then (w / clamp(sum)) * scale as + * two separately rounded ops (the div-row and mul-scalar kernels). */ + const bool lane = sgitg == 0 && tiisg < args.n_used; + device volatile const float *pr = (device volatile const float *)probs; + const float w = lane ? pr[sel_tg[tiisg]] : 0.0f; + float sum = simd_sum(w); + if (sgitg == 0 && tiisg == 0) { + sum = clamp(sum, 6.103515625e-5f, INFINITY); + stage[0] = sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < args.n_used) { + stage[1 + tid] = w / stage[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tid < args.n_used) { + selected[tid] = sel_tg[tid]; + weights[tid] = stage[1 + tid] * args.scale; + } + if (args.fused_matvec) { + threadgroup_barrier(mem_flags::mem_threadgroup); + atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device); + if (tid == 0) atomic_store_explicit(completion, 0u, memory_order_relaxed); + } +} + +/* Shared expert gate/up (Q8_0, kernel_mul_mv_q8_0_f32's walk and reduction + * tree at the dispatch's nsg) with V4.1's rounding: gate and up are rounded + * to BF16 before SwiGLU and the product after it. */ +kernel void kernel_dsv41_shared_gate_up_swiglu_q8_0( + constant ds4_metal_args_mul_mv & args, + device const char * src0_gate, + device const char * src0_up, + device const char * src1, + device char * dst_mid, + constant float &clamp_value, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short NR0 = N_R0_Q8_0; + const short NSG = FC_mul_mv_nsg; + constexpr short NW = N_SIMDWIDTH; + constexpr short NQ = 8; + + const int nb = args.ne00 / QK8_0; + const int r0 = tgpig.x * NR0; + device const float *y = (device const float *)src1; + + device const block_q8_0 *ag[NR0]; + device const block_q8_0 *au[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + const uint64_t offset0 = (r0 + row) * args.nb01; + ag[row] = (device const block_q8_0 *)(src0_gate + offset0); + au[row] = (device const block_q8_0 *)(src0_up + offset0); + } + + float sumg[NR0] = { 0.f }; + float sumu[NR0] = { 0.f }; + const short ix = tiisg / (NW / NQ); + const short il = tiisg % (NW / NQ); + const int ib0 = sgitg * NQ + ix; + float yl[NQ]; + device const float *yb = y + ib0 * QK8_0 + il * NQ; + + for (int ib = ib0; ib < nb; ib += NSG * NQ) { + FOR_UNROLL (short i = 0; i < NQ; ++i) yl[i] = yb[i]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + device const int8_t *qg = ag[row][ib].qs + il * NQ; + device const int8_t *qu = au[row][ib].qs + il * NQ; + float sg = 0.f; + float su = 0.f; + FOR_UNROLL (short i = 0; i < NQ; ++i) { + sg += qg[i] * yl[i]; + su += qu[i] * yl[i]; + } + sumg[row] += sg * ag[row][ib].d; + sumu[row] += su * au[row][ib].d; + } + yb += NSG * NQ * QK8_0; + } + + threadgroup float *shmem_f32 = (threadgroup float *)shmem; + threadgroup float *sh_gate[NR0]; + threadgroup float *sh_up[NR0]; + FOR_UNROLL (short row = 0; row < NR0; ++row) { + sh_gate[row] = shmem_f32 + NW * row; + sh_up[row] = shmem_f32 + NW * (NR0 + row); + if (sgitg == 0) { + sh_gate[row][tiisg] = 0.0f; + sh_up[row][tiisg] = 0.0f; + } + sumg[row] = simd_sum(sumg[row]); + sumu[row] = simd_sum(sumu[row]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + FOR_UNROLL (short row = 0; row < NR0; ++row) { + if (tiisg == 0) { + sh_gate[row][sgitg] = sumg[row]; + sh_up[row][sgitg] = sumu[row]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device float *mid_f32 = (device float *)dst_mid; + FOR_UNROLL (short row = 0; row < NR0 && r0 + row < args.ne01; ++row) { + const float gate = simd_sum(sh_gate[row][tiisg]); + const float up = simd_sum(sh_up[row][tiisg]); + if (tiisg == 0 && sgitg == 0) { + float g = dsv41_bf16(gate); + float u = dsv41_bf16(up); + if (clamp_value > 1.0e-6f) { + g = min(g, clamp_value); + u = clamp(u, -clamp_value, clamp_value); + } + const float silu = g / (1.0f + exp(-g)); + mid_f32[r0 + row] = dsv41_bf16(silu * u); + } + } +} + diff --git a/metal/dsv4_hc.metal b/metal/dsv4_hc.metal index c161a03a8b..5b22130ae9 100644 --- a/metal/dsv4_hc.metal +++ b/metal/dsv4_hc.metal @@ -1123,8 +1123,9 @@ struct ds4_metal_args_hc_norm_mix { // slice, preserving every simd_sum tree), and the matvec keeps the original // per-row accumulation order with y = x*scale computed on the fly, which // rounds identically to the materialized normalized row. The host wrapper -// gates this to n == 16384 && out_dim == 24, where the virtual-thread count -// is exactly 1024 and the mv tail loop is empty. +// gates this to n in {16384, 20480} && out_dim == 24, where the virtual-thread +// count is exactly 1024 and the mv tail loop is empty (n is a multiple of the +// NSG*NF*NB = 4096-value stride). kernel void kernel_dsv4_hc_rms_norm_mix_f16( constant ds4_metal_args_hc_norm_mix & args, device const char * x, diff --git a/tests/test_deepseek41_metal.c b/tests/test_deepseek41_metal.c index 055a70b677..69ac34c424 100644 --- a/tests/test_deepseek41_metal.c +++ b/tests/test_deepseek41_metal.c @@ -316,6 +316,338 @@ static int check_hc_scaled(void) { return 1; } +/* The fused decode HC glue (norm+mix, split+collapse+BF16+norm+BF16, + * expand+BF16) must be byte-identical to the standalone dispatch sequence + * the V4.1 graph ran before it, at V4.1's shape (4 x 5120, 24 mixes). */ +static int check_hc_fuse(void) { + enum { D = 5120, HC = 4, N = HC * D, OUT = 24, ITERS = 20 }; + const float hc_eps = 1e-6f, rms_eps = 1e-20f; + const size_t fn_bytes = (size_t)N * OUT * sizeof(_Float16); + const size_t scale_off = fn_bytes, base_off = scale_off + 64, normw_off = base_off + 128; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (normw_off + D * sizeof(float) + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + _Float16 *fn = model; + float *scale = (float *)((char *)model + scale_off); + float *base = (float *)((char *)model + base_off); + float *normw = (float *)((char *)model + normw_off); + for (size_t i = 0; i < (size_t)N * OUT; i++) fn[i] = (_Float16)(random_value() / 64); + for (int i = 0; i < 3; i++) scale[i] = 0.5f + random_value() / 8; + for (int i = 0; i < OUT; i++) base[i] = random_value() / 4; + for (int i = 0; i < D; i++) normw[i] = 1.0f + random_value() / 8; + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *residual = upload(NULL, N * 4), *block = upload(NULL, D * 4), *pre = upload(NULL, 16); + ds4_gpu_tensor *flat = upload(NULL, N * 4); + ds4_gpu_tensor *mix[2] = {upload(NULL, OUT * 4), upload(NULL, OUT * 4)}; + ds4_gpu_tensor *split[2] = {upload(NULL, OUT * 4), upload(NULL, OUT * 4)}; + ds4_gpu_tensor *x[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *norm[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *out[2] = {upload(NULL, N * 4), upload(NULL, N * 4)}; + ds4_gpu_tensor *pre_next[2] = {upload(NULL, 16), upload(NULL, 16)}; + CHECK(residual && block && pre && flat && mix[0] && mix[1] && split[0] && split[1] && + x[0] && x[1] && norm[0] && norm[1] && out[0] && out[1] && pre_next[0] && pre_next[1]); + float *res = ds4_gpu_tensor_contents(residual), *blk = ds4_gpu_tensor_contents(block); + float *p = ds4_gpu_tensor_contents(pre); + double elapsed[2] = {0}; + for (int round = 0; round < 6; round++) { + for (int i = 0; i < N; i++) res[i] = bf16(random_value() * (round == 5 ? 0x1p-14f : 1.0f)); + for (int i = 0; i < D; i++) blk[i] = bf16(random_value()); + for (int i = 0; i < HC; i++) p[i] = 0.5f + random_value() / 8 + hc_eps; + for (int mode = 0; mode < 2; mode++) { + const double begin = monotonic_seconds(); + CHECK(ds4_gpu_begin_commands()); + if (mode == 0) { + CHECK(ds4_gpu_rms_norm_plain_tensor(flat, residual, N, rms_eps)); + CHECK(ds4_gpu_matmul_f16_tensor(mix[0], model, mapped, 0, N, OUT, flat, 1)); + CHECK(ds4_gpu_hc_split_sinkhorn_tensor(split[0], mix[0], model, mapped, + scale_off, base_off, HC, ITERS, hc_eps)); + CHECK(ds4_gpu_hc_weighted_sum_tensor(x[0], residual, pre, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(x[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(norm[0], x[0], model, mapped, normw_off, D, rms_eps)); + CHECK(ds4_gpu_dsv41_quantize(norm[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_hc_expand_split_tensor(out[0], block, residual, split[0], D, HC)); + CHECK(ds4_gpu_dsv41_quantize(out[0], N, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_tensor_copy(pre_next[0], 0, split[0], 0, 16)); + } else { + CHECK(ds4_gpu_hc_rms_norm_mix_f16_tensor(mix[1], residual, model, mapped, 0, N, OUT, rms_eps)); + CHECK(ds4_gpu_dsv41_hc_collapse_norm(split[1], x[1], norm[1], mix[1], pre, residual, + model, mapped, scale_off, base_off, normw_off, D, HC, ITERS, hc_eps, rms_eps)); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], block, residual, split[1], pre_next[1], NULL, NULL, D)); + } + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (round) elapsed[mode] += (monotonic_seconds() - begin) * 1000.0 / 5; + } + CHECK(!memcmp(ds4_gpu_tensor_contents(mix[0]), ds4_gpu_tensor_contents(mix[1]), OUT * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(split[0]), ds4_gpu_tensor_contents(split[1]), OUT * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(x[0]), ds4_gpu_tensor_contents(x[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(norm[0]), ds4_gpu_tensor_contents(norm[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(pre_next[0]), ds4_gpu_tensor_contents(pre_next[1]), 16)); + const float *s = ds4_gpu_tensor_contents(split[1]); + for (int i = 0; i < OUT; i++) CHECK(isfinite(s[i])); + } + fprintf(stderr, "HC fuse: 10 dispatches %.3f ms -> 3 dispatches %.3f ms, outputs byte-identical\n", + elapsed[0], elapsed[1]); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(mix[i]); ds4_gpu_tensor_free(split[i]); ds4_gpu_tensor_free(x[i]); + ds4_gpu_tensor_free(norm[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(pre_next[i]); + } + ds4_gpu_tensor_free(residual); ds4_gpu_tensor_free(block); ds4_gpu_tensor_free(pre); ds4_gpu_tensor_free(flat); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + +static void fill_q8_0(uint8_t *dst, size_t rows, size_t k) { + for (size_t r = 0; r < rows; r++) for (size_t b = 0; b < k / 32; b++) { + uint8_t *block = dst + (r * (k / 32) + b) * 34; + const _Float16 d = (_Float16)(random_value() / 512 + 0.01f); + memcpy(block, &d, 2); + for (int i = 0; i < 32; i++) block[2 + i] = (uint8_t)(int8_t)(random_value() * 30); + } +} + +/* The fused decode MoE glue (router + select, shared gate/up/SwiGLU, shared + * down + sum + HC expand) must be byte-identical to the 22-dispatch + * standalone sequence at V4.1's shape (5120 x 384 F32 router, 6 of 384, + * scale 1.5, Q8_0 shared expert of 2304). */ +static int check_moe_fuse(void) { + enum { D = 5120, HC = 4, N = HC * D, E = 384, K = 6, FF = 2304 }; + const float scale = 1.5f, clamp = 10.0f; + const size_t router_bytes = (size_t)D * E * 4, bias_bytes = E * 4; + const size_t gate_bytes = (size_t)FF * (D / 32) * 34, down_bytes = (size_t)D * (FF / 32) * 34; + const size_t router_off = 0, bias_off = router_bytes, gate_off = bias_off + bias_bytes; + const size_t up_off = gate_off + gate_bytes, down_off = up_off + gate_bytes; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (down_off + down_bytes + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + float *router = model, *bias = (float *)((char *)model + bias_off); + for (size_t i = 0; i < (size_t)D * E; i++) router[i] = random_value() / 64; + /* Repeated bias values: with a zero input every score ties within its + * group and the canonical order (ascending index among equals) is what + * the standalone argsort produces. */ + for (int i = 0; i < E; i++) bias[i] = (float)(i % 5) / 10; + fill_q8_0((uint8_t *)model + gate_off, FF, D); + fill_q8_0((uint8_t *)model + up_off, FF, D); + fill_q8_0((uint8_t *)model + down_off, D, FF); + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *norm = upload(NULL, D * 4), *routed = upload(NULL, D * 4); + ds4_gpu_tensor *residual = upload(NULL, N * 4), *split = upload(NULL, 24 * 4); + ds4_gpu_tensor *logits[2] = {upload(NULL, E * 4), upload(NULL, E * 4)}; + ds4_gpu_tensor *probs[2] = {upload(NULL, E * 4), upload(NULL, E * 4)}; + ds4_gpu_tensor *selected[2] = {upload(NULL, K * 4), upload(NULL, K * 4)}; + ds4_gpu_tensor *weights[2] = {upload(NULL, K * 4), upload(NULL, K * 4)}; + ds4_gpu_tensor *gate = upload(NULL, FF * 4), *up = upload(NULL, FF * 4); + ds4_gpu_tensor *mid[2] = {upload(NULL, FF * 4), upload(NULL, FF * 4)}; + ds4_gpu_tensor *shared[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *block[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *out[2] = {upload(NULL, N * 4), upload(NULL, N * 4)}; + ds4_gpu_tensor *pre[2] = {upload(NULL, 16), upload(NULL, 16)}; + CHECK(norm && routed && residual && split && gate && up); + for (int i = 0; i < 2; i++) + CHECK(logits[i] && probs[i] && selected[i] && weights[i] && mid[i] && shared[i] && block[i] && out[i] && pre[i]); + float *x = ds4_gpu_tensor_contents(norm), *r = ds4_gpu_tensor_contents(routed); + float *res = ds4_gpu_tensor_contents(residual), *s = ds4_gpu_tensor_contents(split); + double elapsed[2] = {0}; + for (int round = 0; round < 8; round++) { + const int ties = round == 3 || round == 7; + for (int i = 0; i < D; i++) x[i] = ties ? 0.0f : bf16(random_value()); + for (int i = 0; i < D; i++) r[i] = bf16(random_value()); + for (int i = 0; i < N; i++) res[i] = bf16(random_value()); + for (int i = 0; i < 24; i++) s[i] = i < 4 ? 0.5f + random_value() / 8 : i < 8 ? 1.0f + random_value() / 4 : 0.25f + random_value() / 16; + for (int mode = 0; mode < 2; mode++) { + const double begin = monotonic_seconds(); + CHECK(ds4_gpu_begin_commands()); + if (mode == 0) { + CHECK(ds4_gpu_matmul_f32_tensor(logits[0], model, mapped, router_off, D, E, norm, 1)); + CHECK(ds4_gpu_router_select_tensor(selected[0], weights[0], probs[0], model, mapped, + bias_off, 0, 0, 0, E, K, scale, 0, 0, true, false, logits[0])); + CHECK(ds4_gpu_matmul_q8_0_tensor(gate, model, mapped, gate_off, D, FF, norm, 1)); + CHECK(ds4_gpu_dsv41_quantize(gate, FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_matmul_q8_0_tensor(up, model, mapped, up_off, D, FF, norm, 1)); + CHECK(ds4_gpu_dsv41_quantize(up, FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_swiglu_tensor(mid[0], gate, up, FF, clamp, 1.0f)); + CHECK(ds4_gpu_dsv41_quantize(mid[0], FF, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_matmul_q8_0_tensor(shared[0], model, mapped, down_off, FF, D, mid[0], 1)); + CHECK(ds4_gpu_dsv41_quantize(shared[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_add_tensor(block[0], routed, shared[0], D)); + CHECK(ds4_gpu_dsv41_quantize(block[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_hc_expand_split_tensor(out[0], block[0], residual, split, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(out[0], N, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_tensor_copy(pre[0], 0, split, 0, 16)); + } else { + CHECK(ds4_gpu_dsv41_router_select(selected[1], weights[1], probs[1], logits[1], norm, + model, mapped, router_off, bias_off, true, D, E, K, scale)); + CHECK(ds4_gpu_dsv41_shared_gate_up_swiglu(mid[1], norm, model, mapped, gate_off, up_off, D, FF, clamp)); + CHECK(ds4_gpu_dsv41_matvec_bf16(shared[1], model, mapped, down_off, DS4_V41_WEIGHT_Q8_0, FF, D, mid[1]) == 1); + CHECK(ds4_gpu_dsv41_hc_expand4(out[1], routed, residual, split, pre[1], shared[1], block[1], D)); + } + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (round) elapsed[mode] += (monotonic_seconds() - begin) * 1000.0 / 7; + } + const int32_t *sel = ds4_gpu_tensor_contents(selected[1]); + for (int i = 0; i < K; i++) CHECK(sel[i] >= 0 && sel[i] < E); + /* Ties: the fused select orders equal scores by ascending index (the + * canonical argsort order); an argsort without that tie-break agrees + * on the set but not necessarily on the order. */ + if (ties) for (int i = 0; i < K; i++) CHECK(sel[i] == 4 + 5 * i); + CHECK(!memcmp(ds4_gpu_tensor_contents(logits[0]), ds4_gpu_tensor_contents(logits[1]), E * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(probs[0]), ds4_gpu_tensor_contents(probs[1]), E * 4)); + if (!ties) { + CHECK(!memcmp(ds4_gpu_tensor_contents(selected[0]), sel, K * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(weights[0]), ds4_gpu_tensor_contents(weights[1]), K * 4)); + } + CHECK(!memcmp(ds4_gpu_tensor_contents(mid[0]), ds4_gpu_tensor_contents(mid[1]), FF * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(shared[0]), ds4_gpu_tensor_contents(shared[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(block[0]), ds4_gpu_tensor_contents(block[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), N * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(pre[0]), ds4_gpu_tensor_contents(pre[1]), 16)); + } + fprintf(stderr, "MoE fuse: 15 dispatches %.3f ms -> 4-5 dispatches %.3f ms, outputs byte-identical\n", + elapsed[0], elapsed[1]); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(logits[i]); ds4_gpu_tensor_free(probs[i]); ds4_gpu_tensor_free(selected[i]); + ds4_gpu_tensor_free(weights[i]); ds4_gpu_tensor_free(mid[i]); ds4_gpu_tensor_free(shared[i]); + ds4_gpu_tensor_free(block[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(pre[i]); + } + ds4_gpu_tensor_free(norm); ds4_gpu_tensor_free(routed); ds4_gpu_tensor_free(residual); + ds4_gpu_tensor_free(split); ds4_gpu_tensor_free(gate); ds4_gpu_tensor_free(up); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + +/* The fused decode attention glue: BF16-on-store matvecs (Q8_0 and F16 at + * the graph's shapes, including the nr0=4 F16 dispatch), the q/kv norms + * with the KV RoPE/FP8/window tail, the heads' BF16 + inverse RoPE, and the + * logits' collapse + norm; each byte-identical to its standalone sequence. */ +static int check_attn_fuse(void) { + enum { D = 5120, LQ = 1280, HD = 512, HEADS = 64, QD = HEADS * HD, HC = 4, IDX = 1024 }; + const float eps = 1e-20f; + struct { uint32_t type, in, out; size_t bytes, off; } w[] = { + {DS4_V41_WEIGHT_Q8_0, D, LQ, 0, 0}, /* attn_q_a */ + {DS4_V41_WEIGHT_Q8_0, D, HD, 0, 0}, /* attn_kv */ + {DS4_V41_WEIGHT_Q8_0, LQ, QD, 0, 0}, /* attn_q_b */ + {DS4_V41_WEIGHT_F16, D, HD, 0, 0}, /* attn_compressor_kv: the nr0 = 4 dispatch */ + {DS4_V41_WEIGHT_F16, HD, 128, 0, 0}, /* indexer_attn_k */ + {DS4_V41_WEIGHT_F16, LQ, IDX, 0, 0}, /* indexer_attn_q_b */ + }; + const int nw = sizeof(w) / sizeof(*w); + size_t total = 0; + for (int i = 0; i < nw; i++) { + w[i].bytes = w[i].type == DS4_V41_WEIGHT_Q8_0 ? (size_t)w[i].out * (w[i].in / 32) * 34 : (size_t)w[i].out * w[i].in * 2; + w[i].off = total; + total += (w[i].bytes + 63) / 64 * 64; + } + const size_t qnorm_off = total, kvnorm_off = qnorm_off + LQ * 4, onorm_off = kvnorm_off + HD * 4; + const size_t page = (size_t)getpagesize(); + const size_t mapped = (onorm_off + D * 4 + page - 1) / page * page; + void *model = NULL; + CHECK(!posix_memalign(&model, page, mapped)); + for (int i = 0; i < nw; i++) { + if (w[i].type == DS4_V41_WEIGHT_Q8_0) fill_q8_0((uint8_t *)model + w[i].off, w[i].out, w[i].in); + else { _Float16 *h = (_Float16 *)((char *)model + w[i].off); for (size_t k = 0; k < (size_t)w[i].out * w[i].in; k++) h[k] = (_Float16)(random_value() / 64); } + } + float *qnorm = (float *)((char *)model + qnorm_off), *kvnorm = (float *)((char *)model + kvnorm_off); + float *onorm = (float *)((char *)model + onorm_off); + for (int i = 0; i < LQ; i++) qnorm[i] = 1.0f + random_value() / 8; + for (int i = 0; i < HD; i++) kvnorm[i] = 1.0f + random_value() / 8; + for (int i = 0; i < D; i++) onorm[i] = 1.0f + random_value() / 8; + CHECK(ds4_gpu_set_model_map(model, mapped)); + + ds4_gpu_tensor *xin[2] = {upload(NULL, D * 4), upload(NULL, LQ * 4)}; /* D-wide and LQ/HD-wide inputs */ + ds4_gpu_tensor *out[2] = {upload(NULL, QD * 4), upload(NULL, QD * 4)}; + ds4_gpu_tensor *qr[2] = {upload(NULL, LQ * 4), upload(NULL, LQ * 4)}; + ds4_gpu_tensor *kv[2] = {upload(NULL, HD * 4), upload(NULL, HD * 4)}; + ds4_gpu_tensor *window[2] = {upload(NULL, 128 * HD * 4), upload(NULL, 128 * HD * 4)}; + ds4_gpu_tensor *heads[2] = {upload(NULL, QD * 4), upload(NULL, QD * 4)}; + ds4_gpu_tensor *residual = upload(NULL, HC * D * 4), *pre = upload(NULL, 16); + ds4_gpu_tensor *x[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + ds4_gpu_tensor *norm[2] = {upload(NULL, D * 4), upload(NULL, D * 4)}; + CHECK(xin[0] && xin[1] && residual && pre); + for (int i = 0; i < 2; i++) CHECK(out[i] && qr[i] && kv[i] && window[i] && heads[i] && x[i] && norm[i]); + float *xd = ds4_gpu_tensor_contents(xin[0]), *xl = ds4_gpu_tensor_contents(xin[1]); + for (int round = 0; round < 4; round++) { + for (int i = 0; i < D; i++) xd[i] = bf16(random_value()); + for (int i = 0; i < LQ; i++) xl[i] = bf16(random_value()); + for (int i = 0; i < nw; i++) { + const ds4_gpu_tensor *in = w[i].in == D ? xin[0] : xin[1]; + CHECK(ds4_gpu_begin_commands()); + if (w[i].type == DS4_V41_WEIGHT_Q8_0) CHECK(ds4_gpu_matmul_q8_0_tensor(out[0], model, mapped, w[i].off, w[i].in, w[i].out, in, 1)); + else CHECK(ds4_gpu_matmul_f16_tensor(out[0], model, mapped, w[i].off, w[i].in, w[i].out, in, 1)); + CHECK(ds4_gpu_dsv41_quantize(out[0], w[i].out, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_matvec_bf16(out[1], model, mapped, w[i].off, w[i].type, w[i].in, w[i].out, in) == 1); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + if (memcmp(ds4_gpu_tensor_contents(out[0]), ds4_gpu_tensor_contents(out[1]), (size_t)w[i].out * 4)) { + fprintf(stderr, "matvec bf16 mismatch: weight %d (type %u, %u -> %u)\n", i, w[i].type, w[i].in, w[i].out); + CHECK(0); + } + } + /* q/kv norms + KV tail, both frequency tables and window slots. */ + const uint32_t pos = 777u + 1000u * round; + const int compressed = round & 1; + float *q0 = ds4_gpu_tensor_contents(qr[0]), *q1 = ds4_gpu_tensor_contents(qr[1]); + float *k0 = ds4_gpu_tensor_contents(kv[0]), *k1 = ds4_gpu_tensor_contents(kv[1]); + for (int i = 0; i < LQ; i++) q0[i] = q1[i] = bf16(random_value()); + for (int i = 0; i < HD; i++) k0[i] = k1[i] = bf16(random_value()); + const uint64_t slot = (uint64_t)(pos % 128u) * HD * 4u; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_rms_norm_weight_tensor(qr[0], qr[0], model, mapped, qnorm_off, LQ, eps)); + CHECK(ds4_gpu_dsv41_quantize(qr[0], LQ, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(kv[0], kv[0], model, mapped, kvnorm_off, HD, eps)); + CHECK(ds4_gpu_dsv41_quantize(kv[0], HD, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_rope(kv[0], HD, 1, 1, pos, compressed, false)); + CHECK(ds4_gpu_dsv41_quantize(kv[0], HD, 1, DS4_V41_FP8_E8M0)); + CHECK(ds4_gpu_tensor_copy(window[0], slot, kv[0], 0, HD * 4)); + CHECK(ds4_gpu_dsv41_qkv_norm_kv_tail(qr[1], kv[1], window[1], slot, model, mapped, qnorm_off, kvnorm_off, LQ, HD, eps, pos, compressed)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(q0, q1, LQ * 4)); + CHECK(!memcmp(k0, k1, HD * 4)); + CHECK(!memcmp((char *)ds4_gpu_tensor_contents(window[0]) + slot, (char *)ds4_gpu_tensor_contents(window[1]) + slot, HD * 4)); + /* heads: BF16 over every value, inverse RoPE on each head's last 64 */ + float *h0 = ds4_gpu_tensor_contents(heads[0]), *h1 = ds4_gpu_tensor_contents(heads[1]); + for (int i = 0; i < QD; i++) h0[i] = h1[i] = random_value() * 3.0f; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_dsv41_quantize(heads[0], QD, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_rope(heads[0], HD, HEADS, 1, pos, compressed, true)); + CHECK(ds4_gpu_dsv41_bf16_rope(heads[1], HD, HEADS, 1, pos, compressed, true)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(h0, h1, QD * 4)); + /* logits: collapse with the carried pre, BF16, output norm, BF16 */ + float *res = ds4_gpu_tensor_contents(residual), *p = ds4_gpu_tensor_contents(pre); + for (int i = 0; i < HC * D; i++) res[i] = bf16(random_value()); + for (int i = 0; i < HC; i++) p[i] = 0.5f + random_value() / 8; + CHECK(ds4_gpu_begin_commands()); + CHECK(ds4_gpu_hc_weighted_sum_tensor(x[0], residual, pre, D, HC)); + CHECK(ds4_gpu_dsv41_quantize(x[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_rms_norm_weight_tensor(norm[0], x[0], model, mapped, onorm_off, D, eps)); + CHECK(ds4_gpu_dsv41_quantize(norm[0], D, 1, DS4_V41_BF16)); + CHECK(ds4_gpu_dsv41_hc_collapse_norm(NULL, x[1], norm[1], NULL, pre, residual, model, mapped, 0, 0, onorm_off, D, HC, 0, 0.0f, eps)); + CHECK(ds4_gpu_end_commands()); + CHECK(ds4_gpu_synchronize()); + CHECK(!memcmp(ds4_gpu_tensor_contents(x[0]), ds4_gpu_tensor_contents(x[1]), D * 4)); + CHECK(!memcmp(ds4_gpu_tensor_contents(norm[0]), ds4_gpu_tensor_contents(norm[1]), D * 4)); + } + fprintf(stderr, "attention fuse: matvec bf16 (6 shapes), qkv norm + KV tail, bf16 + rope, logits collapse: byte-identical\n"); + for (int i = 0; i < 2; i++) { + ds4_gpu_tensor_free(xin[i]); ds4_gpu_tensor_free(out[i]); ds4_gpu_tensor_free(qr[i]); ds4_gpu_tensor_free(kv[i]); + ds4_gpu_tensor_free(window[i]); ds4_gpu_tensor_free(heads[i]); ds4_gpu_tensor_free(x[i]); ds4_gpu_tensor_free(norm[i]); + } + ds4_gpu_tensor_free(residual); ds4_gpu_tensor_free(pre); + ds4_gpu_cleanup(); free(model); + CHECK(ds4_gpu_init()); + return 1; +} + #endif static int check_engram(void) { @@ -1279,6 +1611,21 @@ int main(int argc, char **argv) { ds4_gpu_cleanup(); return ok ? 0 : 1; } + if (argc == 2 && !strcmp(argv[1], "--hc-fuse")) { + const int ok = ds4_gpu_init() && check_hc_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } + if (argc == 2 && !strcmp(argv[1], "--moe-fuse")) { + const int ok = ds4_gpu_init() && check_moe_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } + if (argc == 2 && !strcmp(argv[1], "--attn-fuse")) { + const int ok = ds4_gpu_init() && check_attn_fuse(); + ds4_gpu_cleanup(); + return ok ? 0 : 1; + } #endif if (argc == 2 && !strcmp(argv[1], "--bf16-linear")) { const int ok = ds4_gpu_init() && check_bf16_linear(); @@ -1315,6 +1662,9 @@ int main(int argc, char **argv) { check_candidates() && check_sparse_gather() && check_indexer_batch() && check_embedding() && check_index_projection() && check_general_topk() && check_causal_topk() && check_compact_carry() && check_attention_output(false) && check_tp_attention(); +#ifdef __APPLE__ + ok = ok && check_hc_fuse() && check_moe_fuse() && check_attn_fuse(); +#endif ds4_gpu_cleanup(); return ok ? 0 : 1; }