metal: gate nax off for gen-17 devices (wrong gemm/qmm results on M5-class)#4
Conversation
…emm/qmm) The gen >= 17 gate let M5-class parts (applegpu_g17s) take the nax path, but the nax steel-gemm and qmm_t kernels are gen-18 targeted and compute wrong results on g17 silicon: fp16 GEMM max|err| ~4 vs fp32 for every shape routed to steel_gemm_fused_nax (both nn and nt, aligned and unaligned tiles); quantized qmm_t is similarly wrong. The non-nax fallback bit-matches stock mlx. On affected devices this silently corrupted transformer prefill (any matmul falling in the routed envelope) and surfaced as model-level output degeneration. Require gen >= 18 on every device class until the g17 path is fixed and correctness-gated. Validation: 5-line plain-gemm probe vs fp32 ground truth; 25-cell M x N sweep all broken pre-fix, all clean post-fix (fp16 rounding, ~0.00098).
|
Upstream context: mlx has an active NAX correctness trail (qmm_t_nax edge-tile overflow fixed in ml-explore#3631, which postdates our base — the overflow line is still present here; ml-explore#3593 adds an MLX_DISABLE_NAX escape hatch; ml-explore#3797 is an open M5 split-K miscompute report). Suggests the re-enable path for g17 is: cherry-pick ml-explore#3631, adopt the ml-explore#3593 hatch, then gate on the 25-cell sweep from this PR's validation. |
|
Added upstream cherry-pick ml-explore#3631 (NAX qmm edge-tile int16 overflow fix) to this PR: it fixes a real correctness bug that applies to gen-18 devices too (the gate only makes qmm_nax dormant on g17), so it belongs regardless. Also measured the g17 re-enable question directly: built a test stack with g17 nax ACTIVE + ml-explore#3631 — the plain steel-gemm-nax breakage persists (same err ~4, matches open ml-explore#3797), and large-N qmm configs stay broken via the same gemm route, while ml-explore#3631 does fix the M>=64 small-N qmm class. Conclusion: gate stays until the steel-gemm-nax bug is fixed (upstream ml-explore#3797 or here); the 25-cell sweep + qmm probe in this thread are the re-enable gates. |
afe6c7d to
1970148
Compare
|
Scope cleanup: dropped my ml-explore#3631 cherry-pick from this PR — ml-explore#3631 and ml-explore#3632 land together via |
…est (#6) * fix(metal): gate NAX to gen-18+ in flattened device header (M5/gen-17 miscompute) The flattened Cmlx copy the SPM build compiles carried the pre-fix gate (gen >= arch=='p'?18:17), leaving NAX enabled on gen-17 M5-class GPUs where the steel-gemm/qmm path miscomputes. Apply the same gen>=18 gate as the mlx core (PrismML-Eng/mlx#4), and add a regression test. Durable fix: bump the mlx submodule past #4 and regenerate via tools/update-mlx.sh. * test: use class setUp() for default-device config (match MLXTests suite convention)
1 participant
What
One-line gating fix (+comment): require
gen >= 18on every device class before enabling the nax path. Previously non-phone parts enabled it atgen >= 17, which put M5-class GPUs (applegpu_g17s) on the nax kernels.Why
The nax steel-gemm and qmm_t kernels compute wrong results on g17 silicon (they are gen-18 targeted):
steel_gemm_fused_nax: max|err| ~4 vs fp32 ground truth (values essentially uncorrelated), for every routed shape — bothnnandnt, aligned and unaligned tiles.qmm_tsimilarly wrong (~400 abs err).Shapes that route to split-K or gemv were unaffected, which made this look shape/composition-dependent downstream: models ran mostly fine at decode (M=1) and corrupted at prefill, surfacing as output degeneration that was chased through converters, packing, schedulers, and recurrence code before landing here.
Validation
x[16,5120]fp16 @ W.T[5120,10240]— pre-fix err 3.94, post-fix 0.000976 (== stock).MLX_METAL_GPU_ARCH=applegpu_g15(or any pre-nax arch) — the env override path is unchanged by this fix.Notes
Gate-first fix: g17 hardware may well be able to run a corrected nax path, but that kernel work should land behind a correctness gate (repro above) rather than a device-class assumption. Until then g17 falls back to the standard steel path.