perf(mlx): opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder

[Gajesh2007]

What

Adds an opt-in, exactly-scoped Gemma 4 26B-A4B expert-QMM tile path behind MLX_GATHER_QMM_EXPERT_SLICES:

• A distinctly named qmm_t_expert_impl (BM32 expert tiles, BM16 fallback rows) taking a private/by-value row count. The shared qmm_t_impl constant-address ABI (const constant int& M) and every ordinary gathered / batched / dense QMM route are byte-for-byte unchanged.
• build_gemma4_sorted_expert_tiles_bm32 — the descriptor builder runs on one 128-thread threadgroup (parallel binary search of each expert's sorted range + Hillis–Steele exclusive scan + strided upper-bound emission) instead of a single GPU thread serially building all descriptors.
• Selector runs after the existing NAX-first route (gather_qmm_rhs_nax keeps priority — a NAX engagement is recorded as non-engagement, never bypassed) and requires: affine BF16 transposed activations, contiguous uint32 sorted indices, 4-bit gs=64 weights, exactly 128 experts, assignment count ∈ {4096, 8192, 16384}, and one of the two exact production geometries (packed gate/up [128, 1408, 352]; down [128, 2816, 88]). Every miss returns through the unchanged legacy route.
• Device-level one-shot request resolution, non-throwing dual-symbol AOT probe + prewarm, and relaxed-atomic per-primitive diagnostics: requested, aotAvailable, naxAvailable, hits, and mutually-exclusive fallback classes (nax, outer_route, quantization, topology, assignment_count, geometry, metallib_unavailable), with attempts = hits + fallbacks maintained exactly.

Before / After — dispatch behavior

flowchart LR
  subgraph Before
    A1[sorted gather QMM<br/>128 experts, gs=64 b=4] --> B1[NAX available?]
    B1 -- yes --> C1[gather_qmm_rhs_nax]
    B1 -- no --> D1[gather_qmm_rhs_*<br/>generic 16-row tiling]
  end
  subgraph After
    A2[sorted gather QMM<br/>128 experts, gs=64 b=4] --> B2[NAX available?]
    B2 -- yes --> C2[gather_qmm_rhs_nax<br/>counted: fallback_nax]
    B2 -- no --> S{exact selector<br/>opt-in + AOT + dtype +<br/>128 experts + 4096/8192/16384<br/>+ exact geometry}
    S -- hit --> R1[build_gemma4_sorted_expert_tiles_bm32<br/>+ affine_gather_qmm_gemma4_expert_tiles_...]
    S -- miss --> D2[gather_qmm_rhs_*<br/>counted per-class fallback]
  end

Loading

Before / After — code

flowchart TD
  subgraph Before
    Q1[quantized.cpp dispatch] --> Q2[qmm_t_impl<br/>shared constant-address M]
    Q2 --> Q3[generic BM16/BM32 tiles]
  end
  subgraph After
    P1[quantized.cpp dispatch] --> P2{selector gate}
    P2 --> P3[quantized.metal<br/>build_gemma4_sorted_expert_tiles_bm32<br/>affine_gather_qmm_gemma4_expert_tiles_bm_32]
    P3 --> P4[gemma4_expert_qmm.h<br/>qmm_t_expert_impl private M]
    P2 -- miss --> P5[qmm_t_impl — unchanged]
    P6[device.cpp/h] --> P7[request resolution, AOT probe/prewarm,<br/>diagnostics snapshot/reset]
    P1 -. counts .-> P6
  end

Loading

Testing

• tests/gpu_tests.cpp: exact-shape arithmetic parity at all three assignment counts, one-field selector misses, NAX priority, AOT-absence fallback, counter invariant/reset, and zero attempts from ordinary routes.
• Built deterministically (JIT off) through scripts/fetch-metallib.sh in the downstream tree, which refuses artifacts missing either R1 symbol; the resulting mlx.metallib is byte-identical (3f9d85ec…) to the artifact used in the production retention matrix.

Measured standing (honest read)

In the downstream v0.8.2 production retention matrix (gemma-4-26B-A4B-it-qat-4bit, M4 Max):

• Standalone profile: dropped (prefill geomean −10.2% vs bracket).
• Paired with direct weighted-unsort (the retained production tuple): prefill +1.8%, TTFT −7.5%, decode +3.3%, arrival E2E +12.0% vs bracket — retained-final.
• Production engagement is zero on measured paths: 810 armed attempts, all fallbacks (723 outer-route, 87 assignment-count). It ships strictly as an opt-in experiment with hit/fallback observability; orphan/never-hitting claims are not made.

Stack: upstream of Layr-Labs/mlx-swift (Swift bindings + mirrors) → Layr-Labs/mlx-swift-lm → Layr-Labs/d-inference (v0.8.2).

[Gajesh2007]

@codex review

[Gajesh2007]

Deep-review round complete — fixes landed as 8d538a08 (pushed). Summary of the adversarial review (agent-verified with a standalone Metal harness, 14/14 GPU experiments bitwise-clean) and what it produced:

Review finding	Resolution
F1 in-tree tests never execute the kernels on GPU	Documented: ON-path runs (MLX_GATHER_QMM_EXPERT_SLICES=1 swift test --filter SortedGatherQuantizedMMTests) execute real HITs — verified locally on this machine (hits == 1 asserted against the fresh source-matched metallib). Header doc now carries the exact invocation.
F2 descriptor builder trusts global sortedness (legacy was order-agnostic)	Kernel now self-checks: per-thread range verification + simd/threadgroup vote; on violation count[0]=0 and the host synchronizes, reads the count, and re-routes to the legacy path — mis-sorted input degrades to correct-but-slow instead of silently wrong.
F3 armed_ plain bool race	std::atomic<bool>, relaxed semantics, explicit order preservation.
F4 bias re-normalization on fallback calls	normalization now happens only inside route == hit.
F5 route-predicate literal duplicated	single static constexpr takes_sorted_rhs_route(...) feeds dispatch and counters.
F6/F7 nits	recorded.

Verification: gpu_tests 261/261 (3531 assertions); xcrun metal -Wall -Wextra -fno-fast-math 0 warnings; symbol gates all present.

Disclosure: the fail-safe re-route requires a CommandEncoder::synchronize() (stream drain) on every requested-hit call before the host chooses the tile vs legacy route. R1 hits are production-inert today (exact-shape selector), but this cost is real for future engaged runs and will be re-measured in the clean perf gate before any engagement flip.

[Gajesh2007]

@codex review

[Gajesh2007]

Final review wave (#2) complete. Every prior finding is closed; the final verifier pass signed off each repo:

• mlx 9b0d1b4c: the sortedness fail-safe is now a sound oracle (strided adjacent-pair scan ⟺ globally non-decreasing — intra-segment inversions can't escape); retracts get a dedicated fallback_sortedness_retracted counter (diagnostics ABI 80→88 appended-only, facades byte-synced, pins updated); new doctests pin the arm/disarm cycle. 262/262 cases, 3550 assertions.
• mlx-swift 38eaa93: mirrors byte-identical to canonical; generator post-condition guarantees one MLX_API guard or fails; facade pivoted to the 88-byte ABI; OFF (6 tests, 2 conditional skips) and ON (6/6, real R1 hits on this machine) both green.
• mlx-swift-lm 06046c3: the LAST textual survivor of the k_eq_v-gated head rule is gone (validateAttentionProducts) — rule eradication is now complete across model, CBv2 layer kinds, MTP validation, and sizing. 599+471 green twice.
• d-inference b0fa08af6: UpdateCommand now falls back ONLY for missing/unreadable config (malformed files fail loudly — final swallow hole); preload-gate margin widened for loaded machines (make provider-test: 1692/181, 0 failures on this box).

Perf-gate status unchanged: attribution rerun deferred to a cooled machine (see the prior comment — the box measured ~1.7× slow during the late-night window). The remaining open perf item is the nested PR's per-hit synchronize() cost on engaged R1 shapes (production-inert today); it gets a dedicated measurement in that gate.

@codex review

[Gajesh2007]

Final merge-readiness pass complete:

• Upstream conflicts resolved by merging current base branches; no history rewrite.
• Heads: mlx 9b0d1b4c, mlx-swift 38eaa93a, mlx-swift-lm 29ff95bc, d-inference 6aa321a71.
• Every prior review finding is fixed and every stale thread is resolved.
• LM preserves upstream paged/CBv2 work and adds focused regressions for all shared-tower compatibility edges. Weighted expert reduction is now limited to scheduled CBv2 prefill, eliminating the direct-prefill regression found by the final ablation.
• Root benchmark artifacts now record typed effective Gemma settings across all phases, reject malformed baselines/sample-count drift, load explicit configs read-only, pin KV backend posture, and force fixed decode token budgets.
• Verification: final downstream make provider-test = 2,017 tests / 207 suites plus 82 XCTest; benchmark contracts 59/59; coordinator tests/build green; UI 499 tests + lint (0 errors) + production build green; pre-push hook green.
• Final same-binary A/B correctness gates: exact output invariance PASS, explicit contiguous backend PASS in every phase, delivered arrival topology PASS. Late timing attribution is explicitly caveated because the host drifted during the bracket; the earlier clean attribution epoch remains the performance record.

@codex review