Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator - #1559
Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator#1559wyanzhao wants to merge 4 commits into
Conversation
…parator The gated delta prefill kernel is latency-bound: each timestep performs two full 32-lane SIMD reductions per value row, and with Dk=128 each lane holds only four state elements. This packs eight independent value rows into a SIMD-group: four lanes own each row and each lane keeps 32 contiguous state elements in registers. Numerics contract: the packed kernel reproduces, bitwise and by construction on any device, an explicit-tree comparator kernel that writes out the ascending butterfly (shuffle_xor 1,2,4,8,16) instead of calling simd_sum; the tests pin packed == comparator with array_equal across shapes, dtypes and lengths. On all current Apple GPUs simd_sum lowers to this same tree (a canary test asserts comparator == existing kernel and would flag any future device or toolchain where the lowering differs, and metal_kernel compiles with MathMode::Safe so the written order is preserved). On the validated M5 toolchain the packed kernel is therefore also bit-identical to the existing kernel, including full-model logits. Blast radius is limited to the shapes the packed kernel serves (Dk == 128, Dv % 8 == 0, scalar gating, no padding mask; generic in batch, head counts and element type). Masked, vector-gate and other-shape paths keep the pre-existing simd_sum kernels byte-for-byte, and the MLX_GDN_PACKED=0 kill-switch routes back to the pre-existing kernel as well. New ops-reference tests cover the masked, vector-gate and Dk=64 generic paths. Timings on M5 Max, Qwen3.6-35B-A3B shapes (B=1, Hk=16, Hv=32, D=128, T=2048): ~1.8-2.0x per layer across measurement rounds (e.g. 2.54-2.56ms -> 1.34-1.38ms on layers 0/20/38); end-to-end prefill 1.064x (median of six saturated adjacent A/B pairs, 5/6 positive).
|
Ran this on an M5 (Apple, Neural-Accelerator / NAX path active) against NAX active (canary)bf16 4096³ GEMM, mlx 0.32.0 wheel (NAX-class), one build/venv per side, mlx-lm varied only:
Both well under the ~9 ms NAX-off floor, so the NAX path is live in both. (A couple of the whole-model runs drifted a bit hot from unrelated background I/O on the box — flagged below; it doesn't touch the correctness verdict.) Losslessness — token-identical ✅ (the headline)Decoded 96 greedy tokens from the same prompt under base (original All four sequences are byte-identical: Kernel-level numerics on model-realistic (L2-normalized q/k) inputs confirm it: Kernel A/B — packed vs unpacked (isolated, one process, canary-clean)Whole-model prefill A/B was too noisy to read a magnitude off — the GDN step is a fraction of total prefill (MoE GEMMs dominate) and thermal/background-I/O drift swamped the signal (base and head traded places across reps). So the timing below is the step kernel in isolation,
Consistent ~2× on the changed kernel, lossless, no regression on the routing that stays on the original path (masked / vector-gate / |
Port the packed Dk=128 prefill specialization from upstream PR ml-explore#1559 onto this fork's fused gated-delta convention, and make it the default path for the shapes it supports. Why --- The generic kernel assigns one 32-lane SIMD-group to each value row. At Dk=128 that leaves every lane holding only four state elements while still spending two full-SIMD reductions per value row per token, so prefill is dominated by reduction traffic rather than by arithmetic. What the packed kernel does --------------------------- It packs eight value rows into a SIMD-group: four lanes own each row, and each lane keeps Dk/4 = 32 contiguous state elements in registers. Both reductions shrink from full 32-lane sums to two shuffles inside a four-lane row group. Launch geometry changes accordingly, from grid (32, Dv, B*Hv) / threadgroup (32, 4, 1) to grid (32, Dv//8, B*Hv) / threadgroup (32, 2, 1). Bitwise contract ---------------- The packed kernel and a new unpacked comparator (_make_gated_delta_kernel_xtree) both write the ascending butterfly (shuffle_xor 1,2,4,8,16) out in source instead of calling simd_sum, so the reduction order is a contract of this file rather than of the simd_sum lowering. The butterfly's first three levels combine partials that live in a single packed lane, and the last two map onto shuffle_xor(1) and shuffle_xor(2) within the four-lane row group; each 4-element partial keeps the comparator's sequential accumulation order. The packed kernel is therefore bit-identical to the comparator on any device by construction, and measured bit-identical to the pre-existing simd_sum kernel on M4. Eligibility and fallback ------------------------ The packed path requires mask is None, scalar gating (a.ndim == 3), Dk == 128, Dv % 8 == 0 and a float32 state. It is otherwise generic in B, Hk, Hv and the input element type. Masked, vector-gate and Dk != 128 shapes keep the original kernels untouched. Unlike the unfused upstream variant there is no g.dtype == mx.float32 eligibility condition: in this fork's API g is not an array, it is computed in-kernel in float32 from (a, A_log, dt_bias). The fused gate arithmetic is lifted verbatim from the generic kernel's g_compute (_FUSED_G) so both paths evaluate the same expression. Kill switch ----------- MLX_GDN_PACKED=0 routes everything back to the pre-existing simd_sum kernels. It is read once at import time, so it must be set before mlx_lm is imported and cannot be flipped from inside a running process. Tests ----- tests/test_gated_delta.py pins packed == comparator and packed == simd_sum kernel across eight shape/dtype cases x T in (1, 7, 64, 257, 2048), plus a canary that simd_sum still lowers to the written-out butterfly, an ops-reference tolerance check, and default-routing, kill-switch and fallback-on-unsupported-shape cases. Random inputs only, no model download required. Context ------- This replaces exo's runtime monkeypatch (src/exo/worker/engines/mlx/patches/gdn_packed.py), whose packed Metal source this kernel is byte-identical to. That patch has been serving since 2026-08-04, measured at +3.7% prompt_tps / +2.6% gen_tps on a 2x M4 pipeline; moving the kernel here measured performance-neutral against it. Ported-from: ml-explore#1559
The README was upstream verbatim, with nothing saying this is a fork or what it changes. Add FORK.md covering the lineage (base branch leo/deepseek-v4, tracking ml-explore/mlx-lm), why the fork exists, the packed gated-delta kernel ported from ml-explore#1559, the MLX_GDN_PACKED kill switch, and the branch map. README.md gets a three-line pointer and is otherwise untouched, so future upstream merges stay conflict-free.
2 participants
Summary
Add a packed Metal kernel for the scalar-gated, unmasked
Dk == 128gated-delta prefill path.
MLX_GDN_PACKED=0restores the existing kernelMotivation
The existing gated-delta prefill kernel assigns one full 32-lane
SIMD-group to each value row. With
Dk == 128, each lane owns only fourstate elements, while every timestep still performs two full-SIMD
reductions per row:
This makes the kernel latency-bound on the target Qwen3.5/3.6 shapes.
Implementation
The packed kernel maps eight independent value rows onto one SIMD-group:
simd_shuffle_xor(..., 1)andsimd_shuffle_xor(..., 2)within each four-lane row group.The first three butterfly levels are evaluated locally inside each packed
lane. Each four-element partial preserves the sequential accumulation order
of the original lane.
Numerical contract
The packed kernel is pinned against a test-only comparator that explicitly
writes the ascending butterfly tree:
This gives the packed implementation a device-independent, source-level
reduction contract. Tests require
array_equalbetween the packed kerneland the comparator across shapes, dtypes, and sequence lengths.
On the validated M5 toolchain,
simd_sumlowers to the same tree, so thepacked kernel is also bit-identical to the existing kernel, including
full-model logits. A canary test compares the explicit-tree comparator with
the existing
simd_sumkernel and will flag a device or toolchain whoselowering differs.
metal_kernelusesMathMode::Safe, preserving thewritten operation order.
Routing and fallback behavior
Dk == 128Dv % 8 == 0g.dtype == float32,state.dtype == float32MLX_GDN_PACKED=0The specialization is otherwise generic in batch size, head counts,
Dv,sequence length, and input element type.
Performance
M5 Max, Qwen3.6-35B-A3B shapes:
2.54-2.56 ms1.34-1.38 ms1.8-2.0x1.064x/+6.4%The end-to-end result is the median of six saturated adjacent A/B pairs;
five of six pairs were positive.
Validation
array_equalto the explicit-tree comparatoracross:
bfloat16,float16, andfloat32inputs;Dvvalues that differ fromDk; andTin{1, 7, 64, 257}.simd_sumkernel.Dk == 64path.tests/test_gated_delta.py: 8/8 passing;pre-commit run --all-files: passing.Checklist
pre-commit run --all-files.MLX_GDN_PACKED=0kill switch.