Add a fused full-attention path for head_dim 256 on NAX devices

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<AccumType, TQ, TDh>;
+    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<AccumType, kRowsPT> max_score;
+    metal::vec<AccumType, kRowsPT> sum_score{0};
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      max_score[i] = Limits<AccumType>::finite_min;
+    }
+
+    if (has_sinks) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < kRowsPT; ++i) {
+        max_score[i] = M_LOG2E_F * static_cast<AccumType>(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<AccumType, TQ, TK>;
+    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<T, 1, 1> 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<T, 2, 1> 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<AccumType>::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<AccumType>::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<AccumType>::finite_min;
+
+        const int base_row = tid.x * BQ + tm;
+        const int base_col = kb * BK;
+
+        constexpr bool is_bool = is_same_v<MaskType, bool>;
+        using melem_t = typename metal::conditional_t<is_bool, bool, AccumType>;
+        using mtile_t = NAXTile<melem_t, TQ, TK>;
+        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<AccumType, kRowsPT> new_max;
+      metal::vec<AccumType, kRowsPT> factor;
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < kRowsPT; ++i) {
+        new_max[i] = max_score[i];
+      }
+
+      Stile.template row_reduce<MaxOp>(new_max);
+      Stile.template row_bin_op<ExpSubOp>(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<SumOp>(sum_score);
+
+      Otile.template row_bin_op<MulOp>(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<T, 1, 2> 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<AccumType, kRowsPT> rcp;
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < kRowsPT; ++i) {
+      rcp[i] = 1.f / sum_score[i];
+    }
+
+    Otile.template row_bin_op<MulOp>(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<AccumType, TQ, TD>;
   otile_t Otile;
 

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) \