Skip to content

metal: gate nax off for gen-17 devices (wrong gemm/qmm results on M5-class) #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bri-prism merged 1 commit into from
Jul 6, 2026
Merged

metal: gate nax off for gen-17 devices (wrong gemm/qmm results on M5-class) #4

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion mlx/backend/metal/device.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -797,7 +797,14 @@ bool is_nax_available() {
auto& d = metal::device(mlx::core::Device::gpu);
auto arch = d.get_architecture().back();
auto gen = d.get_architecture_gen();
can_use_nax &= gen >= (arch == 'p' ? 18 : 17);
// gen-17 desktop parts (M5-class, applegpu_g17*) take the nax path under
// the old `gen >= 17` gate but compute WRONG results in the nax steel-gemm
// and qmm_t kernels (measured 2026-07-05: fp16 GEMM max|err| ~4 vs fp32 for
// every shape routed to steel_gemm_fused_nax; quantized qmm_t ~400 abs err;
// non-nax fallback bit-matches stock mlx). Require gen >= 18 on every
// device class until the g17 path is fixed and correctness-gated.
(void)arch;
can_use_nax &= gen >= 18;
return can_use_nax;
};
static bool is_nax_available_ = _check_nax();
Expand Down