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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#4,832 1 comentario 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
78/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
cpp
Área
backend, performance

Línea de trabajo

Lee mlx/backend/metal/jit_kernels.cpp en get_steel_attention_nax_kernel y compara los argumentos de plantilla de attention_nax_dsplit con instantiate_attn_dsplit en mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal. Reproduce el problema con el caso bf16 de SDPA proporcionado en un M5 usando una compilación JIT; se considera terminado cuando el kernel dsplit se compila y devuelve resultados que coinciden con los de una compilación precompilada.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
C++
Estrellas
29.1k
Forks
2.4k
Merge medio
3 d 23 h
PR fusionados (30 d)
90

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de ml-explore/mlx

Todos los issues de ml-explore/mlx

Issues similares

Más issues de C++

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.