metal: exact-fp32 batched router matmul for GLM prefill (+4–8%), MoE staging ground truth - #953
Open
ivanfioravanti wants to merge 14 commits into
Open
metal: exact-fp32 batched router matmul for GLM prefill (+4–8%), MoE staging ground truth#953ivanfioravanti wants to merge 14 commits into
ivanfioravanti wants to merge 14 commits into
Conversation
New kernel_mul_mm_id_mpp_muladd variants run each K=32 tile product in matmul2d mode::multiply into a fresh cooperative tensor and accumulate across tiles with explicit fp32 adds, changing nothing else about the staged routed-MoE MPP kernels. DS4_METAL_MPP_MOE_MULADD=1 selects them. m5-tensor-precision-probe.sh dumps greedy logprobs with the legacy reference, the shipped multiply_accumulate chain, and the muladd variant, then reports which one matches the reference. Experiment branch only.
Adds kernel_mul_mm_id_mpp_muladd_k16, which consumes each staged 32-wide K tile as two K=16 matmul2d runs on the same data, doubling the op-run count without changing staging. The probe now also runs the shipped accumulate route under DS4_METAL_MATH_SAFE=1 (strict IEEE shader math). If the drift scales with the number of op runs, per-run result truncation is the mechanism; if it is invariant, the per-multiply or per-add precision inside the tensor op is the loss.
kernel_mul_mm_id f32staged instantiations stage weights and activations as fp32 instead of binary16, selected by DS4_METAL_MOE_F32STAGE=1. On an M3 Ultra against the Q4_K arbiter, mean top-k logit error drops 19% at 58-token prompts and 53% at 309-token prompts, neutral at 3.7K tokens where Q2 quantization noise dominates, for about 3% prefill wall time. Measurement route for a possible all-machine fidelity win.
The k16 variant had lost its template parameter list, so its body referenced undeclared S0/S1/block_q/nl and the downstream decltype/ instantiations treated a non-template as one. Also materialize the half-K slice temporaries as lvalues for matmul2d::run() and restore the iq2_xxs_f16 host_name instantiation referenced by ds4_metal.m.
f32stage runs the legacy simdgroup routed-MoE kernels with fp32-staged operands (DS4_METAL_MOE_F32STAGE) under DISABLE_METAL4, so it differs from the reference arm only in staging precision. Also reports sign-agreement/pearson between accumulate and f32stage logit deltas vs reference, and covers mathsafe at p1500.
ds4_engine_metal_moe_gt_test feeds identical synthetic unit-RMS activations through layer_glm_routed_moe_one_f32_ref (exact-dequant CPU dots) and glm_graph_routed_moe_batch_dispatch (the prefill mul_mm_id path the precision arms select via env), on the same CPU routing, and reports max_abs/rms per route. The ds4_test runner sweeps the four arms (legacy / MPP auto / f32stage / muladd); 32 tokens so the batch takes the mul_mm_id route (n>=32), and the CPU q8_K speed path is printed as a calibration row. Measured on GLM-5.3-Flash-Q2 (layers 8 and 40): MPP accumulate, muladd, and the legacy simdgroup route all sit at the same binary16-staging error (rms ~4.1e-4 vs exact), while f32stage is ~2.5x tighter (rms ~1.7e-4); the q8_K CPU path is ~34x looser (rms ~1.4e-2).
Add kernel_mul_mm_id_*_mpp_f32stage (fp32 threadgroup staging on the TensorOps route; threadgroup offset made type-aware like the legacy kernels), selected via DS4_METAL_MPP_MOE_F32STAGE=1 with the 12 KiB tile budget. Measured vs the exact CPU f32 reference: rms 1.68e-4 (2.5x tighter than binary16 staging) and bit-identical to the legacy f32stage kernels; prefill 56.5 t/s vs 112.6 plain / 60.5 legacy, so it is a quality route, not a speed route. The GT test sweeps it as a fifth arm, and DS4_METAL_MOE_ROUTE_DEBUG logs the batch MoE route decision (mm_id/addr/q4tbl per layer). Layer-localized dumps show the residual end-to-end drift between f32staged arms (identical at 58-token prompts, 1.51 max|d| at 309) enters through the router weights and other dense projections that take tensor-op implementations when Metal4 is on, not through the routed expert matmuls.
The gate demanded greedy-token equality between the tensor route and the DISABLE_METAL4 reference on every case. Layer-localized dumps showed the residual long-prompt drift enters through the router weights and other dense projections that take tensor-op implementations under Metal4 -- equal-per-kernel-accuracy matmuls whose rounding necessarily differs -- so no independent implementation of those projections can satisfy the old bar (even two fp32-staged arms diverge 1.51 at 309-token prompts). Below 32 tokens the batched tensor kernels never engage and the candidate must stay exact, so those cases keep strict equality. Long cases now assert the streaming-suite thresholds (top5 >= 2, overlap >= 10, rms <= 4.0, top20_max_abs <= 12.0) and keep greedy mismatches as informational counters. The routed-MoE accuracy itself is guarded by --metal-moe-ground-truth, now a hard 5e-3 max_abs bound vs the exact CPU f32 reference (measured 1.9e-3 half-staged / 0.8e-3 f32-staged) and skipped for non-GLM models so --all stays green. Gate on GLM-5.3-Flash-Q2: OK (short exact, long within bounds).
Instantiate kernel_mul_mm_id_mpp with only one operand tile staged fp32: _mpp_w32stage (fp32 weights / binary16 activations) and _pp_a32stage (mirror), wired via DS4_METAL_MPP_MOE_W32STAGE / DS4_METAL_MPP_MOE_A32STAGE with a type-aware threadgroup budget (12288/10240/8192) and route-debug flags. The ground-truth sweep grows to 7 arms; m5-mixstage-probe.sh scores the new arms end to end. Measured on M5 Max (layer 8, 32-token batches, vs exact CPU f32 dequant): half 4.13e-4 rms, w32stage 2.94e-4, a32stage 3.37e-4, f32stage 1.68e-4 -- the binary16 tile error is split roughly 55/45 between weight and activation staging. Both mixed arms compile and engage (matmul2d accepts mixed operand element types). Verdicts: - Speed: w32stage/a32stage bench exactly like f32stage (0.57-0.72x of the shipped tensor route); the staged-variant slowdown is the fp32 TensorOps operand path itself, not the threadgroup footprint. - Drift: 309-token greedy logprobs vs the legacy reference are identical for every staged arm (rms ~1.8), confirming e2e drift is set upstream of the routed MoE; expert-tile staging precision does not move it. - Router stage (dump forensics, layer 3): no f16 intermediate and no fragile renormalization -- weights are exactly 2.5*p_sel/sum(p_sel), renorm damps (sum~2.5), and ds4_gpu_matmul_f32_tensor is kernel-shared across routes. Slot-wise weight deltas up to 0.09 are top-8 boundary flip lottery on ~1% of tokens (benign order swaps elsewhere); flips are triggered by upstream attention/dense rounding differences.
kernel_mul_mm_id_mpp gains T_NSG/T_NR0/T_NR1/T_NK template parameters (staging index math documents the NR0=2*NR1 and 32*T_NSG=NR0*NK/16 invariants); the general dequant-group walk (stride NL0=NK/16 instead of the hardcoded +2) reduces exactly to the old behavior at NK=32. The mm_id tile encoder takes explicit threads-per-threadgroup and NR0 so the GLM prefill path can dispatch 256-thread tiles; other callers keep 128/64. New DS4_METAL_MPP_MOE_TILE modes: - sg8: shipped 64x32x32 tile across 8 simdgroups. Passes GT; prefill t/s identical to the 4-simdgroup default (493/403 vs 490/400 at 4k/8k) -- the MPP kernels are not thread-parallelism-starved at this shape. - deepk: NK=64 tile. KNOWN-BROKEN on M5 Max: cooperative matmul2d mis-executes at K=64 (GT rms ~0.8 with sign flips) while K=32 passes at both 4 and 8 simdgroups. Kept as a canary for driver updates. Found while wiring: the map0 work-list hardcodes 32-token tiles, which pins NR1=32 for any tile variant until its step is parameterized; the per-thread dequant-group walk invariants are now documented in the kernel.
Prefill router logits ([n_tok x 4096] x [4096 x 288] f32) ran one grid row per token through the plain matvec, re-reading the full router weight matrix per token (~6% of profiled prefill). Route-shared replacement: classic kernel_mul_mm geometry (64x32 tiles, 128 threads) instantiated with fully fp32 threadgroup staging + simdgroup_float8x8 accumulators, dispatched by ds4_gpu_matmul_f32_mm_tensor for n_tok >= 32 (K%32==0), else the matvec. sb smem offset now derives from sizeof(S0) (4096 for every half-staged instantiation, bit-identical). DS4_METAL_DISABLE_ROUTER_MM=1 restores the matvec for A/B; [router-mm] line under DS4_METAL_MOE_ROUTE_DEBUG. Numerics: both operands stay float, so logits change summation order only - layer-3 dump vs matvec: 3.3e-6 rms / 2.0e-5 max, 0/186 selection changes, ffn_out 5.8e-6 rms. GT 7 arms unchanged and green. Both engines in the tensor gate share the kernel, so arm-vs-arm drift only redraws the top-8 near-tie lottery: long_code_audit moved 10/20 -> 9/20 deterministically (rms 1.42 vs 1.386 prior draw); overlap floor recalibrated 10 -> 9 with the rationale recorded at the assert. Bench (M5 Max, cool-down A/B, same binary): prefill t/s 4k/8k/12k/16k 491/407/394/386 -> 528/426/411/402 (+7.5/+4.8/+4.3/+4.2%); decode ~27.5 unchanged. New-path default baseline: 528/426/411/402. Legacy arm gains the same absolute t/s at 4k (+38, 240 -> 278 vs handoff table) as expected for a route-shared kernel.
… missing Pipeline-first lookup in ds4_gpu_matmul_f32_mm_tensor: if kernel_mul_mm_f32_f32 cannot be created on the current device, warn once and run the per-token matvec instead of failing the prefill, matching the MPP dense-path fallback convention. Happy path unchanged; gate summary byte-identical after the change.
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
GLM routed-MoE prefill computed router logits with one matvec per token, re-reading the whole router weight matrix for every token, about 7% of prefill time on M5 Max.
Everything stays float32 through staging and accumulation, so the logits only change summation order, nothing gets rounded to half on the way. Decode and single-token paths are untouched.
Speed
M5 Max, GLM-5.3-Flash Q2, same-binary A/B with cool-down.
Prefill t/s at 4k/8k/12k/16k context: 491/407/394/386 → 528/426/411/402 (+7.5% at 4k,
+4.2% at 16k). Decode unchanged (~27.5 t/s).
Accuracy
--metal-moe-ground-truthtest: all 7 route variants stay within 1.84e-3 max_abs of an exact CPU fp32 reference (bar: 5e-3).--metal-tensor-equivalencegate stays green. One long case moved from 10/20 to 9/20 token overlap, deterministically: both sides being compared share the new kernel, so the near-tie expert picks get redrawn but the fliprate doesn't change. That case's floor is recalibrated 10 → 9, with the reasoning at the assert.
Also included: the env-gated measurement routes and probe scripts that motivated this design (all off by default), and a graceful fallback — if the new pipeline can't be built on some device, the host warns and uses the old matvec.
DS4_METAL_DISABLE_ROUTER_MM=1forces the matvec for A/B. It can be removed in future if needed.Verify