Gemma 4: D512 two-pass SDPA + GQA dedup kernels - #13
Open
davidtai wants to merge 6 commits into
Open
Conversation
Chunk M1 of the Gemma 4 mlxfast port (ledger §C.1). Ported from Layr-Labs/mlxfast-gemma4-26b-a4b-engine 944c572 (final vendored tree), which carried it from engine commits 796aa221 (Validate submission 6ce1e46e-fe7e-4996-be64-cfb9501fc8f5, softmax) and 57087ca2 (Accept submission 82c69b6c-37e9-4f38-9d20-6b122e7ceb57, sdpa_vector). Mechanism: `#pragma unroll` on the compile-time-trip-count loops — N_READS in both softmax kernels, qk_per_thread / v_per_thread / elem_per_thread in the SDPA vector, vector-2pass and vector-2pass-reduce kernels. No arithmetic, no accumulation order, no operand shapes change; these are loop-structure hints only, so every output stays bit-identical. Files: - mlx/backend/metal/kernels/softmax.h (+12) - mlx/backend/metal/kernels/sdpa_vector.h (+14) Already upstream in 0.32.2: nothing. Both hunks are new here. `softmax.h` is byte-identical between the engine's fork base (d5a2404) and this branch's tip, so it applied unchanged. `sdpa_vector.h` drifted 143 lines across 0.32.0 -> 0.32.2, but every pragma landed by three-way merge against d5a2404 with no conflict. Co-authored-by: fkiene <46886660+fkiene@users.noreply.github.com> Co-authored-by: jungjipdo <130676635+jungjipdo@users.noreply.github.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NS86aXV85ai4eLRw1ADEyX
… fused GEMM epilogues
Chunk M2 of the Gemma 4 mlxfast port (ledger §C.1).
Ported from Layr-Labs/mlxfast-gemma4-26b-a4b-engine 944c572 (final
vendored tree), which carried CAUSAL-CLOAD from engine commit 4f44957c
(Validate submission 5a454a6a-26b6-4af0-ab20-d256cfe328bd) and
NAX-SKIP-EMPTY-001 from f68023e0 (Validate submission
1e9f5531-235a-4920-b516-c08a9908d864).
Mechanism (CAUSAL-CLOAD): in the addmm epilogue, recognise the
composed-prefill causal-bias operand by its signature — bf16 accumulate,
!transpose_a && transpose_b, fdc == 1, ldc == N + 1, M <= N, all-zero C
batch strides — and synthesize its two constants per accumulator element
instead of loading them: widened bfloat16 lowest finite (0xFF7F) strictly
above the causal diagonal at N - M, widened bfloat16 negative zero on and
below it. A row stride of N + 1 cannot arise from a contiguous or
broadcast operand of the declared output width, so the signature is
unambiguous. The addend still enters through the same TransformAdd with
the same widening as the loaded operand it replaces, so every stored word
is bit-identical; every other addmm keeps the loaded-operand epilogue.
Adds the `kCausalBiasSynthEligible` constexpr gate (so complex64 never
instantiates the branch), the `c_bstride_zero` batch-stride check, and the
`gemm_epilogue_causal_synth` NAX tile helper.
Files:
- mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused.h (+57/-2)
- mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h (+72/-2)
Already upstream in MLX 0.32.2, so NOT re-applied:
- NAX-SKIP-EMPTY-001 in its entirety. This branch's tip already carries the
empty-output-extent elision unconditionally — `if constexpr (!kAlignedM ||
!kAlignedN) { if (!has_output) …}` in steel/gemm/gemm_nax.h and the
`(kAlignedM.value || sgp_sm > 0) && (kAlignedN.value || sgp_sn > 0)`
epilogue guard in steel_gemm_fused_nax.h. The engine's version of the
same optimisation wraps those guards in a `DARKBLOOM_GEMMA4_NAX_SKIP_EMPTY`
kill-switch macro; re-applying it would only make an already-live
optimisation conditional. steel/gemm/gemm_nax.h is therefore untouched by
this commit, and the macro definition and both guard rewrites were dropped
from steel_gemm_fused_nax.h — only the CAUSAL-CLOAD body was kept inside
upstream's guard.
Conflicts resolved (three-way against the engine's fork base d5a2404):
- steel/gemm/gemm_nax.h ×3 — kept this branch's unguarded skip-empty.
- steel_gemm_fused_nax.h ×1 — kept this branch's epilogue guard, took the
engine's causal-synth body inside it.
steel_gemm_fused.h merged with no conflict (0 lines of drift vs d5a2404).
Co-authored-by: Amal-David <11647194+Amal-David@users.noreply.github.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NS86aXV85ai4eLRw1ADEyX
…kernel Chunk M3 of the Gemma 4 mlxfast port (ledger §C.1). Ported from Layr-Labs/mlxfast-gemma4-26b-a4b-engine 944c572 (final vendored tree), aggregating engine commits 465ce5ce, 58d03e6a, 3b51325b and f68023e0 (Accept/Validate submissions 91c92a7e-fba5-4fe7-abdc-a0e61d3e2a87, 6bf82ab1-f553-41de-a1d6-4041e4a0a352, c8726cfe-26bb-4ef9-a089-099d5243a5c8 and 1e9f5531-235a-4920-b516-c08a9908d864). Mechanism: `affine_gather_qmm_rhs_nax` computes a full SM x SN tile per simdgroup and then discards, at `store_slice`, every row outside the current expert segment `[seg_lo, seg_hi)`. This hoists that band ahead of the K-loop and skips the discarded work instead of computing it: - `seg_empty` — the whole band is dead, so the simdgroup runs no A load and no MMA at all (band granularity); - `seg_partial` (aligned-M only) — 16-row fragment-row granularity: only fragment rows intersecting the band call `gather_rhs_load_frag_row` / `gather_rhs_mma_frag_row`, each running the stock path's exact op sequence for that row. Cooperative weight loads and every `threadgroup_barrier` stay unconditional, so barrier convergence is preserved; `offset`/`offset_next` are threadgroup uniform and `seg_*` simdgroup uniform, so no intra-simdgroup divergence is introduced. Gated by `kGatherRhsSegmentElide`; with it off only the stock path runs. Also brings the engine's `qmm_t_nax_tgp_impl` and `tile_matmad_nax` additions in the same file. Files: - mlx/backend/metal/kernels/quantized_nax.h (+271/-34) Already upstream in MLX 0.32.2, and therefore SUPERSEDED rather than re-applied: this branch's tip had independently added the band-granular half of the same elision as `sg_active`, computed from `m_lo_lim`/ `m_hi_lim` — expressions textually identical to the engine's `seg_lo`/ `seg_hi`. `seg_empty` is exactly `!sg_active`, and `seg_partial` is the finer tier upstream does not have, so the engine's form subsumes it. The now-unreferenced `m_lo_lim`/`m_hi_lim`/`sg_active` trio was removed. Conflicts resolved (three-way against the engine's fork base d5a2404), all five inside `affine_gather_qmm_rhs_nax`: - K-loop head ×2 and unaligned-K tail ×1 — took the engine's `seg_partial`/ `seg_empty` structure over this branch's `if (sg_active)`. - Btile load reformat ×2 — pure whitespace; took the engine's wrapping. - store block ×1 — took the engine's `if (!seg_empty)` + `seg_lo`/`seg_hi` spelling of this branch's `m_lo_lim`/`m_hi_lim` slice. NOT VALIDATED ON DEVICE. This port was produced build-only; the numerics of the fragment-row path have not been re-measured against this branch's kernels. Co-authored-by: Amal-David <11647194+Amal-David@users.noreply.github.com> Co-authored-by: i34-9 <313589706+i34-9@users.noreply.github.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NS86aXV85ai4eLRw1ADEyX
Chunk M4 of the Gemma 4 mlxfast port (ledger §C.1); M4a and M4b are
combined here — see "Deviation" below.
Ported from Layr-Labs/mlxfast-gemma4-26b-a4b-engine 944c572 (final
vendored tree), aggregating engine commits 61c33ed1, fd981eea, eb82740f,
3ab9bd49, 41ab2004, 46e41ab9, 2afa9b80, 1356b788, 15b14f54, cbbcb92d,
bdbb9947, 38269aa8, 6042598b, 712602f4, 81314800, cdddcf15, 60318626,
f6904ac8, c2284eeb and b0631d8e (Accept/Validate submissions on the
ranked Gemma 4 26B-A4B track).
Mechanism — new tiers, all bit-exact restatements of the stock `qdot` /
`qmv_impl` arithmetic with a different load or dispatch shape:
- `qdot_affine4_registered{,_word}`, `qdot_affine8_registered{,_word}`,
`qdot_affine4_pair{,_word}`, `qdot_affine8_pair`, `qdot_affine4_loaded`,
`qdot_affine4_loaded_pair`, `qdot_affine4_g64_word` — the bits==4 /
values_per_thread==8 (and byte-weight) arms of `qdot` against packed
weight words already in registers: same nibble masks, same two 4-term
sums, same accumulation order over i, same `scale * accum + sum * bias`
close. Only the load shape differs (one 4-byte load instead of two
2-byte loads).
- `qmv_affine4_g64_pair_impl`, `_triple_stream_impl`, `_quad_stream_impl`,
`qmv_affine8_g64_pair_impl`, `_quad_stream_impl`, and
`qmv_affine4_g64_singles_impl` — 1/2/3/4 same-expert assignments served
from ONE weight stream; each (output, input) pair keeps its own
accumulator and K-loop order, so every output element's add sequence
matches the incumbent per-arm kernel.
- `qmv_fast_crossrow_affine4_g64{,_wide,_m}`, `qmv_fast_singlerow_affine2_g64`
— cross-row tight-grid bodies for the batch-8 decode plane.
- `mma8_lane`/`mma8_lo`/`mma8_hi`/`mma8_runsum4` +
`gemma4_qmv_mma8_affine4_g64_impl` — fp32 `simdgroup_float8x8` body for
the M=8 decode cohort on 4-bit affine g64 (A = raw weight codes 8x8,
B = x-transpose 8x8, C zeroed per g64 group).
- `gather_qmv_gemma4_down_tile` + the `affine_gather_qmv` dispatch rewrite
— RUN-QUAD leader election over the flattened 64-assignment route table,
reading the EXPERT-PREFIX-BOUNDS-001 packed route word (bit 31 = format
flag, bits 0-7 expert, 8-13 run offset, 14-19 run length) with a
linear-scan fallback when the flag is clear, plus the y-tile-coarsened
arm for the K = 704 down plane. Both arms are compile-time flippable
(`gemma4_down_tile`) and bit-identical by construction.
- Two `qmv_impl` loop bounds change from `k < in_vec_size - block_size` to
`k <= …`, so an exactly block-aligned input runs its last full block on
the fast path instead of the `qdot_safe` tail; the tail's `remaining`
clamp already covers k == in_vec_size.
Files:
- mlx/backend/metal/kernels/quantized.h (+2281/-55)
Already upstream in MLX 0.32.2 and preserved unchanged by the three-way
merge: this branch's `qmv_wide` family (Layr-Labs/mlx-swift 606d28c
"expose qmv_wide to Swift runtime", 4 references) and the
`has_global_scale` template parameter added across the affine kernels
(10 references) both live in the same `qmv_affine*` region the engine's
tiers were written into. Neither was reverted; the engine's fork base
(d5a2404) predates both.
Conflict resolved (one, three-way against d5a2404): the declaration
immediately preceding `[[kernel]] void affine_gather_qmv` — the engine
inserted its `qdot_affine4_g64_word` + `qmv_affine4_g64_singles_impl` +
`gather_qmv_gemma4_down_tile` block there while this branch had widened
the following template to `template <typename T, int group_size, int bits,
bool has_global_scale = false>`. Kept this branch's widened signature and
inserted the engine's block ahead of it.
Deviation from the planned chunking: the ledger suggests splitting this
into M4a (the `qdot`/`_pair`/`_stream` primitives) and M4b (the
Gemma-4-specific tiers). The diff does not split at hunk boundaries — one
888-line hunk contains both `qmv_affine4_g64_pair_impl` and the
`mma8_*`/`gemma4_qmv_mma8_*` family — so a split would have required
sub-hunk surgery on generated kernel text with no device validation
available. Kept as one commit.
NOT VALIDATED ON DEVICE. Build-only port; none of these tiers has been
re-measured against this branch's kernels.
Co-authored-by: 0xkydo <95952950+0xkydo@users.noreply.github.com>
Co-authored-by: Amal-David <11647194+Amal-David@users.noreply.github.com>
Co-authored-by: DashiellB <65423051+DashiellB@users.noreply.github.com>
Co-authored-by: brandonegg <13079136+brandonegg@users.noreply.github.com>
Co-authored-by: delordemm1 <46292455+delordemm1@users.noreply.github.com>
Co-authored-by: ercumentyildirim <43972346+ercumentyildirim@users.noreply.github.com>
Co-authored-by: exakoss <67432899+exakoss@users.noreply.github.com>
Co-authored-by: i34-9 <313589706+i34-9@users.noreply.github.com>
Co-authored-by: ivanfioravanti <1069210+ivanfioravanti@users.noreply.github.com>
Co-authored-by: jungjipdo <130676635+jungjipdo@users.noreply.github.com>
Co-authored-by: newjordan <11369410+newjordan@users.noreply.github.com>
Co-authored-by: polymorf <127736+polymorf@users.noreply.github.com>
Co-authored-by: rinaldofesta <5622471+rinaldofesta@users.noreply.github.com>
Co-authored-by: rube-de <8930910+rube-de@users.noreply.github.com>
Co-authored-by: samfenwick <45273188+samfenwick@users.noreply.github.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NS86aXV85ai4eLRw1ADEyX
C1 rung 2. Gemma 4 26B-A4B's five global attention layers are 16 query
heads / 2 KV heads at head_dim 512. `has_fused_kernel`'s vector head-dim
list is {64, 96, 128, 192, 256}, so every decode call on those layers is
rejected and `use_fallback` sends it to the unfused
`scale.q -> matmul -> softmax -> matmul` graph (fast.cpp). That graph
unflattens the query to [B, kv_heads, gqa, 1, D] and broadcasts K and V
over the gqa axis, so its two matmuls are batched gemvs over 16 batch
entries against 2 distinct planes: each key plane and each value plane is
streamed once PER QUERY HEAD -- eight times per layer -- and a
[B, 16, 1, kL] bf16 score plane is materialised, written once and read
twice.
`sdpa_vector_2pass_1` and `sdpa_vector_2pass_2` are already templated on
D and V, so no new kernel body is needed: this instantiates them at
512/512 and admits the dim.
* kernels/scaled_dot_product_attention.metal: a 2-pass-only instantiation
macro plus `..._2pass(type, 512, 512)` and
`..._aggregation(type, 512)`. The single-pass `sdpa_vector` twin is
deliberately NOT instantiated -- it holds q, k and o at D/32 floats each
and launches at 1024 threads per threadgroup, which at D = 512 is 48
live floats against the Metal maximum thread count, an occupancy claim
the split-K kernel (32 x gqa_factor threads, 32 live floats) does not
make.
* scaled_dot_product_attention.cpp: `eval_gpu` routes every D = 512
vector call to `sdpa_vector_2pass` at ANY key length, since there is no
single-pass instantiation to fall back to. The 2-pass form is
length-generic: blocks that see no key leave `sums = 0` and
`maxs = finite_min`, which the merge pass folds in with weight
`exp(finite_min - max) == 0`.
Prefill is untouched: at query length > 8 the call takes the full
attention branch, whose head-dim list is unchanged, so it keeps falling
back exactly as before. MTP verify rectangles are also untouched --
`query_sequence_length * gqa_factor > 32` rejects them at gqa 8 and any
L > 4.
NOT bit-exact against the unfused graph: the split-K kernel carries an
online (running-max) softmax and folds `blocks` partials in a second
pass, so the reduction order over the key axis differs. The bar is
greedy-token parity.
Switch DARKBLOOM_GEMMA4_D512_DECODE_2PASS, default ON; any of
{0, false, no, off} removes the admission and restores the unfused graph
byte for byte.
Also instantiates `sdpa_vector_2pass_1_gqa` at 512/HPT=2 and extends the
`_gqa` kernel-name condition to reach it, behind
DARKBLOOM_GEMMA4_D512_DECODE_2PASS_DEDUP, DEFAULT OFF. The plain 2-pass
kernel gives each query head its own simdgroup and so still reads each
K/V byte gqa_factor times; the dedup variant reads it gqa_factor / HPT
times and is the only form that actually removes the redundant stream.
It is off by default because at D = 512 a thread holds
HPT * (D / 32) * 2 = 64 live floats at 256 threads per threadgroup, and a
pipeline whose `maxTotalThreadsPerThreadgroup` came back under 256 would
make `check_kernel_threadgroup_size` throw rather than degrade. Turn it
on once that is measured on the device.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014L39jXT8ReGzjxUoUKLfan
…up limit The DEDUP arm failed pipeline creation on the device: [metal::Device] Unable to load kernel sdpa_vector_2pass_1_gqa_bfloat16_t_512_512_nomask_qnt_nc_nosinks_128: Threadgroup memory size (32896) exceeds the maximum threadgroup memory allowed (32768) Cause. `sdpa_vector_2pass_1_gqa` publishes its cross-simdgroup merge plane as `threadgroup U o_sh[G * HPT * V]`, U = float. At (G 8, HPT 2, V 512) that is 8192 floats = 32,768 B on its own, and `se_sh` + `mx_sh` (16 floats each) put it 128 B over the limit. Exactly the reported 32,896. `blocks` is not a term in that expression, so tuning MLX_SDPA_BLOCKS could not have moved it; the shipped (64, HPT 8) and (128, HPT 4) instantiations both land at 16,640 B, which is why the limit had never been reached before. Fix. A `SPLIT` template parameter: the plane is published in SPLIT passes of V / SPLIT columns, so it allocates `G * HPT * V / SPLIT` floats. The 512 instantiation takes SPLIT = 2 -> 16,512 B, in line with the shipped two. The existing instantiations take SPLIT = 1 and are unchanged. SPLIT does not change the arithmetic: * each lane keeps the same register slice, and the shared plane is only a scratch relabelling of that slice -- write and read use the identical lane mapping (`simd_lid * v_per_pass`), never the global column index, so the plane's internal layout never has to match the output column order; * `gmax` and `denom` are computed once, on pass 0, from the full `mx_sh` / `se_sh` arrays, in the same order over s; * every `acc[i]` keeps its accumulation order over s inside its pass, and no output element is touched by more than one pass; * SPLIT = 1 is the shipped body instruction for instruction -- the publish of the plane and of the scalars stay in one loop, there is still exactly one barrier before the merge, and the extra write-after-read barrier is guarded by `p > 0`. Also adds three `static_assert`s so this class of failure cannot reach a device again: SPLIT must divide V and V / 32, and the threadgroup allocation must fit 32,768 B. Verified both ways with the offline gate -- `xcrun metal -c` passes at SPLIT = 2 with no warnings, and temporarily setting the 512 instantiation back to SPLIT = 1 reproduces the device failure as a compile error naming `sdpa_vector_2pass_1_gqa<float, 512, 512, 8, 2, 1>`. Symbol check on the resulting .air: `sdpa_vector_2pass_1_gqa_*_512_512` present for all three types at <512, 512, 8, 2, 2>, and the shipped 64/128 kernels still at <..., 1>. DEDUP stays DEFAULT OFF. The remaining unmeasured claim is registers: a thread holds q[2][16] + o[2][16] + kr[16] + vr[16] + acc[16] = 112 live floats at 256 threads per threadgroup, and a pipeline whose `maxTotalThreadsPerThreadgroup` came back under 256 would throw from `check_kernel_threadgroup_size`. Turn it on with DARKBLOOM_GEMMA4_D512_DECODE_2PASS_DEDUP=1 for its own arm. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Gajesh2007
approved these changes
Sep 5, 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
2 participants
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.
Gemma 4: D512 two-pass SDPA + GQA dedup kernels
This is the base (kernel) member of a four-repo stack for Gemma 4 26B-A4B. It adds
the MLX kernels the measured decode path needs.
mlx-swiftpins this commit.1. What this PR changes
vector SDPA does not reach the measured decode rate at head dim 512.
per query group.
attention.
kernel instead of materializing it.
expert gather.
Each item is a default-on path. The C++ env gate
=0restores the stock path.2. Why
The single largest decode lever is the D=512 two-pass SDPA
(
DARKBLOOM_GEMMA4_D512_DECODE_2PASS, +9.4%; with_DEDUP, +10.4%). It is an MLXkernel reached through
MLXFast.scaledDotProductAttention. It is not an engineSwift switch. So the engine cannot reach the measured decode on stock MLX.
Base port serial decode, one prompt, 17,408 tokens, 1,024 output, batch 1, M5 Max:
3. How it was tested
4. What is intentionally NOT changed
=0env gate to the stock path.5. Stack
Merge bottom-up: mlx -> mlx-swift -> mlx-swift-lm -> d-inference.
Pin dependencies:
Source/Cmlx/mlx= bb794a7.6. Known limits
falls per ~6,000 tokens. HumanEval and MBPP gate this class (see the mlx-swift-lm
member).
🤖 Generated with Claude Code
https://claude.ai/code/session_016mnuocRN7JMaSmAjSWBRWw