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

オープン 初心者向け
  • 何か問題がありますか？
#4,832 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 3 日以内に返信

まだ誰も着手していません。

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
78/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
cpp
領域
backend, performance

調査の方向性

mlx/backend/metal/jit_kernels.cpp の get_steel_attention_nax_kernel を読み、attention_nax_dsplit のテンプレート引数を mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal の instantiate_attn_dsplit と比較してください。M5 で JIT ビルドを使用し、提供された bf16 SDPA ケースで再現してください。完了条件は、dsplit カーネルがコンパイルされ、プリコンパイル済みビルドと一致する結果を返すことです。

索引モデルが issue の本文から書いたものです。

説明

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.

主要言語
C++
スター
29.1k
フォーク
2.4k
平均マージ
3日 23時間
マージ済み PR（30日）
90

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

ml-explore/mlx のほかの issue

ml-explore/mlx の issue をすべて見る

似ている issue

C++ の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。