feat(sdpa): fused shared-prefix cascade decode op, fa_verify head_dim 64/128, return_lse

CHANGELOG.md

@@ -20,6 +20,11 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   `KQ_QMM_SPLITK_NAX`, `KQ_MV_EXT_SB`, `KQ_MV_EXT_NX`, `KQ_MV_EXT_HD`):
   the NAX split-K path lifts the collapsed M9-16 band 65-76%; the rest
   measured flat to negative on M5 and stay off by default.
+- `sdpa_decode_gqa_cascade`: fused shared-prefix batched decode; one KV
+  walk serves the shared prefix for every batch row, private suffixes read
+  per row. 1.6-4.2x vs per-row calls at P 14k-32k, hd128/hd256.
+- `sdpa_fa_verify` head_dim 64/128 tiles; `return_lse` on
+  `sdpa_decode_gqa` and `sdpa_fa_verify`.
 
 ## [0.3.7]
 

bindings.cpp

@@ -199,7 +199,52 @@ NB_MODULE(_ext, m) {
 
   m.def(
       "sdpa_decode_gqa",
-      &mlx_kquant::sdpa_decode_gqa,
+      [](mx::array q,
+         mx::array k,
+         mx::array v,
+         float scale,
+         const std::optional<mx::array>& sinks,
+         int splits,
+         int tile_c,
+         const std::optional<mx::array>& starts,
+         const std::optional<mx::array>& k_scales,
+         const std::optional<mx::array>& k_biases,
+         const std::optional<mx::array>& v_scales,
+         const std::optional<mx::array>& v_biases,
+         bool return_lse,
+         mx::StreamOrDevice s) -> nb::object {
+        if (return_lse) {
+          auto outs = mlx_kquant::sdpa_decode_gqa_lse(
+              std::move(q),
+              std::move(k),
+              std::move(v),
+              scale,
+              sinks,
+              splits,
+              tile_c,
+              starts,
+              k_scales,
+              k_biases,
+              v_scales,
+              v_biases,
+              s);
+          return nb::make_tuple(outs[0], outs[1]);
+        }
+        return nb::cast(mlx_kquant::sdpa_decode_gqa(
+            std::move(q),
+            std::move(k),
+            std::move(v),
+            scale,
+            sinks,
+            splits,
+            tile_c,
+            starts,
+            k_scales,
+            k_biases,
+            v_scales,
+            v_biases,
+            s));
+      },
       "q"_a,
       "k"_a,
       "v"_a,
@@ -213,6 +258,7 @@ NB_MODULE(_ext, m) {
       "v_scales"_a = nb::none(),
       "v_biases"_a = nb::none(),
       nb::kw_only(),
+      "return_lse"_a = false,
       "stream"_a = nb::none(),
       R"(
         Decode/verify GQA attention tuned for long KV caches: the key axis
@@ -244,19 +290,45 @@ NB_MODULE(_ext, m) {
                 kL). Out-of-range values read as an empty row (zero output).
 
         Returns:
-            array: attention output [B, n_q_heads, qL, D].
+            array: attention output [B, n_q_heads, qL, D]. With
+            ``return_lse=True``, a tuple ``(out, lse)`` where lse
+            [B, n_q_heads, qL] float32 is the natural-log softmax
+            normalizer per query row (the merge weight for combining
+            attention over disjoint key regions).
       )");
 
   m.def(
       "sdpa_fa_verify",
-      &mlx_kquant::sdpa_fa_verify,
+      [](mx::array q,
+         mx::array k,
+         mx::array v,
+         float scale,
+         int q_len,
+         int splits,
+         bool return_lse,
+         mx::StreamOrDevice s) -> nb::object {
+        if (return_lse) {
+          auto outs = mlx_kquant::sdpa_fa_verify_lse(
+              std::move(q),
+              std::move(k),
+              std::move(v),
+              scale,
+              q_len,
+              splits,
+              s);
+          return nb::make_tuple(outs[0], outs[1]);
+        }
+        return nb::cast(mlx_kquant::sdpa_fa_verify(
+            std::move(q), std::move(k), std::move(v), scale, q_len, splits, s));
+      },
       "q"_a,
       "k"_a,
       "v"_a,
       "scale"_a,
       "q_len"_a,
       "splits"_a = 0,
       nb::kw_only(),
+      "return_lse"_a = false,
       "stream"_a = nb::none(),
       R"(
         Speculative-verify attention on the GPU matrix units for a GQA-folded
@@ -272,8 +344,8 @@ NB_MODULE(_ext, m) {
 
         Args:
             q (array): folded queries [1, n_kv_heads, G*q_len, D],
-                float16/bfloat16; D = 256 or 512; G*q_len <= 64 at D=256,
-                <= 32 at D=512.
+                float16/bfloat16; D = 64, 128, 256 or 512; G*q_len <= 64
+                except <= 32 at D=512.
             k (array): keys [1, n_kv_heads, kL, D]; head/seq strided is fine
                 (read in place), the head_dim must be contiguous.
             v (array): values [1, n_kv_heads, kL, D].
@@ -284,7 +356,86 @@ NB_MODULE(_ext, m) {
             splits (int): key-axis split count; 0 picks the default.
 
         Returns:
-            array: attention output [1, n_kv_heads, G*q_len, D].
+            array: attention output [1, n_kv_heads, G*q_len, D]. With
+            ``return_lse=True``, a tuple ``(out, lse)`` where lse
+            [1, n_kv_heads, G*q_len] float32 is the natural-log softmax
+            normalizer per folded row (cascade merge weight).
+      )");
+
+  m.def(
+      "sdpa_decode_gqa_cascade",
+      [](mx::array q,
+         mx::array k_shared,
+         mx::array v_shared,
+         mx::array k_priv,
+         mx::array v_priv,
+         float scale,
+         const std::optional<mx::array>& starts,
+         int splits_shared,
+         int splits_priv,
+         int tile_c,
+         bool return_lse,
+         mx::StreamOrDevice s) -> nb::object {
+        auto outs = mlx_kquant::sdpa_decode_gqa_cascade(
+            std::move(q),
+            std::move(k_shared),
+            std::move(v_shared),
+            std::move(k_priv),
+            std::move(v_priv),
+            scale,
+            starts,
+            splits_shared,
+            splits_priv,
+            tile_c,
+            return_lse,
+            s);
+        if (return_lse) {
+          return nb::make_tuple(outs[0], outs[1]);
+        }
+        return nb::cast(outs[0]);
+      },
+      "q"_a,
+      "k_shared"_a,
+      "v_shared"_a,
+      "k_priv"_a,
+      "v_priv"_a,
+      "scale"_a,
+      "starts"_a = nb::none(),
+      "splits_shared"_a = 0,
+      "splits_priv"_a = 0,
+      "tile_c"_a = 0,
+      nb::kw_only(),
+      "return_lse"_a = false,
+      "stream"_a = nb::none(),
+      R"(
+        Fused shared-prefix (cascade) decode attention: every batch row
+        attends one COMMON prefix, stored once, plus its own private
+        suffix. The shared region is walked ONCE for all B*gqa query rows
+        on the matrix-unit row tile; the private region runs per row (with
+        optional left-pad ``starts``); both partial sets fold through a
+        single merge pass. Equivalent to ``sdpa_decode_gqa`` over the
+        concatenated KV, reading the prefix once instead of B times.
+
+        Args:
+            q (array): queries [B, n_q_heads, 1, D], float16/bfloat16;
+                D in {64, 128, 256, 512}; gqa <= 16; B*gqa <= 64 (<= 32 at
+                D=512).
+            k_shared (array): shared prefix keys [1, n_kv_heads, P, D].
+            v_shared (array): shared prefix values [1, n_kv_heads, P, D].
+            k_priv (array): private suffix keys [B, n_kv_heads, Sp, D],
+                Sp >= 1.
+            v_priv (array): private suffix values [B, n_kv_heads, Sp, D].
+            scale (float): query scale (typically 1/sqrt(D)).
+            starts (array, optional): int32 [B] per-row private-region key
+                start offsets (left-padded private suffixes).
+            splits_shared (int): shared-region split count; 0 = default.
+            splits_priv (int): private-region split count; 0 = default.
+            tile_c (int): private-pass staged tile height; 0 picks by
+                head_dim.
+
+        Returns:
+            array: attention output [B, n_q_heads, 1, D]. With
+            ``return_lse=True``, a tuple ``(out, lse)``.
       )");
 
   m.def(

metal/kq_sdpa.metal

@@ -108,6 +108,14 @@ instantiate_kq_sdpa_gqa_p2(float16_t, 64, 32, 4)
       type,                                                            \
       D)
 
+instantiate_kq_sdpa_fa_verify(bfloat16_t, 64, 32)
+instantiate_kq_sdpa_fa_verify(float16_t, 64, 32)
+instantiate_kq_sdpa_fa_verify(bfloat16_t, 64, 64)
+instantiate_kq_sdpa_fa_verify(float16_t, 64, 64)
+instantiate_kq_sdpa_fa_verify(bfloat16_t, 128, 32)
+instantiate_kq_sdpa_fa_verify(float16_t, 128, 32)
+instantiate_kq_sdpa_fa_verify(bfloat16_t, 128, 64)
+instantiate_kq_sdpa_fa_verify(float16_t, 128, 64)
 instantiate_kq_sdpa_fa_verify(bfloat16_t, 256, 32)
 instantiate_kq_sdpa_fa_verify(float16_t, 256, 32)
 instantiate_kq_sdpa_fa_verify(bfloat16_t, 256, 64)

metal/mlx/backend/metal/kernels/kq_sdpa.h

@@ -21,6 +21,12 @@ constant bool gqa_has_starts [[function_constant(4)]];
 // cooperative tile stage dequants into sK/sV -- everything downstream of
 // the stage is unchanged. Compiled out when false.
 constant bool gqa_kv_q8 [[function_constant(5)]];
+constant bool gqa_write_lse [[function_constant(6)]];
+// Cascade merge: a second partials set (the shared-prefix region, written by
+// the fa row-tile pass in its folded [1, Hkv, B*gqa, splits2, D] layout)
+// folds into the same online-softmax reduction as the primary set. Compiled
+// out when false.
+constant bool gqa_cascade [[function_constant(7)]];
 
 template <typename T, int D, int V = D>
 [[kernel]] void kq_sdpa_vector_2pass_1(
@@ -1029,6 +1035,12 @@ template <typename T, int D>
     const device float* sinks [[buffer(3)]],
     device T* out [[buffer(4)]],
     const constant int& n_q_heads [[buffer(5)]],
+    device float* out_lse [[buffer(6)]],
+    const device float* partials2 [[buffer(7)]],
+    const device float* sums2 [[buffer(8)]],
+    const device float* maxs2 [[buffer(9)]],
+    const constant int& cascade_splits [[buffer(10)]],
+    const constant int& cascade_gqa [[buffer(11)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 tpg [[threadgroups_per_grid]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
@@ -1042,12 +1054,28 @@ template <typename T, int D>
   sums += base * gqa_splits;
   maxs += base * gqa_splits;
 
+  // Second set: fa folded layout, row = kv*(B*gqa) + b*gqa + g.
+  int splits2 = 0;
+  if (gqa_cascade) {
+    const int kv = head_idx / cascade_gqa;
+    const int g = head_idx % cascade_gqa;
+    const size_t row = ((size_t)kv * tpg.y + batch_idx) * cascade_gqa + g;
+    splits2 = cascade_splits;
+    partials2 += row * splits2 * D;
+    sums2 += row * splits2;
+    maxs2 += row * splits2;
+  }
+
   threadgroup float ws[128];
+  threadgroup float ws2[128];
 
   float m = Limits<float>::finite_min;
   for (int s = simd_lid; s < gqa_splits; s += 32) {
     m = max(m, maxs[s]);
   }
+  for (int s = simd_lid; s < splits2; s += 32) {
+    m = max(m, maxs2[s]);
+  }
   m = simd_max(m);
   if (gqa_has_sinks) {
     m = max(m, sinks[head_idx]);
@@ -1059,6 +1087,11 @@ template <typename T, int D>
     ws[s] = w;
     denom += w * sums[s];
   }
+  for (int s = simd_lid; s < splits2; s += 32) {
+    const float w = fast::exp(maxs2[s] - m);
+    ws2[s] = w;
+    denom += w * sums2[s];
+  }
   denom = simd_sum(denom);
   if (gqa_has_sinks) {
     denom += fast::exp(sinks[head_idx] - m);
@@ -1072,8 +1105,19 @@ template <typename T, int D>
       acc[e] += w * partials[s * D + e * 32 + simd_lid];
     }
   }
+  for (int s = 0; s < splits2; s++) {
+    const float w = ws2[s];
+    for (short e = 0; e < EPT; e++) {
+      acc[e] += w * partials2[s * D + e * 32 + simd_lid];
+    }
+  }
   out += base * D;
   for (short e = 0; e < EPT; e++) {
     out[e * 32 + simd_lid] = static_cast<T>(denom == 0 ? 0.0f : acc[e] / denom);
   }
+  // Natural-log softmax normalizer (sinks included when present): the
+  // cascade merge weight for combining disjoint key regions.
+  if (gqa_write_lse && simd_lid == 0) {
+    out_lse[base] = denom == 0 ? -INFINITY : (fast::log(denom) + m);
+  }
 }

mlx_kquant/__init__.py

@@ -64,6 +64,7 @@
     rmsnorm_multi3,
     route_shed,
     sdpa_decode_gqa,
+    sdpa_decode_gqa_cascade,
     sdpa_fa_verify,
     sdpa_vector,
     shared_event_create,
@@ -120,6 +121,7 @@
     "rmsnorm_multi3",
     "route_shed",
     "sdpa_decode_gqa",
+    "sdpa_decode_gqa_cascade",
     "sdpa_fa_verify",
     "sdpa_vector",
     "shared_event_create",