[BUG] JIT build: attention_nax_dsplit fails to compile (extra bv template argument)

Open Beginner friendly
  • Is something wrong?
#4,832 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
78/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
cpp
Domain
backend, performance

Research direction

Read mlx/backend/metal/jit_kernels.cpp at get_steel_attention_nax_kernel and compare its attention_nax_dsplit template arguments with instantiate_attn_dsplit in mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal. Reproduce with the provided bf16 SDPA case on an M5 using a JIT build; done means the dsplit kernel compiles and returns results matching a precompiled build.

Written by the indexing model from the issue text.

Description

Summary
On an M5 (NAX) GPU, a JIT build of MLX fails to compile the head-dim-split attention kernel the first time a fused SDPA takes it (head dim 256, query length >= 1024, bf16, causal or array mask). The process then aborts:

[metal::Device] Unable to build metal library from source
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h:950:122: error: reference to overloaded function could not be resolved; did you mean to call it?
template [[host_name("steel_attention_dsplit_bfloat16_bq64_bk32_bd256_bv256_wm4_wn2_maskbfloat16")]] [[kernel]] decltype(attention_nax_dsplit<bfloat16_t, 64, 32, 256, 256, 4, 2, bfloat16_t>) attention_nax_dsplit<bfloat16_t, 64, 32, 256, 256, 4, 2, bfloat16_t>;
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h:525:66: note: possible target for call
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h:950:192: error: explicit instantiation of 'attention_nax_dsplit' does not refer to a function template, variable template, member function, member class, or static data member

attention_nax_dsplit is declared template <typename T, int BQ, int BK, int BD, int WM, int WN, typename MaskType = float, typename AccumType = float>, with no BV. The precompiled instantiation (instantiate_attn_dsplit in steel_attention_nax.metal) passes dtype, bq, bk, bd, wm, wn, mtype, float, which matches. But get_steel_attention_nax_kernel in mlx/backend/metal/jit_kernels.cpp passes type, bq, bk, bd, bv, wm, wn, mask_type for both attention_nax and attention_nax_dsplit, so in the dsplit case every argument after bd shifts by one. This is still the case on main.

Builds that use nojit_kernels.cpp aren't affected. mlx-swift always builds jit_kernels.cpp, so every mlx-swift app on an M5 hits this.

To Reproduce

We hit this through mlx-swift 0.33.0, which vendors MLX v0.32.4 (36969c0), running a bf16 Qwen 3.5 model (head dim 256) with a prompt of more than 1024 tokens in one forward pass. The equivalent in Python should need an MLX built with -DMLX_METAL_JIT=ON, on an M5:

import mlx.core as mx

q = mx.random.normal((1, 16, 1024, 256)).astype(mx.bfloat16)
k = mx.random.normal((1, 4, 1024, 256)).astype(mx.bfloat16)
v = mx.random.normal((1, 4, 1024, 256)).astype(mx.bfloat16)
o = mx.fast.scaled_dot_product_attention(q, k, v, scale=256**-0.5, mask="causal")
mx.eval(o)

Expected behavior
The kernel compiles and SDPA returns the same result as the precompiled build. Possible fix: drop bv from the template arguments in get_steel_attention_nax_kernel when split_d is true.

Desktop (please complete the following information):

  • OS Version: macOS 27.2, Apple M5 Max
  • Version: MLX v0.32.4 (36969c0) through mlx-swift 0.33.0; the same code is on main

Additional context
Head dim 512 (for example Gemma 4's global layers) takes the same dsplit path for long causal prefills. Inputs under 1024 query tokens take the unfused path and don't crash.

Dominant language
C++
Stars
29.1k
Forks
2.4k
Avg merge
3d 20h
Merged PRs (30d)
103

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from ml-explore/mlx

All issues in ml-explore/mlx

Similar issues

More C++ issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.