v41: fuse the single-box decode glue (HC, MoE, attention): +20% on M3 Ultra on top of #1041, bit-exact - #1042
Open
adriangalilea wants to merge 7 commits into
Open
adriangalilea wants to merge 7 commits into
adriangalilea wants to merge 7 commits into
Conversation
…yer) DeepSeek's production decode runs each half-layer's HC work in one kernel (Mega-mHC); ds4's V4.1 graph mirrored the reference op by op. Three fusions, all byte-identical to the standalone sequence they replace: - kernel_dsv41_hc_collapse_norm4: split/sinkhorn + pre-weighted collapse + BF16 + weighted RMSNorm + BF16 (5 -> 1), keeping the standalone kernels' thread mapping, simd trees and 1/sqrt scale. - kernel_dsv41_hc_expand4_bf16: post/comb expand + BF16, optional carry of split[0..3] into pre (2-3 -> 1). - kernel_dsv4_hc_rms_norm_mix_f16 now accepts V4.1's 4x5120 HC row: the norm runs 1024 threads there too and 20480 is a whole multiple of the matvec's 4096-value stride, so the exact replica and empty tail hold. Single box only (TP keeps the unfused path). Rollback: DS4_METAL_DISABLE_V41_HC_FUSE. tests/test_deepseek41_metal --hc-fuse compares every output of both paths byte for byte (6 rounds, M4 Pro).
DeepSeek's production decode runs the router as one "Mega-Gate" kernel; ds4's V4.1 graph ran it as the F32 logits matvec plus ten generic select dispatches (softplus, sqrt, bias add, argsort, gather, sum, clamp, div, scale), and the shared expert as three Q8_0 matvecs with three BF16 passes, SwiGLU, another BF16 pass, then the routed + shared sum and its rounding. Three fusions, byte-identical to those sequences: - kernel_dsv41_router_select: kernel_mul_mv_f32_f32_4's matvec (nsg=8, nr0=2), then the last-arriving threadgroup does sqrt(softplus), the bias, the canonical (score desc, idx asc) top-k of the argsort path and the normalized weights with the six-lane sum_rows reduction, the clamp, the row division and the scale as separately rounded ops. 11 -> 1. - kernel_dsv41_shared_gate_up_swiglu_q8_0: both Q8_0 matvecs at the dispatch's nsg, gate/up rounded to BF16 before SwiGLU and the product after it. 6 -> 1. - kernel_dsv41_shared_down_hc_expand4_q8_0: the down matvec, shared = bf16(.), block = bf16(routed + shared), the post/comb expand with its rounding and the pre carry; runs in the layer's FFN tail. 6 -> 1. Single box only. Rollback: DS4_METAL_DISABLE_V41_MOE_FUSE (the HC fusion switch also disables it). tests/test_deepseek41_metal --moe-fuse compares every output of both paths byte for byte, including a zero-input round where all scores tie within bias groups (canonical order check).
…oup elsewhere The last-arriving-threadgroup select read stale logits on the M3 Ultra (nondeterministic probs, kernel test caught it on the box); V4 gates the same pattern to M5. Pre-M5 runs the standalone F32 matvec dispatch and one selecting threadgroup: 11 -> 2 instead of 11 -> 1, byte-identical.
…perts Attention (single box), byte-identical to the standalone sequences (tests/test_deepseek41_metal --attn-fuse): - Q8_0 and F16 single-row matvecs round to BF16 on the store (kernel_dsv41_mul_mv_q8_0_f32_bf16, kernel_dsv41_mul_mv_f16_f32_4_bf16: the standalone walks and reduction trees, nsg from the dispatch, the nr0=4 F16 shape included); ds41_matmul(…, round) uses them, removing the rounding dispatch after q_a, q_b, kv, output_b, the shared down, the compressor and indexer projections. - kernel_dsv41_qkv_norm_kv_tail: the q LoRA and KV weighted norms with their roundings, then the KV RoPE, FP8 (E8M0) block quantization and the raw-window store in one dispatch (6 -> 1). 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 reduction tree unchanged. - kernel_dsv41_bf16_rope: the rounding pass over all heads and their inverse RoPE (2 -> 1). - The logits' collapse + output norm reuse kernel_dsv41_hc_collapse_norm4 without a split (4 -> 1). MoE: the shared expert's down projection now runs before the routed experts (rounded on its store) and the routed + shared sum with its rounding is folded into the FFN tail's expand; a fused tail that read shared_mid after the routed kernel saw clobbered values on the M3 Ultra (greedy output differed) while the same kernel was exact in isolation. That late kernel is gone. Rollbacks: DS4_METAL_DISABLE_V41_ATTN_FUSE, DS4_METAL_DISABLE_V41_MATVEC_BF16.
The fused select orders equal scores by ascending index; an argsort without that tie-break agrees on the selected set but not necessarily on its order, so the tie rounds check the fused order alone.
adriangalilea
changed the title
This was referenced Sep 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
The V4.1 single-token decode graph mirrors the reference
model.pyop by op: per layer, 19 dispatches of hyper-connection glue (norm, mix matvec, split/sinkhorn, collapse, BF16 rounding, weighted norm, BF16, expand, BF16, …), 22 of MoE glue (router matvec + ten generic select dispatches; shared expert gate/up/down with three BF16 passes, SwiGLU and its rounding; the routed + shared sum and its rounding) and ~14 of attention glue (projections each followed by a BF16 rounding dispatch, two norms, RoPE, FP8 quantization, window copy, the heads' rounding + inverse RoPE). DeepSeek's production decode runs these as a handful of fused kernels per layer (tech report §3.2: Mega-mHC, Mega-Gate, fused RoPE-attention-cast).This PR fuses ds4's sequences while keeping every reduction tree and every rounding point of the standalone kernels, so the output is byte-identical. Six commits, each its own step:
kernel_dsv41_hc_collapse_norm4(split/sinkhorn + collapse with the previous sublayer'spre+ BF16 + weighted RMSNorm + BF16),kernel_dsv41_hc_expand4_bf16(post/comb expand + BF16, optionalprecarry, optional folded routed + shared sum), and the existing exact norm+mix fusionkernel_dsv4_hc_rms_norm_mix_f16accepting V4.1's 4×5120 row (1024 norm threads and an empty matvec tail there too).kernel_dsv41_router_selectreproduces sqrt(softplus), the bias, the canonical (score desc, idx asc) top-6 and the normalized weights with the six-lanesum_rowsreduction, the clamp, the row division and the scale as separately rounded ops. The single-dispatch form (last-arriving threadgroup selects, DeepSeek's Mega-Gate shape) is enabled on M5 only, like V4'skernel_dsv4_router_project_select_fused: on the M3 Ultra the other groups' logits were not reliably visible to the last one (nondeterministic output, caught by the on-box kernel test). Elsewhere the standalone F32 matvec runs first and one threadgroup selects.kernel_dsv41_shared_gate_up_swiglu_q8_0(both Q8_0 matvecs at the dispatch's nsg, gate/up rounded before SwiGLU, the product after), the down projection rounded on its store, and the routed + shared sum with its rounding folded into the FFN tail's expand. The down projection deliberately runs before the routed experts: a fused tail that readshared_midafter the routed kernel saw clobbered values on the M3 Ultra (greedy output differed) while the same kernel was byte-exact in isolation.kernel_dsv41_mul_mv_q8_0_f32_bf16/kernel_dsv41_mul_mv_f16_f32_4_bf16(the standalone walks and reduction trees, nsg from the dispatch, the nr0=4 F16 shape included) round on the store, removing the rounding dispatch after q_a, q_b, kv, output_b, the shared down, the compressor and indexer projections;kernel_dsv41_qkv_norm_kv_taildoes the q LoRA and KV weighted norms with their roundings, then the KV RoPE, FP8 (E8M0) block quantization and the raw-window store (6 → 1);kernel_dsv41_bf16_ropeis the rounding pass over all heads plus the inverse RoPE (2 → 1); the logits' collapse + output norm reuse the collapse kernel without a split (4 → 1).Single box only; the TP graph keeps the unfused sequences. Rollbacks:
DS4_METAL_DISABLE_V41_HC_FUSE(everything),_MOE_FUSE,_ROUTER_FUSE,_SHARED_FUSE,_ATTN_FUSE,DS4_METAL_DISABLE_V41_MATVEC_BF16,DS4_METAL_DISABLE_V41_ROUTER_SINGLE_DISPATCH.Tests
tests/test_deepseek41_metal --hc-fuse,--moe-fuse,--attn-fuse(also in the default suite on Apple): random data at V4.1's shapes, the standalone dispatch sequences through the publicds4_gpu_*entry points vs the fused kernels,memcmpon every output (mix, split, x, norm, expanded rows, carriedpre; logits, probs, selected, weights, mid, shared, block; six matvec shapes, q/kv rows, window row, heads, the logits' x/norm). Pass on M4 Pro and M3 Ultra. Ties in the router (a zero-input round) are checked against the canonical order the fused select defines.Numbers
Mac Studio M3 Ultra 512 GB,
DeepSeek-V4.1-Flash-Q4.gguf(resident),./ds4 -m … -p <2k of promessi_sposi> -n 256 --temp 0, on top of #1041 + #1035 (DS4_ENGRAM_PARALLEL_DECODE=1):+18-20 %, output byte-identical. GPU busy per token (stage timestamps, #828) 36.9 → 31.2 ms; pre_attn 2.81 → 2.13, post_attn 2.16 → 1.16, after_moe 0.65 → 0.24, moe_route 2.70 → 1.15, moe_shared + FFN tail 3.07 + 0.65 → 2.50 + 0.24, attn_kv 0.69 → 0.27, attn_project 4.31 → 3.77. Prefill unchanged (the batched path is untouched).
Independent of #1041; measured on top of it because that is what the box runs.