Fix implicit thread address space qualifier becoming explicit in metal 4.1 - #3963

Backport ml-explore#3963 to the Prism fork, retaining its integral_constant additions. Explicitly qualify thread member functions and remove inferred address spaces from type deductions.

On macOS 27, MLX 0.32.1 compiles in-process sources at Metal 4.1 (4.0
on macOS 26). Metal 4.1 gives decltype of a local an explicit `thread`
qualifier, and the MPP trait that gates get_destination_cooperative_tensor
strips cv, reference and pointer but not the address space, so no
overload matched in nax_mma:

  error: no matching member function for call to
      'get_destination_cooperative_tensor'
  note: candidate template ignored: requirement 'false' was not
      satisfied [with LeftOperandType = thread metal::cooperative_tensor<...>]

Under VLLM_METAL_BUILD_FROM_SOURCE=1 the loader caught the error, logged
a warning and kept the tiled kernel, so M5 prefill silently lost NAX.
The prebuilt metallib was unaffected: it compiles at the 26.2 deployment
floor, which defaults to Metal 4.0. The MPP headers are not the cause;
the same headers accept the kernel at 4.0 and reject it at 4.1.

The operand types now go through metal::remove_addrspace_t, as in MLX's
steel NAX kernels (ml-explore/mlx#3963) and the usage example in the MPP
header. At Metal 4.0 the wrap is the identity.

M5 Pro, macOS 27.0, tools/nax_prefill_parity.py: PARITY PASS (60 cases)
in source mode (Metal 4.1) and from the prebuilt metallib (Metal 4.0),
worst 0.87 ulp max / 0.021 ulp mean; every case prints the same as the
prebuilt kernel before this change.

Tests: the NAX source compiles with the build's flags at the deployment
floor and at the SDK's own version (its newest Metal version) on any Mac
whose SDK and Metal toolchain can build NAX; on NAX hardware get_ops()
comes up with the NAX kernel loaded instead of the fallback.

## Problem

On macOS 27 / Xcode 27 the metal build fails in
`scratchy-target-metal`'s build script:

```
mlx_steel_attn/mma.h:268:12: error: reference to type 'thread vec<...>' could not bind to an lvalue of type 'vec<...>'
panicked at crates/targets/metal/build.rs:167:13: `xcrun metal` failed for .../attention_steel.metal
```

The cause is the **Metal toolchain** (Metal 4.1 language, toolchain
`v27.1.266.1`), not the OS itself. Metal 4.1 puts an unqualified `this`
in the new *generic* address space; it used to be implicitly `thread`.
The vendored MLX steel/NAX headers return `thread T&` into members from
unqualified member functions, and that no longer compiles.

## Fix

This ports the upstream fix,
[ml-explore/mlx#3963](ml-explore/mlx#3963):

- Add the `thread` qualifier to member functions (`frag_at`, `load`,
`store`, `next`, ...). It has to spread to every caller, the way `const`
does. Most of these edits were applied with `-Xclang -fixit
-Wmetal-addr-spaces` under the pre-4.1 rules, where the missing
qualifier is only a warning.
- Wrap `decltype`s in `metal::remove_addrspace_t` in the
`integral_constant` binops and in the
`get_destination_cooperative_tensor` template arguments. Same manual
edits as upstream.

Only shader headers change; there are no Rust changes.

## Verification (macOS 27.0, Metal toolchain 27.1)

- `cargo build -Fmetal,model/gemma4-moe,quant/mlx --release`: failed
before, builds now.
- `scr chat -m mlx-community/gemma-4-26b-a4b-it-4bit --device metal`:
prints "The capital of France is Paris."
- The NAX shaders compiled at runtime (`quantized_qmm_nax`, `nax_probe`,
and `attention_steel_nax_paged` with its runtime prefix) give 0 errors
and 0 `-Wmetal-addr-spaces` warnings under both `-std=metal4.0` (what
the runtime uses) and `-std=metal4.1`.
- **Unchanged for older toolchains (macOS 26 / Xcode 26):** compiled
under the pre-4.1 rules (`-std=metal3.2`, or `4.0` for NAX), all 49
offline-compiled shaders produce **byte-identical AIR** before and after
this change. The qualifier only states what the older compiler already
assumed. I had no Xcode 26 toolchain to build with directly.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 5.5 (1M context) <[email protected]>

MLX 0.32.1 (ml-explore/mlx#3963) made its Metal kernels use
`metal::remove_addrspace_t`, which the Metal stdlib in Xcode 15.4 --
the macos-14 runner image's default -- does not provide, so the MLX
build fails there. Select Xcode 16.2, the newest on that image, for
the macos-14 jobs in CI and the AOT release build. The macOS 14
floor comes from the deployment target and is still asserted by the
NIF floor check.

…compile) (#1059)

### Summary
On macOS 27, MLX 0.32.1 asks the runtime Metal compiler for MSL 4.1
(ml-explore/mlx#4052, e190d045). MSL 4.1 makes the implicit `thread`
address space of a local variable explicit in `decltype`. In
`pagedattention_nax.metal`, `nax_mma` passes `decltype(ct_a)` /
`decltype(ct_b)` to `matmul2d::get_destination_cooperative_tensor`.
Under MSL 4.1 those types become `thread cooperative_tensor<...>`, fail
MPP's requires-clause, and the NAX library can't be built:

```
program_source:57:27: error: no matching member function for call to 'get_destination_cooperative_tensor'
MPPTensorOpsMatMul2d.h:562:3: note: candidate template ignored: requirement 'false' was not satisfied [with LeftOperandType = thread ...]
```

With `VLLM_METAL_BUILD_FROM_SOURCE=1` on an M5 running macOS 27, this
logs `NAX prefill attention initialization failed; using the non-NAX
fallback` and prefill quietly falls back to the tiled kernel.

MLX hit the same issue in its own NAX GEMM/attention headers. It fixed
it in ml-explore/mlx#3963 (ef5fc0fa, shipped in 0.32.1) by wrapping the
types in `metal::remove_addrspace_t`. This PR applies the same one-line
fix to our NAX kernel, which was adapted from MLX `steel_attention_nax`
before that fix landed.

### Change
```diff
   auto ct_c = op.template get_destination_cooperative_tensor<
-      decltype(ct_a), decltype(ct_b), CT>();
+      metal::remove_addrspace_t<decltype(ct_a)>,
+      metal::remove_addrspace_t<decltype(ct_b)>, CT>();
```

### Impact
- MSL ≤ 4.0 (macOS 26, or MLX 0.32.0): `decltype(ct_a)` has no address
space, so `remove_addrspace_t<T>` is `T`. The types and codegen are the
same as before.
- Prebuilt/AOT metallib (`python -m vllm_metal.metal.build`,
`-mmacosx-version-min=26.2`, default `-std`): this already compiled and
still does. The change also protects against a future default of MSL
4.1.
- Non-M5 devices: the NAX library is never compiled (`nax_supported()`
is false), so nothing changes.
- Shader-only change. Local prebuilt metallibs will show as stale and
need the usual rebuild.

### Test plan (M5 Mac17,2, macOS 27.0.1 26A434, Xcode
metalfe-32023.921.6, MLX 0.32.1; re-verified on main 94ce7d6)
- [x] Offline compile, `xcrun metal -fno-fast-math
-mmacosx-version-min=26.2 -c`:
- original source: passes with `-std=metal4.0` and the default;
**fails** with `-std=metal4.1` (error above);
  - patched source: passes with `metal4.0`, `metal4.1` and the default.
- [x] Runtime compile (`VLLM_METAL_BUILD_FROM_SOURCE=1`): logs `NAX
prefill attention kernels loaded (M5 tensor units)`, no fallback
warning, `nax_ready() == True`.
- [x] `python tools/nax_prefill_parity.py`: PARITY PASS (60 cases,
fp16/bf16 × head_size 64–512 × block_size 8/16/32 × fast/masked).
- [x] `pytest tests/test_nax_loader.py tests/test_paged_attention.py
tests/test_paged_attention_sinks.py`: 77 passed (re-run on main 94ce7d6,
macOS 27.0.1 26A434).
- [x] `pytest tests/test_gqa_paged_decode.py`: 350 passed (main
94ce7d6).

Signed-off-by: suntp <[email protected]>
Co-authored-by: suntp <[email protected]>

…ber functions)

Metal 4.1 (Xcode 27, metalfe-32023) makes the implicit address space of
`this` in member functions `__metal_generic` instead of `thread`. Returning
`thread T&` to a member then no longer compiles, which broke 17 of 40
kernels, including every quantized matmul (quantized, quantized_nax,
fp_quantized*) and all steel GEMM/attention kernels.

- Backport ml-explore/mlx#3963 (ef5fc0fab2) to the vendored mlx 0.30.6.
  Hunks for code that has since diverged upstream (sort.h LessThan, the
  newer NAXFrag-based nax.h, the Conv3D loaders) do not apply and are
  not needed here.
- Apply the same rule to the older NAXSubTile/NAXTile structs in
  steel/gemm/nax.h and steel/attn/nax.h (32 signatures each), which
  #3963 does not cover because upstream rewrote them.
- Regenerate mlx-generated/ (JIT sources and metal/ headers) from the
  patched kernels. Only the patched-vs-unpatched delta is applied, so
  existing hand edits in mlx-generated/ are kept.

All 39 kernels in the metallib build list compile with Metal 32023.921.
The change is a no-op on older toolchains (the qualifier was already
the implicit default there).

Co-Authored-By: Claude Opus 5.5 <[email protected]>