perf(metal): E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill + trust mode (unit 1.22-1.32x, e2e +15.2% @8k)

[Gajesh2007]

Problem

Qwen3.6-35B-A3B MoE prefill routes its expert QMMs through the legacy gather_qmm path: E=256 experts with real multinomial top-8 routing produce expert row-boundaries that almost never align to 16-row tiles, so the legacy kernel pays tile inflation on every routed matmul (the exact diagnosis already fixed for Gemma E=128 by the expert-tile route — which is hard-wired to E=128 and rejects Qwen's expert count and geometries).

On top of that, the route's sortedness-retract fail-safe reads a GPU counter on the host: one full stream drain per routed gather = 120 drains per 512-token chunk (3 gathers × 40 MoE layers), which erases the kernel win end-to-end.

Change

1. Templatize the sorted expert-tile descriptor builder on expert count (template <int NE> build_sorted_expert_tiles_bm32, quantized.h): NE=128 keeps the Gemma4 host name; NE=256 instantiates build_sorted_expert_tiles_bm32_e256 (quantized.metal). The tile kernel itself is unchanged and shared.
2. Admit E=256 with Qwen geometries in classify_gemma4_expert_qmm (gemma4_expert_qmm.h): fused gate_up [256,1024,2048], split gate/up [256,512,2048], down [256,2048,512]; geometry table tied to expert count with cross-rejection both directions.
3. Per-E descriptor kernel + dispatch in try_gemma4_expert_qmm (quantized.cpp), with a documented max_tile_count upper-bound proof.
4. MLX_GATHER_QMM_EXPERT_SLICES=trust (device.cpp/h, quantized.cpp): skips only the retract-counter host readback — the grid is already over-dispatched with device early-exit, so trust changes no dispatch. Under trust, a genuinely mis-sorted sorted-flag call yields undefined output for that matmul instead of legacy fallback; the only sorted-flag producer (mlx-swift-lm SwitchGLU) sorts on-device immediately before the call. TODO(1.3) documents the device-side fallback that would make trust unconditional.
5. AOT gate fail-closed on all 3 symbols (both builders + shared tile kernel, device.cpp): a metallib missing any symbol disables the route entirely.
6. Includes the gpu::eval UAF fix (ab5e1dd) — required: the retract check synchronizes mid-eval, and without the fix the E=256 shapes segfault deterministically. Cherry-picked separately to main as fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval #5 (prod-relevant, should merge first).

Measured (M4 Max 40c, release, prod 20 GB Qwen3.6 artifact)

Unit microbench (isolated gatherQuantizedMM, sorted, realistic multinomial top-8 histogram over 256 experts, T=512/M=4096, median of 25):

unit	legacy	tiles+sync	tiles+trust	trust speedup
gate_up fused K2048 N1024	3.010 ms	2.443	2.284	1.32x
gate split K2048 N512	1.602	1.332	1.217	1.32x
down K512 N2048	1.537	1.368	1.208	1.27x
gate_up fused T1024 (M=8192)	4.394	3.719	3.563	1.23x
down T1024	2.264	1.990	1.855	1.22x

MoE routed block T=512 (shipping split-legacy 4.741 ms → fused+tiles+trust 3.492 ms) = −26.3%. End-to-end prefill (with the mlx-swift-lm fuseGateUp PR): 8k prompts 1243 → 1433 tok/s (+15.2%), 2k +7.4%, 32k +6.6%. With the retract drain left on, the unit win nets ≈0 e2e (1364 vs 1364 @8k) — that is what trust recovers. Uniform tile-aligned histograms flip the unit result (legacy wins ~15%), confirming the tile-inflation diagnosis: the win comes exactly from expert boundaries not aligning to 16-row tiles.

Behavior

flowchart LR
  subgraph Before
    A1["Qwen E=256 routed QMM"] --> B1["route rejects E!=128 → legacy gather_qmm<br/>16-row tile inflation on every expert boundary"] --> C1["+ retract fail-safe = host drain ×120/chunk<br/>8k prefill 1243 tok/s"]
  end
  subgraph After
    A2["Qwen E=256 routed QMM"] --> B2["E=256 descriptor builder → shared tile kernel<br/>unit 1.22–1.32x"] --> C2["trust skips retract readback only<br/>8k prefill 1433 tok/s (+15.2%)"]
  end

Loading

Code

flowchart LR
  subgraph Before
    D1["quantized.h: build_gemma4_sorted_expert_tiles_bm32 (E=128 hard-wired)"] --> E1["gemma4_expert_qmm.h: classify rejects E=256"] --> F1["quantized.cpp: retract sync always drains"]
  end
  subgraph After
    D2["quantized.h: template<int NE> builder — NE=128 keeps Gemma name, NE=256 new symbol"] --> E2["gemma4_expert_qmm.h: E∈{128,256} topology + per-E geometry table, cross-rejected"] --> F2["quantized.cpp: per-E descriptor select + tile-count bound; trust skips readback"]
    F2 --> G2["device.cpp: env incl. trust; 3-symbol AOT fail-closed gate"]
  end

Loading

Tests

• tests/gpu_tests.cpp "Qwen 3.6 expert QMM pure route table": hits ×3 geometries ×{4096,8192,16384}; topology/geometry/assignment/quantization/metallib misses; E×geometry cross-rejection both directions.
• GPU kernel suite (via mlx-swift SortedGatherQuantizedMMTests, route on): 10/10 — Gemma suite intact; Qwen closed-form fixtures all 3 geometries at M∈{4096,8192,16384} incl. empty experts, unaligned multi-boundary tiles, 255×1-row fragmentation; random-tensor equivalence vs legacy and vs dequantized gatherMM; sortedness-violation retract→legacy with correct output; T=128 legacy fallback by design.

Stack

Merge order (top to bottom):

1. fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval #5 — gpu::eval UAF fix (base main) — prod-relevant crash fix, merge first
2. perf(metal): E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill + trust mode (unit 1.22-1.32x, e2e +15.2% @8k) #6 — E=256 expert tiles + trust (base codex/gemma4-autoresearch-v0.8.2, the line mlx-swift pins Source/Cmlx/mlx to) · perf(metal): E=256 expert-tile route + trust + gpu::eval UAF fix — darkbloom-base mirror #7 — same commits mirrored onto darkbloom-base (the line d-inference pins libs/mlx to) ← this PR
3. perf(mlx-swift): Cmlx bump to E=256 expert-tile route (Qwen 3.6) + regenerated sources + Qwen kernel tests mlx-swift#12 — Cmlx bump to 76663df + Qwen kernel tests (base main)
4. perf(cbv2-mtp): Qwen3.6 GDN capture-verify + target-prefix acceptance (1.53x vs serial verify) mlx-swift-lm#106 — GDN capture-verify — merged (squash d06da67)
5. perf(qwen35): fuse routed-expert gate_up + weightedExpertSum — E=256 tile-route consumer mlx-swift-lm#107 — fuseGateUp consumer, +15.2% prefill @8k (base main)
6. perf(qwen36): bump mlx-swift-lm to GDN capture-verify + target-prefix MTP (mlx-swift-lm#106) d-inference#616 — capture-verify pin bump (base master)
7. perf(qwen36): E=256 expert-tile prefill + fused gate_up + trust d-inference#617 — prefill submodule pins + metallib contract + trust env projection (base perf/qwen36-mtp-capture-verify = Fix typo in CMakeLists.txt ml-explore/mlx#616's head)

Chain: mlx UAF fix → mlx tiles (+mirror) → mlx-swift → mlx-swift-lm#106 (merged) → mlx-swift-lm prefill → d-inference#616 → d-inference prefill.

[Gajesh2007]

Closing in favor of #7: consolidating on darkbloom-base as the single integration line (repo convention — #1/#3/#4 all merged there, and #4 already folded codex/gemma4-autoresearch-v0.8.2 into it). #7 carries commit-for-commit identical tile work PLUS the gpu::eval UAF fix on darkbloom-base. mlx-swift consumers should repin Cmlx to the darkbloom-base lineage when #7 merges (tracked in mlx-swift#12).