From 84d84368cf69e310fd611dcab70162bbf6f84647 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Mon, 13 Jul 2026 12:52:53 -0700 Subject: [PATCH 1/2] Add a fused full-attention path for head_dim 256 on NAX devices use_fallback currently allows the fused full-attention path only for head_dim in {64, 80, 128}. Models with head_dim 256 (e.g. the Qwen3.5/3.6 family's full-attention layers) fall back to the unfused graph, which materializes the [n_heads, qL, kL] score tensor and computes the full rectangle before masking. A naive bd=256 instantiation of steel_attention_nax runs at half the efficiency of the bd=128 kernel on every shape (22 vs 48 TF on M5 Max): the per-simdgroup accumulator working set doubles (16 + 4 f32 output and score fragments), and accumulator bytes are the resource that gates tensor-unit throughput at this width (a probe with fp16 accumulators recovers +51%; static register pressure, scheduling fences, accumulation chain order, instruction footprint and Q reload traffic were each ruled out by measurement). This adds a head-dim-split configuration instead (bq=64, bk=32, wm=4, wn=2): the second warp dimension splits D, so each simdgroup computes Q@K.T partials over one half of D (exchanged pair-by-pair through a 16KB threadgroup buffer in per-lane-linear layout, which needs no coordinate math because both halves share the fragment-to-lane map), runs softmax redundantly on the full score tile, and accumulates P@V for its half of Dv. That puts the accumulator set (2 S + 8 O fragments, 320B/lane) at exact bd=128 parity, and throughput follows: 42.8 TF at 8192^2 causal, 90% of the bd=128 kernel. Q fragments stay register-resident across the KV loop, and mma chains keep a single accumulator (alternating accumulators measures -25%). Routing and raggedness (all numbers M5 Max, 16 Q / 2 KV heads, bf16, causal): - The fused path is taken for causal shapes with qL >= 1024 and no array mask, where it beats the unfused graph everywhere measured, offset rectangles included: 2048x8192 5.77 vs 8.92 ms, 4096^2 3.45 vs 8.87 ms, 1024x8192 3.58 vs 4.65 ms, 2048^2 1.19 vs 2.39 ms. Below ~1k queries against a long cache there are too few query blocks to fill the machine (512x8192 measured 1.08x on an earlier variant), so those keep the fallback. - Ragged lengths are padded up to the tile sizes at dispatch and the padded query rows sliced off on the way out: the unaligned (align_Q/align_K = false) pipelines slow every block down ~3x at this width, and real decoder prefills are essentially never tile-aligned. With qL_off computed from the true lengths the causal bound c <= r + qL_off never reaches a padded key column; pads are zero-filled so padded value rows cannot poison P @ V; cost is ~0.6 ms against a ~49 ms cliff. - float32 with tf32 disabled keeps the fallback, since only the NAX kernel is instantiated at this width. Transient memory at 4096^2 drops 663 MB -> 75 MB (no score materialization), and the fused output is closer to an f32 reference than the unfused bf16 graph (relative L2 1.6e-3 vs 2.3e-3). End to end (Qwen3.6-35B-A3B 4-bit, 8k ragged prompt): default 2048-token chunked prefill: 4039 -> 4193 tok/s (+3.8%) single-chunk prefill: 4209 -> 4822 tok/s (+14.6%) Tests: fallback shapes and fused shapes (aligned and ragged squares, a boundary rectangle and a ragged rectangle) against the primitives reference in float32 and bfloat16, plus a subprocess case with MLX_ENABLE_TF32=0 covering the float32 fallback predicate (a mismatch there fails as a kernel-load error). --- .../steel/attn/kernels/steel_attention_nax.h | 345 +++++++++++++++++- .../attn/kernels/steel_attention_nax.metal | 1 + .../metal/scaled_dot_product_attention.cpp | 116 +++++- python/tests/test_fast_sdpa.py | 74 ++++ 4 files changed, 521 insertions(+), 15 deletions(-) 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..287ee7a8fa 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 @@ -127,7 +127,10 @@ template < // Prepare MMA tiles constexpr short kU = 16; - constexpr int kNWarps = WM * WN; + // With WN == 2 the second warp dimension splits the head dim (each + // simdgroup owns half of D for Q@K.T and half of Dv for P@V); only WM + // warps split the Q sequence. + constexpr int kNWarps = (WN == 2) ? WM : WM * WN; static_assert( BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -140,6 +143,346 @@ template < constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); + + if constexpr (WN == 2) { + // Head-dim split path. Per simdgroup the accumulator working set is + // half of the WN == 1 kernel (TD/2 output fragments), which is the + // resource that gates bd=256 throughput. The Q@K.T partial sums over + // each half of D are exchanged through threadgroup memory; both halves + // then run softmax redundantly on the full S tile and each computes + // P@V for its own half of Dv. + constexpr int TDh = TD / 2; + constexpr int BDh = BD / 2; + + const short row_group = simd_group_id / 2; + const short d_half = simd_group_id & 1; + + using otile_t = NAXTile; + otile_t Otile; + Otile.clear(); + + const short tm = kU * TQ * row_group; + Q += tm * int(params->Q_strides[2]) + d_half * BDh; + K += d_half * BDh; + V += d_half * BDh; + O += tm * int(params->O_strides[2]) + d_half * BDh; + + constexpr short kRowsPT = otile_t::kRowsPerThread; + + metal::vec max_score; + metal::vec sum_score{0}; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + } + + if (has_sinks) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = M_LOG2E_F * static_cast(sinks[tidl.y]); + sum_score[i] = 1; + } + } + + int kb_lim = params->NK; + int kb_min_causal = params->NK; + + if (do_causal) { + int q_max = (tid.x + 1) * BQ + params->qL_off; + kb_lim = (q_max + BK - 1) / BK; + kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); + } + + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; + + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + + // One slot per (row group, half): a fragment pair in per-lane-linear + // layout. Both halves share the fragment-to-lane mapping, so the + // exchange needs no coordinate math. + threadgroup AccumType s_xchg[WM][2][2 * kEPF * 32]; + + // Keep the simdgroup's Q half resident in registers for the whole KV + // loop (TDh fragments of T are cheap; see the bd=256 notes). + NAXTile Qtiles[TDh]; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + const int Q_load_off = id * kU; + if (!align_Q && is_last_q) { + Qtiles[id].load_rows( + Q + Q_load_off, int(params->Q_strides[2]), lim_rows_q); + } else { + Qtiles[id].load(Q + Q_load_off, int(params->Q_strides[2])); + } + } + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); + const short sm = simd_coord.y; + const short sn = simd_coord.x; + + // Loop over KV seq length + for (int kb = 0; kb < kb_lim; kb++) { + const int is_last_k = (kb == (params->NK_aligned)); + + stile_t Stile; + Stile.clear(); + + // S = Q @ K.T, this half of D only, exchanged pair by pair. + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Ktile.load_rows( + K + K_load_off, + int(params->K_strides[2]), + lim_rows_k - ik * kU); + } else { + Ktile.load(K + K_load_off, int(params->K_strides[2])); + } + + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + + // Exchange the partial pair and reduce. + threadgroup AccumType* slot = s_xchg[row_group][d_half]; + thread auto& s0 = Stile.frag_at(0, ik); + thread auto& s1 = Stile.frag_at(0, ik + 1); + const short base = short(simd_lane_id) * (2 * kEPF); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + i] = s0[i]; + slot[base + kEPF + i] = s1[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s0[i] += peer[base + i]; + s1[i] += peer[base + kEPF + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // Scale S + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= float(scale2); + } + + // Mask out length sequence + if (!align_K && is_last_k) { + constexpr auto neg_inf = Limits::finite_min; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const short col_pos = ik * kU + sn; + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; + } + } + } + } + + // Mask out if causal + if (do_causal && kb >= kb_min_causal) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + params->qL_off + tm; + const int base_col = kb * BK; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = base_row + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = (r < c) ? neg_inf : fg[loc]; + } + } + } + } + + // Other masking as needed + if (has_mask) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + tm; + const int base_col = kb * BK; + + constexpr bool is_bool = is_same_v; + using melem_t = typename metal::conditional_t; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; + + if (base_row + kU <= params->qL && base_col + BK <= params->kL) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } + + // Do softmax (redundantly per half; the row statistics are cheap) + metal::vec new_max; + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + new_max[i] = max_score[i]; + } + + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + sum_score[i] = sum_score[i] * factor[i]; + } + + Stile.template row_reduce(sum_score); + + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O = P @ V, this half of Dv only. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Vtile.load_rows( + V + V_load_off, + int(params->V_strides[2]), + lim_rows_k - ik * kU); + } else { + Vtile.load(V + V_load_off, int(params->V_strides[2])); + } + + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + // Next block + K += BK * int(params->K_strides[2]); + V += BK * int(params->V_strides[2]); + } + + // Normalize output + threadgroup_barrier(mem_flags::mem_none); + + metal::vec rcp; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + rcp[i] = 1.f / sum_score[i]; + } + + Otile.template row_bin_op(rcp); + + if (!align_Q && is_last_q) { + if (lim_rows_q <= 0) + return; + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); + } else { + Otile.store(O, int(params->O_strides[2])); + } + return; + } using otile_t = NAXTile; otile_t Otile; diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index d40df50a45..e0d219a48b 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal @@ -12,6 +12,7 @@ attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float) #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index d6f74bc2a9..f297317a00 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -18,30 +18,92 @@ namespace { void sdpa_full_self_attention_nax( const Stream& s, metal::Device& d, - const array& q, - const array& k, - const array& v, + const array& q_in, + const array& k_in, + const array& v_in, const float scale, - array& o, + array& o_in, bool do_causal_, const std::optional& mask, const std::optional& sinks) { using namespace mlx::steel; + int bd = q_in.shape(-1); int wm = 4; - int wn = 1; + // bd=256 uses the head-dim-split kernel (wn=2): halving the per-simdgroup + // accumulator working set is what recovers tensor-unit throughput at this + // head width. + int wn = bd == 256 ? 2 : 1; - int bd = q.shape(-1); int bq = 64; + // bk=32 everywhere: with the head-dim split at bd=256 this puts the + // per-simdgroup accumulator set (S 2 + O 8 fragments) at exact bd=128 + // parity, which is what restores tensor-unit throughput (42.8 vs + // 47.7 TF at 8192^2). int bk = 32; - int B = q.shape(0); - int H = q.shape(1); - int D = q.shape(3); - int gqa_factor = q.shape(1) / k.shape(1); + int B = q_in.shape(0); + int H = q_in.shape(1); + int D = q_in.shape(3); + int gqa_factor = q_in.shape(1) / k_in.shape(1); + + int qL = q_in.shape(2); + int kL = k_in.shape(2); + + // The causal offset must keep describing the true diagonal even when the + // inputs get padded below. + const int qL_off = kL - qL; + + // For head_dim 256 the unaligned pipelines (align_Q/align_K = false) + // spill registers and slow every block down ~3x, not just the ragged edge + // ones. Pad ragged causal inputs up to the tile sizes and take the + // aligned pipeline instead: with qL_off computed from the true lengths, + // the causal bound c <= r + qL_off never reaches a padded key column for + // any true query row, and the padded query rows are dropped by the copy + // back into o. The pads are zero-filled so the padded value rows cannot + // poison P @ V with NaNs (their probabilities are exactly zero). + std::optional q_pad, k_pad, v_pad, o_pad; + if (bd == 256 && do_causal_ && !mask.has_value() && + ((qL % bq) || (kL % bk))) { + auto& enc = metal::get_command_encoder(s); + auto pad_seq = [&](const array& src, int Lp) { + array dst( + {src.shape(0), src.shape(1), Lp, src.shape(3)}, + src.dtype(), + nullptr, + {}); + array zero(0, src.dtype()); + fill_gpu(zero, dst, s); + array dst_slice(src.shape(), dst.dtype(), nullptr, {}); + dst_slice.copy_shared_buffer( + dst, dst.strides(), dst.flags(), dst_slice.size(), 0); + copy_gpu_inplace(src, dst_slice, CopyType::GeneralGeneral, s); + enc.add_temporary(std::move(zero)); + enc.add_temporary(std::move(dst_slice)); + enc.add_temporary(dst); + return dst; + }; - int qL = q.shape(2); - int kL = k.shape(2); + if (qL % bq) { + int qLp = bq * ((qL + bq - 1) / bq); + q_pad = pad_seq(q_in, qLp); + o_pad = array({B, H, qLp, D}, o_in.dtype(), nullptr, {}); + o_pad->set_data(allocator::malloc(o_pad->nbytes())); + enc.add_temporary(*o_pad); + qL = qLp; + } + if (kL % bk) { + int kLp = bk * ((kL + bk - 1) / bk); + k_pad = pad_seq(k_in, kLp); + v_pad = pad_seq(v_in, kLp); + kL = kLp; + } + } + + const array& q = q_pad ? *q_pad : q_in; + const array& k = k_pad ? *k_pad : k_in; + const array& v = v_pad ? *v_pad : v_in; + array& o = o_pad ? *o_pad : o_in; const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; @@ -131,7 +193,7 @@ void sdpa_full_self_attention_nax( /* int qL_rem = */ (qL - NQ_aligned * bq), /* int kL_rem = */ (kL - NK_aligned * bk), - /* int qL_off = */ (kL - qL), + /* int qL_off = */ qL_off, /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)}, /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)}, @@ -161,6 +223,15 @@ void sdpa_full_self_attention_nax( MTL::Size group_dims = MTL::Size(32, wm, wn); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + + // Drop the padded query rows. + if (o_pad) { + array o_slice(o_in.shape(), o_pad->dtype(), nullptr, {}); + o_slice.copy_shared_buffer( + *o_pad, o_pad->strides(), o_pad->flags(), o_slice.size(), 0); + copy_gpu_inplace(o_slice, o_in, CopyType::GeneralGeneral, s); + compute_encoder.add_temporary(std::move(o_slice)); + } } void sdpa_full_self_attention_metal( @@ -623,8 +694,25 @@ bool ScaledDotProductAttention::use_fallback( (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || query_head_dim == 256)) || (query_head_dim == 192 && value_head_dim == 128); + // head_dim 256 is only instantiated for the NAX kernel, so it must match + // the exact condition under which sdpa_full_self_attention_metal routes + // to it. It is further gated to the shapes where the fused kernel beats + // the unfused graph (measured on M5 Max): large, nearly-square causal + // blocks. On offset rectangles (kL >> qL, i.e. chunked prefill against a + // long cache) the causal block-skip amortizes away and the unfused + // graph's GEMMs win; below ~2k queries there are too few query blocks to + // fill the machine. + const bool takes_nax_full_path = metal::is_nax_available() && + (env::enable_tf32() || q.dtype() != float32); + // The head-dim-split kernel beats the unfused graph on every measured + // causal shape with qL >= 1024, rectangles included (M5 Max); below that + // there are too few query blocks to fill the machine against a long + // cache (512x8192 measures 1.08x of unfused). + const bool sdpa_full_supported_256 = takes_nax_full_path && do_causal && + !has_arr_mask && query_sequence_length >= 1024; const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128); + (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128 || + (query_head_dim == 256 && sdpa_full_supported_256)); const bool sdpa_full_supported_mask = !has_mask || has_arr_mask || (query_sequence_length <= key_sequence_length && do_causal); diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 7bd867084e..3782591498 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -117,6 +117,80 @@ def mlx_primitives_sdpa(q, k, v, scale, mask=None): class TestFastSDPA(mlx_tests.MLXTestCase): + def test_sdpa_full_head_dim_256_tf32_disabled(self): + # float32 with tf32 disabled must NOT be routed to the fused path: + # only the NAX kernel is instantiated for head_dim 256, so a + # predicate mismatch shows up as a kernel-load failure. TF32 is a + # process-wide setting, so exercise it in a subprocess. + import os + import subprocess + import sys + + # qL >= 2048 so that, predicate mismatch aside, the shape would + # otherwise be routed to the fused path. + script = ( + "import mlx.core as mx\n" + "q = mx.random.normal(shape=(1, 4, 2048, 256))\n" + "k = mx.random.normal(shape=(1, 2, 2048, 256))\n" + "v = mx.random.normal(shape=(1, 2, 2048, 256))\n" + "out = mx.fast.scaled_dot_product_attention(" + "q, k, v, scale=0.0625, mask='causal')\n" + "mx.eval(out)\n" + "print('ok')\n" + ) + env = dict(os.environ, MLX_ENABLE_TF32="0") + result = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("ok", result.stdout) + + def test_sdpa_full_head_dim_256(self): + # On NAX devices, large nearly-square causal blocks take the fused + # path, with ragged lengths padded up to the tile size inside the + # dispatch; everything else takes the unfused fallback. All of it + # must be correct. + D = 256 + Nq, Nkv = 8, 2 + scale = D**-0.5 + mx.random.seed(0) + cases = [ + # (qL, kL, mask): unfused fallback shapes + (17, 17, None), + (17, 17, "causal"), + (128, 128, "causal"), + (512, 512, "causal"), + # fused on NAX: aligned square, ragged square (padded), aligned + # rectangle at the routing boundary, ragged rectangle (padded) + (2048, 2048, "causal"), + (2049, 2049, "causal"), + (2048, 2560, "causal"), + (2049, 2560, "causal"), + ] + for dtype in (mx.float32, mx.bfloat16): + for qL, kL, mask in cases: + with self.subTest(dtype=dtype, qL=qL, kL=kL, mask=mask): + q = (5e-1 * mx.random.normal(shape=(1, Nq, qL, D))).astype(dtype) + k = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype(dtype) + v = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype(dtype) + k_rep = mx.repeat(k, Nq // Nkv, axis=1) + v_rep = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, k_rep, v_rep, scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + self.assertEqual(out.shape, ref.shape) + if dtype == mx.float32: + # The fused shapes run through tf32 tensor ops when + # MLX_ENABLE_TF32 is on (the default). + tol = 1e-3 if qL >= 2048 else 1e-4 + else: + tol = 5e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 From 1aee415d61379735ee637494f07f2e4158fde612 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Mon, 13 Jul 2026 13:41:20 -0700 Subject: [PATCH 2/2] Format NAX attention routing predicate --- mlx/backend/metal/scaled_dot_product_attention.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index f297317a00..4e3e12aa2a 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -702,8 +702,8 @@ bool ScaledDotProductAttention::use_fallback( // long cache) the causal block-skip amortizes away and the unfused // graph's GEMMs win; below ~2k queries there are too few query blocks to // fill the machine. - const bool takes_nax_full_path = metal::is_nax_available() && - (env::enable_tf32() || q.dtype() != float32); + const bool takes_nax_full_path = + metal::is_nax_available() && (env::enable_tf32() || q.dtype() != float32); // The head-dim-split kernel beats the unfused graph on every measured // causal shape with qL >= 1024, rectangles included (M5 Max); below that // there are too few query blocks to fill the machine against a long