[jonathan308]

Summary

get_steel_attention_nax_kernel passed bv to both attention_nax and attention_nax_dsplit. attention_nax gained a BV template parameter in the upstream merge (#38); attention_nax_dsplit has none (its value head dim is BD). With bv in the list, every later argument shifts by one, an integer lands in the MaskType slot, and the runtime compile always fails:

[metal::Device] Unable to build metal library from source
steel_attention_nax.h:950: error: explicit instantiation of 'attention_nax_dsplit' does not refer to a function template ...

JIT builds (mlx-swift, and so Darkbloom provider 0.10.2) compile this kernel at runtime on NAX GPUs (M5) for full attention with head dim 256 or 512 and at least 1,024 query tokens. The error trapped the process, and provider 0.10.2 crashed on M5 Macs serving Gemma 4 and Qwen 3.5/3.6/3.8. Builds that use the precompiled metallib were unaffected, because steel_attention_nax.metal instantiates the split kernel without bv.

The fix passes the split kernel the same arguments as its metallib instantiation. The non-split kernel is unchanged.

Before and after

flowchart LR
  subgraph Before
    A1["get_steel_attention_nax_kernel(split_d)"] --> B1["attention_nax_dsplit<T, bq, bk, bd, bv, wm, wn, mask>"] --> C1["compile error -> throw -> process trap"]
  end
  subgraph After
    A2["get_steel_attention_nax_kernel(split_d)"] --> B2["attention_nax_dsplit<T, bq, bk, bd, wm, wn, mask>"] --> C2["library + pipeline built"]
  end

Loading

Validation

• Runtime compile of the exact JIT source (metal::utils() + metal::steel_attention_nax() + template definition, language 4.1), on an M3 Ultra and an M5 Ultra: the shipped dsplit instantiations for head dim 256 and 512 fail with the error above; the fixed ones compile, and on the M5 their pipelines build (256 threads, 16 KB threadgroup memory).
• End to end on an M5 Ultra (mlx-swift eedca444, real MLXFast.scaledDotProductAttention, bf16, against a float32 reference):
  • before: Trace/BPT trap: 5 with the error above on the first head-dim-256 call;
  • after: 6 of 6 shapes pass (max abs error ≤ 0.0089): D256 causal 1024, D256 GQA 16/4 causal 2048, D256 causal 1024 after 512 cached keys, D256 array mask 1024, D512 GQA 16/8 causal 2048, D512 causal 1024 × 32 heads.
• This repository's CI builds without JIT, so it does not compile jit_kernels.cpp; mlx-swift and d-inference CI do.

Pinned for the provider hotfix from branch hotfix/0.10.2-jit-nax-dsplit (31f550af, on ff4bb35b); this PR brings the same commit to main.