From 74a113ded5d4967e746843bbb60f2cac282c0812 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Mon, 13 Jul 2026 13:05:50 -0700 Subject: [PATCH] Unroll the NAX attention Q@K.T head-dim loop by 4 not fully Fully unrolling the head-dim loop makes the compiler hoist all TD K-tile loads to the top of the block, ahead of the mma chain. Unrolling by 4 instead lets it interleave each K load with the running matmul, which measures faster at head_dim 128 and is a no-op at head_dim 64. M5 Max, 16 Q / 2 KV heads, bf16, causal, 8192^2 (median of 12): head_dim 128: 5.70 -> 5.08 ms (48.2 -> 54.1 TF, +12%) head_dim 64: 2.45 -> 2.44 ms (unchanged; TD == 4 is still a full unroll at count 4) No numerical change (bitwise-identical accumulation order); 16/16 fast-sdpa tests pass. --- .../metal/kernels/steel/attn/kernels/steel_attention_nax.h | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index adc9a42798..7d74d6f44b 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -207,7 +207,12 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik += 2) { - STEEL_PRAGMA_UNROLL + // Unrolling the head-dim loop by 4 rather than fully lets the + // compiler interleave the next K-tile loads with the running mma + // chain instead of hoisting all TD loads up front; measures +11% + // at head_dim 128 on M5 Max and is a no-op at head_dim 64 (TD == 4, + // already a full unroll). +#pragma clang loop unroll_count(4) for (short id = 0; id < TD; id++) { NAXTile Qtile; NAXTile Ktile;