From de7ee4d62a29960b40173fda2b3945871b320ebc Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Mon, 27 Jul 2026 23:08:16 -0700
Subject: [PATCH 1/5] feat(sdpa): sparse page-gather decode walk (fc gqa_paged,
 per-kv-head page list)

---
 bindings.cpp                              |  33 +++++++
 metal/mlx/backend/metal/kernels/kq_sdpa.h |  29 +++++-
 mlx_kquant/__init__.py                    |   2 +
 src/kquant.h                              |  20 +++-
 src/kquant_sdpa.cpp                       | 107 +++++++++++++++++++++-
 tests/test_sdpa.py                        |  73 +++++++++++++++
 6 files changed, 254 insertions(+), 10 deletions(-)

diff --git a/bindings.cpp b/bindings.cpp
index 91b9292..47cf8f3 100644
--- a/bindings.cpp
+++ b/bindings.cpp
@@ -362,6 +362,39 @@ NB_MODULE(_ext, m) {
             normalizer per folded row (cascade merge weight).
       )");
 
+  m.def(
+      "sdpa_decode_gqa_paged",
+      &mlx_kquant::sdpa_decode_gqa_paged,
+      "q"_a,
+      "k"_a,
+      "v"_a,
+      "scale"_a,
+      "pages"_a,
+      "splits"_a = 0,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"(
+        Sparse page-gather decode attention: attend ONLY the key/value
+        pages listed per (batch, kv-head), walking the selected pages
+        through the decode kernel instead of the full cache. The page
+        unit is the head dim's staged tile height: 32 rows at head_dim
+        64/128, 16 at 256, 8 at 512.
+
+        Args:
+            q (array): queries [B, n_q_heads, 1, D], float16/bfloat16.
+            k (array): keys [B, n_kv_heads, S, D] (full cache view).
+            v (array): values [B, n_kv_heads, S, D].
+            scale (float): query scale (typically 1/sqrt(D)).
+            pages (array): int32 [B, n_kv_heads, n_pages] page indices in
+                [0, ceil(S / page_size)); no duplicates. The final partial
+                page is tail-clamped to S automatically.
+            splits (int): key-axis split count; 0 buckets by the SELECTED
+                key count.
+
+        Returns:
+            array: attention output [B, n_q_heads, 1, D].
+      )");
+
   m.def(
       "sdpa_decode_gqa_cascade",
       [](mx::array q,
diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h
index a135afb..6402a25 100644
--- a/metal/mlx/backend/metal/kernels/kq_sdpa.h
+++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h
@@ -27,6 +27,10 @@ constant bool gqa_write_lse [[function_constant(6)]];
 // folds into the same online-softmax reduction as the primary set. Compiled
 // out when false.
 constant bool gqa_cascade [[function_constant(7)]];
+// Page-gather decode: the key walk follows a per-(batch, kv-head) list of
+// selected C-row pages (sparse top-k attention) instead of the contiguous
+// [0, N) axis. Compiled out when false.
+constant bool gqa_paged [[function_constant(8)]];
 
 template <typename T, int D, int V = D>
 [[kernel]] void kq_sdpa_vector_2pass_1(
@@ -232,6 +236,8 @@ template <typename T, int D, int C = 32, int NE = 4, int QPS = 1>
     const constant size_t& ks_seq_stride [[buffer(19)]],
     const constant size_t& vs_head_stride [[buffer(20)]],
     const constant size_t& vs_seq_stride [[buffer(21)]],
+    const device int* pages [[buffer(22)]],
+    const constant int& n_pages [[buffer(23)]],
     uint3 tptg [[threads_per_threadgroup]],
     uint3 tidtg [[thread_position_in_threadgroup]],
     uint3 tid [[threadgroup_position_in_grid]],
@@ -322,7 +328,22 @@ template <typename T, int D, int C = 32, int NE = 4, int QPS = 1>
   const int flat = (tidtg.z * gqa_factor + tidtg.y) * 32 + lane;
   const int n_threads = 32 * gqa_factor * tptg.z;
 
-  for (int kt = kt0; kt < k1; kt += C) {
+  // Tile walk: contiguous chunks of the key axis, or (paged) this
+  // threadgroup's slice of the selected-page list. A page tile's tail
+  // guard is N itself (pages are C-aligned windows of the full cache).
+  int t0 = kt0 / C;
+  int t1 = (k1 + C - 1) / C;
+  const device int* prow = pages;
+  if (gqa_paged) {
+    const int pchunk = (n_pages + gqa_splits - 1) / gqa_splits;
+    t0 = split_idx * pchunk;
+    t1 = min(t0 + pchunk, n_pages);
+    prow = pages + (size_t)(batch_idx * num_kv_heads + kv_head_idx) * n_pages;
+  }
+  const int kend = gqa_paged ? N : k1;
+
+  for (int t = t0; t < t1; t++) {
+    const int kt = gqa_paged ? prow[t] * C : t * C;
     threadgroup_barrier(mem_flags::mem_threadgroup);
     // Cooperative tile load; zero-fill the tail so stale threadgroup data
     // can never reach the accumulators. The q8 path works in uint4 units
@@ -337,7 +358,7 @@ template <typename T, int D, int C = 32, int NE = 4, int QPS = 1>
         const int c16 = u % DU;
         const int kg = kt + row;
         const int i = row * D4 + c16 * 4;
-        if (kg < k1) {
+        if (kg < kend) {
           const uint4 kw =
               ((const device uint4*)(kwbase + (size_t)kg * k_seq_stride))[c16];
           const uint4 vw =
@@ -383,7 +404,7 @@ template <typename T, int D, int C = 32, int NE = 4, int QPS = 1>
         const int row = i / D4;
         const int col = i % D4;
         const int kg = kt + row;
-        if (kg < k1) {
+        if (kg < kend) {
           sK[i] = ((const device T4*)(kbase + (size_t)kg * k_seq_stride))[col];
           sV[i] = ((const device T4*)(vbase + (size_t)kg * v_seq_stride))[col];
         } else {
@@ -419,7 +440,7 @@ template <typename T, int D, int C = 32, int NE = 4, int QPS = 1>
         }
         s[p] = simd_shuffle(s[p], NL * ty);
         const bool valid =
-            kg < k1 && kg <= lim[p] && (!gqa_has_starts || kg >= row_start);
+            kg < kend && kg <= lim[p] && (!gqa_has_starts || kg >= row_start);
         mqk[p][cc] = valid ? s[p] : Limits<float>::finite_min;
         m_tile[p] = max(m_tile[p], mqk[p][cc]);
       }
diff --git a/mlx_kquant/__init__.py b/mlx_kquant/__init__.py
index ccfe9e6..241b02c 100644
--- a/mlx_kquant/__init__.py
+++ b/mlx_kquant/__init__.py
@@ -65,6 +65,7 @@
     route_shed,
     sdpa_decode_gqa,
     sdpa_decode_gqa_cascade,
+    sdpa_decode_gqa_paged,
     sdpa_fa_verify,
     sdpa_vector,
     shared_event_create,
@@ -122,6 +123,7 @@
     "route_shed",
     "sdpa_decode_gqa",
     "sdpa_decode_gqa_cascade",
+    "sdpa_decode_gqa_paged",
     "sdpa_fa_verify",
     "sdpa_vector",
     "shared_event_create",
diff --git a/src/kquant.h b/src/kquant.h
index 41e1394..6519f5a 100644
--- a/src/kquant.h
+++ b/src/kquant.h
@@ -242,6 +242,19 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     bool return_lse = false,
     mx::StreamOrDevice s = {});
 
+// Sparse page-gather decode: attend only the C-row pages listed per
+// (batch, kv-head). pages is int32 [B, n_kv_heads, n_pages] with page
+// indices into the key axis (page size = the head dim's staged tile:
+// 32 at D<=128, 16 at D=256, 8 at D=512). qL == 1 only; fp16/bf16 KV.
+mx::array sdpa_decode_gqa_paged(
+    mx::array q,
+    mx::array k,
+    mx::array v,
+    float scale,
+    mx::array pages,
+    int splits = 0,
+    mx::StreamOrDevice s = {});
+
 // sdpa_fa_verify returning {out, lse}: lse [B, Hkv, n_rows] float32 is the
 // natural-log softmax normalizer per folded row (cascade merge weight).
 std::vector<mx::array> sdpa_fa_verify_lse(
@@ -760,7 +773,8 @@ class KQuantSDPAGQA : public mx::Primitive {
       bool has_sinks,
       bool has_starts,
       bool has_kv_q8 = false,
-      bool return_lse = false)
+      bool return_lse = false,
+      bool paged = false)
       : mx::Primitive(stream),
         scale_(scale),
         splits_(splits),
@@ -768,7 +782,8 @@ class KQuantSDPAGQA : public mx::Primitive {
         has_sinks_(has_sinks),
         has_starts_(has_starts),
         has_kv_q8_(has_kv_q8),
-        return_lse_(return_lse) {}
+        return_lse_(return_lse),
+        paged_(paged) {}
 
   void eval_cpu(
       const std::vector<mx::array>& inputs,
@@ -791,6 +806,7 @@ class KQuantSDPAGQA : public mx::Primitive {
   int tile_c_;
   bool has_sinks_;
   bool has_starts_;
+  bool paged_ = false;
   bool has_kv_q8_;
   bool return_lse_;
 };
diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp
index f8867c3..787abdd 100644
--- a/src/kquant_sdpa.cpp
+++ b/src/kquant_sdpa.cpp
@@ -213,6 +213,7 @@ void KQuantSDPAGQA::eval_gpu(
   const bool sinks = has_sinks_;
   const bool starts = has_starts_;
   const bool kv_q8 = has_kv_q8_;
+  const bool paged = paged_;
   const size_t starts_idx = 3 + (sinks ? 1 : 0);
   const size_t qkv_idx = starts_idx + (starts ? 1 : 0);
   kq_sdpa_check_layout("sdpa_decode_gqa", q, k, v);
@@ -226,10 +227,17 @@ void KQuantSDPAGQA::eval_gpu(
   int gqa_factor = n_q_heads / n_kv_heads;
   // Auto splits: coarse buckets (a per-kL value would mint a new pipeline
   // specialization every decode step). Measured on M5 Max: more splits win as
-  // depth grows; ~512-1024 keys per chunk is the sweet spot.
+  // depth grows; ~512-1024 keys per chunk is the sweet spot. The paged walk
+  // rides the pages operand at the end of the input list and buckets on the
+  // SELECTED key count, not the cache depth.
+  int n_pages = 0;
+  if (paged) {
+    n_pages = static_cast<int>(inputs.back().shape(2));
+  }
   int splits = splits_;
   if (splits == 0) {
-    splits = kL <= 8192 ? 16 : kL <= 24576 ? 32 : kL <= 49152 ? 64 : 128;
+    const int span = paged ? n_pages * tile_c_ : kL;
+    splits = span <= 8192 ? 16 : span <= 24576 ? 32 : span <= 49152 ? 64 : 128;
   }
 
   size_t k_head_stride =
@@ -261,6 +269,7 @@ void KQuantSDPAGQA::eval_gpu(
   bool has_kv_q8 = kv_q8;
   bool has_lse = write_lse;
   bool has_cascade = false;
+  bool has_paged = paged;
   mx::metal::MTLFCList fc = {
       {&splits, MTL::DataType::DataTypeInt, 2},
       {&has_sinks, MTL::DataType::DataTypeBool, 3},
@@ -268,6 +277,7 @@ void KQuantSDPAGQA::eval_gpu(
       {&has_kv_q8, MTL::DataType::DataTypeBool, 5},
       {&has_lse, MTL::DataType::DataTypeBool, 6},
       {&has_cascade, MTL::DataType::DataTypeBool, 7},
+      {&has_paged, MTL::DataType::DataTypeBool, 8},
   };
 
   // Pass 1: one threadgroup per (kv-head, batch, split); the whole GQA group
@@ -278,7 +288,8 @@ void KQuantSDPAGQA::eval_gpu(
     std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) +
         "_c" + std::to_string(tile_c_) + (qL > 1 ? "_p2" : "");
     std::string hash = kname + "_s" + std::to_string(splits) +
-        (has_starts ? "_st1" : "_st0") + (has_kv_q8 ? "_q8" : "");
+        (has_starts ? "_st1" : "_st0") + (has_kv_q8 ? "_q8" : "") +
+        (paged ? "_pg" : "");
     auto kernel = kq_get_kernel(d, kname, hash, fc);
     // Register-heavy pipeline: some GPUs cap it below the dispatch width, and
     // Metal turns an oversized dispatch into silent garbage, not an error.
@@ -328,6 +339,9 @@ void KQuantSDPAGQA::eval_gpu(
     ce.set_bytes(ks_seq_stride, 19);
     ce.set_bytes(vs_head_stride, 20);
     ce.set_bytes(vs_seq_stride, 21);
+    // Page list (dummy when the paged walk is compiled out).
+    ce.set_input_array(paged ? inputs.back() : sums, 22);
+    ce.set_bytes(n_pages, 23);
     MTL::Size group_dims(32, gqa_factor, qL > 1 ? (qL + 1) / 2 : 1);
     MTL::Size grid_dims(n_kv_heads, B, splits);
     ce.dispatch_threadgroups(grid_dims, group_dims);
@@ -579,6 +593,7 @@ void KQuantSDPACascade::eval_gpu(
         {&f, MTL::DataType::DataTypeBool, 3},
         {&has_starts, MTL::DataType::DataTypeBool, 4},
         {&f, MTL::DataType::DataTypeBool, 5},
+        {&f, MTL::DataType::DataTypeBool, 8},
     };
     std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) +
         "_c" + std::to_string(tile_c_);
@@ -620,6 +635,9 @@ void KQuantSDPACascade::eval_gpu(
     ce.set_bytes(zero, 19);
     ce.set_bytes(zero, 20);
     ce.set_bytes(zero, 21);
+    const int pzero = 0;
+    ce.set_input_array(sums1, 22);
+    ce.set_bytes(pzero, 23);
     MTL::Size group_dims(32, gqa_factor, 1);
     MTL::Size grid_dims(n_kv_heads, B, s_pr);
     ce.dispatch_threadgroups(grid_dims, group_dims);
@@ -843,7 +861,8 @@ bool KQuantSDPAGQA::is_equivalent(const mx::Primitive& other) const {
   const auto& o = static_cast<const KQuantSDPAGQA&>(other);
   return scale_ == o.scale_ && splits_ == o.splits_ && tile_c_ == o.tile_c_ &&
       has_sinks_ == o.has_sinks_ && has_starts_ == o.has_starts_ &&
-      has_kv_q8_ == o.has_kv_q8_ && return_lse_ == o.return_lse_;
+      has_kv_q8_ == o.has_kv_q8_ && return_lse_ == o.return_lse_ &&
+      paged_ == o.paged_;
 }
 
 static std::vector<mx::array> sdpa_decode_gqa_impl(
@@ -1048,6 +1067,86 @@ mx::array sdpa_decode_gqa(
       s_)[0];
 }
 
+mx::array sdpa_decode_gqa_paged(
+    mx::array q,
+    mx::array k,
+    mx::array v,
+    float scale,
+    mx::array pages,
+    int splits,
+    mx::StreamOrDevice s_) {
+  auto s = mx::to_stream(s_);
+  const char* op = "[mlx_kquant.sdpa_decode_gqa_paged] ";
+  if (q.ndim() != 4 || k.ndim() != 4 || v.ndim() != 4) {
+    throw std::invalid_argument(std::string(op) + "q/k/v must be 4-D.");
+  }
+  int B = q.shape(0);
+  int n_q_heads = q.shape(1);
+  int qL = q.shape(2);
+  int D = q.shape(3);
+  int n_kv_heads = k.shape(1);
+  if (qL != 1) {
+    throw std::invalid_argument(
+        std::string(op) + "query length must be 1 (decode).");
+  }
+  if (D != 64 && D != 128 && D != 256 && D != 512) {
+    throw std::invalid_argument(
+        std::string(op) + "only head_dim 64/128/256/512 is supported.");
+  }
+  auto dt = q.dtype();
+  if (dt != mx::float16 && dt != mx::bfloat16) {
+    throw std::invalid_argument(
+        std::string(op) + "q must be float16 or bfloat16.");
+  }
+  if (k.dtype() != dt || v.dtype() != dt) {
+    throw std::invalid_argument(
+        std::string(op) + "q, k, v must share a dtype (no quantized KV).");
+  }
+  if (n_kv_heads == 0 || n_q_heads % n_kv_heads != 0) {
+    throw std::invalid_argument(
+        std::string(op) + "n_q_heads must be a multiple of n_kv_heads.");
+  }
+  int gqa_factor = n_q_heads / n_kv_heads;
+  if (gqa_factor > 16) {
+    throw std::invalid_argument(std::string(op) + "gqa_factor must be <= 16.");
+  }
+  if (splits < 0 || splits > 128) {
+    throw std::invalid_argument(
+        std::string(op) + "splits must be in [0, 128].");
+  }
+  // Page unit is the head dim's staged tile height.
+  const int tile_c = D <= 128 ? 32 : D == 256 ? 16 : 8;
+  if (pages.dtype() != mx::int32 || pages.ndim() != 3 || pages.shape(0) != B ||
+      pages.shape(1) != n_kv_heads || pages.shape(2) < 1) {
+    throw std::invalid_argument(
+        std::string(op) +
+        "pages must be int32 [B, n_kv_heads, n_pages] with n_pages >= 1 "
+        "(page indices into the key axis; page size = " +
+        std::to_string(tile_c) + " rows at this head_dim).");
+  }
+
+  auto q_c = mx::contiguous(q, false, s);
+  auto k_c = k.strides().back() == 1 ? k : mx::contiguous(k, false, s);
+  auto v_c = v.strides().back() == 1 ? v : mx::contiguous(v, false, s);
+  auto p_c = mx::contiguous(pages, false, s);
+
+  auto prim = std::make_shared<KQuantSDPAGQA>(
+      s,
+      scale,
+      splits,
+      tile_c,
+      /*has_sinks=*/false,
+      /*has_starts=*/false,
+      /*has_kv_q8=*/false,
+      /*return_lse=*/false,
+      /*paged=*/true);
+  auto out_shape = q.shape();
+  std::vector<mx::array> inputs = {
+      std::move(q_c), std::move(k_c), std::move(v_c), std::move(p_c)};
+  return mx::array(
+      std::move(out_shape), dt, std::move(prim), std::move(inputs));
+}
+
 std::vector<mx::array> sdpa_decode_gqa_lse(
     mx::array q,
     mx::array k,
diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py
index b388734..fe7ac75 100644
--- a/tests/test_sdpa.py
+++ b/tests/test_sdpa.py
@@ -663,3 +663,76 @@ def test_sdpa_cascade_fused_validation():
         kq.sdpa_decode_gqa_cascade(
             q, k_sh, v_sh, k_pr[:, :, :0, :], v_pr[:, :, :0, :], scale
         )  # empty private region
+
+
+@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16])
+@pytest.mark.parametrize("D", [64, 128, 256])
+def test_sdpa_paged_matches_selected_reference(D, dtype):
+    # page-gather decode == f32 attention over exactly the selected pages
+    import numpy as np
+
+    np.random.seed(31)
+    B, Hq, Hkv, S, npages = 2, 16, 8, 4093, 24
+    page = 32 if D <= 128 else 16
+    tot = (S + page - 1) // page
+    scale = 1.0 / (D**0.5)
+    q, k, v = _make(B, Hq, Hkv, 1, S, D, dtype, seed=32, strided=False)
+    pg = np.stack(
+        [
+            np.stack(
+                [
+                    np.sort(np.random.choice(tot, size=npages, replace=False))
+                    for _ in range(Hkv)
+                ]
+            )
+            for _ in range(B)
+        ]
+    ).astype(np.int32)
+    pages = mx.array(pg)
+    got = kq.sdpa_decode_gqa_paged(q, k, v, scale, pages)
+    kr = mx.repeat(k, Hq // Hkv, axis=1).astype(mx.float32)
+    vr = mx.repeat(v, Hq // Hkv, axis=1).astype(mx.float32)
+    sc = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2)
+    mask = np.zeros((B, Hkv, S), dtype=bool)
+    for b in range(B):
+        for h in range(Hkv):
+            for pp in pg[b, h]:
+                mask[b, h, pp * page : min((pp + 1) * page, S)] = True
+    mask = np.repeat(mask, Hq // Hkv, axis=1)[:, :, None, :]
+    bias = mx.array(np.where(mask, 0.0, -np.inf).astype(np.float32))
+    ref = mx.softmax(sc + bias, axis=-1) @ vr
+    _eval_or_skip(got, ref)
+    err = float(mx.abs(got.astype(mx.float32) - ref).max())
+    assert err < 2e-2, f"paged vs selected ref err={err}"
+
+
+def test_sdpa_paged_all_pages_is_dense():
+    # selecting every page must reproduce the dense decode call
+    import numpy as np
+
+    B, Hq, Hkv, D, S = 2, 32, 8, 128, 2048
+    scale = 1.0 / (D**0.5)
+    q, k, v = _make(B, Hq, Hkv, 1, S, D, mx.bfloat16, seed=33, strided=False)
+    tot = S // 32
+    pg = np.broadcast_to(np.arange(tot, dtype=np.int32), (B, Hkv, tot)).copy()
+    got = kq.sdpa_decode_gqa_paged(q, k, v, scale, mx.array(pg))
+    ref = kq.sdpa_decode_gqa(q, k, v, scale)
+    _eval_or_skip(got, ref)
+    rel = _rel(got, ref)
+    assert rel < REL_BOUND[mx.bfloat16], f"all-pages vs dense rel {rel:.3e}"
+
+
+def test_sdpa_paged_validation():
+    B, Hq, Hkv, D, S = 2, 16, 8, 128, 1024
+    scale = 1.0 / (D**0.5)
+    q, k, v = _make(B, Hq, Hkv, 1, S, D, mx.bfloat16, seed=34, strided=False)
+    good = mx.zeros((B, Hkv, 4), dtype=mx.int32)
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_paged(q, k, v, scale, good.astype(mx.float32))
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_paged(
+            q, k, v, scale, mx.zeros((B, Hkv + 1, 4), dtype=mx.int32)
+        )
+    q2 = mx.concatenate([q, q], axis=2)
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_paged(q2, k, v, scale, good)

From 079fb6bd2a7c4dd08108c53183a2b188f38f0c67 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Tue, 28 Jul 2026 00:50:07 -0700
Subject: [PATCH 2/5] feat(sdpa): starts on the paged decode builder
 (left-padded batch rows)

---
 bindings.cpp        |  4 +++-
 src/kquant.h        |  4 ++++
 src/kquant_sdpa.cpp | 18 ++++++++++++++--
 tests/test_sdpa.py  | 50 +++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/bindings.cpp b/bindings.cpp
index 47cf8f3..36f57e7 100644
--- a/bindings.cpp
+++ b/bindings.cpp
@@ -371,6 +371,7 @@ NB_MODULE(_ext, m) {
       "scale"_a,
       "pages"_a,
       "splits"_a = 0,
+      "starts"_a = nb::none(),
       nb::kw_only(),
       "stream"_a = nb::none(),
       R"(
@@ -378,7 +379,8 @@ NB_MODULE(_ext, m) {
         pages listed per (batch, kv-head), walking the selected pages
         through the decode kernel instead of the full cache. The page
         unit is the head dim's staged tile height: 32 rows at head_dim
-        64/128, 16 at 256, 8 at 512.
+        64/128, 16 at 256, 8 at 512. Optional starts (int32 [B])
+        restricts row b to keys [starts[b], N) for left-padded batches.
 
         Args:
             q (array): queries [B, n_q_heads, 1, D], float16/bfloat16.
diff --git a/src/kquant.h b/src/kquant.h
index 6519f5a..d334560 100644
--- a/src/kquant.h
+++ b/src/kquant.h
@@ -246,6 +246,9 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
 // (batch, kv-head). pages is int32 [B, n_kv_heads, n_pages] with page
 // indices into the key axis (page size = the head dim's staged tile:
 // 32 at D<=128, 16 at D=256, 8 at D=512). qL == 1 only; fp16/bf16 KV.
+// Optional `starts` (int32 [B]) restricts row b to keys [starts[b], N)
+// -- left-padded batches; pad positions inside selected pages score
+// -inf, so selecting a partially padded page stays exact.
 mx::array sdpa_decode_gqa_paged(
     mx::array q,
     mx::array k,
@@ -253,6 +256,7 @@ mx::array sdpa_decode_gqa_paged(
     float scale,
     mx::array pages,
     int splits = 0,
+    const std::optional<mx::array>& starts = std::nullopt,
     mx::StreamOrDevice s = {});
 
 // sdpa_fa_verify returning {out, lse}: lse [B, Hkv, n_rows] float32 is the
diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp
index 787abdd..9e4a152 100644
--- a/src/kquant_sdpa.cpp
+++ b/src/kquant_sdpa.cpp
@@ -1074,6 +1074,7 @@ mx::array sdpa_decode_gqa_paged(
     float scale,
     mx::array pages,
     int splits,
+    const std::optional<mx::array>& starts,
     mx::StreamOrDevice s_) {
   auto s = mx::to_stream(s_);
   const char* op = "[mlx_kquant.sdpa_decode_gqa_paged] ";
@@ -1136,13 +1137,26 @@ mx::array sdpa_decode_gqa_paged(
       splits,
       tile_c,
       /*has_sinks=*/false,
-      /*has_starts=*/false,
+      /*has_starts=*/starts.has_value(),
       /*has_kv_q8=*/false,
       /*return_lse=*/false,
       /*paged=*/true);
   auto out_shape = q.shape();
   std::vector<mx::array> inputs = {
-      std::move(q_c), std::move(k_c), std::move(v_c), std::move(p_c)};
+      std::move(q_c), std::move(k_c), std::move(v_c)};
+  if (starts.has_value()) {
+    auto st = *starts;
+    if (st.size() != static_cast<size_t>(B)) {
+      throw std::invalid_argument(
+          std::string(op) + "starts must have one element per batch row.");
+    }
+    if (st.dtype() != mx::int32) {
+      throw std::invalid_argument(std::string(op) + "starts must be int32.");
+    }
+    st = mx::reshape(st, {B}, s);
+    inputs.push_back(mx::contiguous(st, false, s));
+  }
+  inputs.push_back(std::move(p_c));
   return mx::array(
       std::move(out_shape), dt, std::move(prim), std::move(inputs));
 }
diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py
index fe7ac75..9988932 100644
--- a/tests/test_sdpa.py
+++ b/tests/test_sdpa.py
@@ -722,6 +722,48 @@ def test_sdpa_paged_all_pages_is_dense():
     assert rel < REL_BOUND[mx.bfloat16], f"all-pages vs dense rel {rel:.3e}"
 
 
+def test_sdpa_paged_starts():
+    # left-padded rows: pad positions inside selected pages score -inf
+    import numpy as np
+
+    np.random.seed(41)
+    B, Hq, Hkv, D, S, npages = 3, 16, 8, 128, 4096, 20
+    page, scale = 32, 1.0 / (D**0.5)
+    pads = [0, 37, 511]
+    q, k, v = _make(B, Hq, Hkv, 1, S, D, mx.float16, seed=42, strided=False)
+    pg = np.stack(
+        [
+            np.stack(
+                [
+                    np.sort(np.random.choice(S // page, size=npages, replace=False))
+                    for _ in range(Hkv)
+                ]
+            )
+            for _ in range(B)
+        ]
+    ).astype(np.int32)
+    # force the pad-boundary page resident so masking inside it is exercised
+    for b in range(B):
+        pg[b, :, 0] = pads[b] // page
+    starts = mx.array(pads, dtype=mx.int32)
+    got = kq.sdpa_decode_gqa_paged(q, k, v, scale, mx.array(pg), starts=starts)
+    kr = mx.repeat(k, Hq // Hkv, axis=1).astype(mx.float32)
+    vr = mx.repeat(v, Hq // Hkv, axis=1).astype(mx.float32)
+    sc = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2)
+    mask = np.zeros((B, Hkv, S), dtype=bool)
+    for b in range(B):
+        for h in range(Hkv):
+            for pp in pg[b, h]:
+                mask[b, h, pp * page : (pp + 1) * page] = True
+        mask[b, :, : pads[b]] = False
+    mask = np.repeat(mask, Hq // Hkv, axis=1)[:, :, None, :]
+    bias = mx.array(np.where(mask, 0.0, -np.inf).astype(np.float32))
+    ref = mx.softmax(sc + bias, axis=-1) @ vr
+    _eval_or_skip(got, ref)
+    err = float(mx.abs(got.astype(mx.float32) - ref).max())
+    assert err < 2e-2, f"paged+starts vs masked ref err={err}"
+
+
 def test_sdpa_paged_validation():
     B, Hq, Hkv, D, S = 2, 16, 8, 128, 1024
     scale = 1.0 / (D**0.5)
@@ -736,3 +778,11 @@ def test_sdpa_paged_validation():
     q2 = mx.concatenate([q, q], axis=2)
     with pytest.raises(ValueError):
         kq.sdpa_decode_gqa_paged(q2, k, v, scale, good)
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_paged(
+            q, k, v, scale, good, starts=mx.zeros((B + 1,), dtype=mx.int32)
+        )
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_paged(
+            q, k, v, scale, good, starts=mx.zeros((B,), dtype=mx.float32)
+        )

From 0da3e36e8b0fcee5d6f5d61f172491904fc6ba18 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Tue, 28 Jul 2026 09:54:48 -0700
Subject: [PATCH 3/5] feat(sdpa): cascade verify width (qL 1-8, end-aligned
 causal private + unclamped shared, lse [B,Hq,qL])

---
 bindings.cpp                              |  8 ++-
 metal/mlx/backend/metal/kernels/kq_sdpa.h |  6 +-
 src/kquant.h                              |  8 ++-
 src/kquant_sdpa.cpp                       | 46 +++++++++-----
 tests/test_sdpa.py                        | 76 +++++++++++++++++++++++
 5 files changed, 120 insertions(+), 24 deletions(-)

diff --git a/bindings.cpp b/bindings.cpp
index 36f57e7..89e7f4f 100644
--- a/bindings.cpp
+++ b/bindings.cpp
@@ -452,9 +452,11 @@ NB_MODULE(_ext, m) {
         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).
+            q (array): queries [B, n_q_heads, qL, D], float16/bfloat16;
+                qL in [1, 8] (verify width: end-aligned causal on the
+                private suffix, full shared visibility); D in
+                {64, 128, 256, 512}; gqa <= 16; B*gqa*qL <= 64 (<= 32 at
+                D=512); gqa*ceil(qL/2) <= 32 at qL > 1.
             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],
diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h
index 6402a25..962e7f5 100644
--- a/metal/mlx/backend/metal/kernels/kq_sdpa.h
+++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h
@@ -1075,12 +1075,14 @@ 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.
+  // Second set: fa folded layout with the query axis innermost,
+  // row = ((kv*B + b)*gqa + g)*qL + t (tpg.z = qL; decode width -> t = 0).
   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;
+    const size_t row =
+        (((size_t)kv * tpg.y + batch_idx) * cascade_gqa + g) * tpg.z + tid.z;
     splits2 = cascade_splits;
     partials2 += row * splits2 * D;
     sums2 += row * splits2;
diff --git a/src/kquant.h b/src/kquant.h
index d334560..c51c4fe 100644
--- a/src/kquant.h
+++ b/src/kquant.h
@@ -225,9 +225,11 @@ mx::array sdpa_fa_verify(
 // query rows, the decode pass covers the private region per row, and both
 // partial sets fold through one merge pass (the splits machinery is the LSE
 // merge). Equivalent to sdpa_decode_gqa over the concatenated KV, reading
-// the prefix once instead of B times. q [B, Hq, 1, D]; B*gqa <= 64 (32 at
-// head_dim 512); gqa <= 16; Sp >= 1. Returns {out} or {out, lse} with
-// return_lse. Metal-only.
+// the prefix once instead of B times. q [B, Hq, qL, D], qL in [1, 8]
+// (verify width: per-row end-aligned causal on the private suffix, full
+// visibility of the shared prefix); B*gqa*qL <= 64 (32 at head_dim 512);
+// gqa <= 16; gqa*ceil(qL/2) <= 32 at qL > 1; Sp >= qL. Returns {out} or
+// {out, lse} (lse [B, Hq, qL]) with return_lse. Metal-only.
 std::vector<mx::array> sdpa_decode_gqa_cascade(
     mx::array q,
     mx::array k_shared,
diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp
index 9e4a152..5c489fe 100644
--- a/src/kquant_sdpa.cpp
+++ b/src/kquant_sdpa.cpp
@@ -548,10 +548,10 @@ void KQuantSDPACascade::eval_gpu(
   int D = q.shape(3);
   int n_kv_heads = k_sh.shape(1);
   int gqa_factor = n_q_heads / n_kv_heads;
-  int n_rows = B * gqa_factor;
   int P = k_sh.shape(2);
   int Sp = k_pr.shape(2);
-  int qL = 1;
+  int qL = q.shape(2);
+  int n_rows = B * gqa_factor * qL;
   float scale = scale_;
 
   int s_sh = splits_shared_;
@@ -596,11 +596,12 @@ void KQuantSDPACascade::eval_gpu(
         {&f, MTL::DataType::DataTypeBool, 8},
     };
     std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) +
-        "_c" + std::to_string(tile_c_);
+        "_c" + std::to_string(tile_c_) + (qL > 1 ? "_p2" : "");
     std::string hash = kname + "_s" + std::to_string(s_pr) +
         (has_starts ? "_st1" : "_st0") + "_casc";
     auto kernel = kq_get_kernel(d, kname, hash, fc);
-    const size_t tg = size_t(32) * gqa_factor;
+    const size_t tg =
+        size_t(32) * gqa_factor * (qL > 1 ? size_t((qL + 1) / 2) : 1);
     if (tg > kernel->maxTotalThreadsPerThreadgroup()) {
       throw std::runtime_error(
           "[mlx_kquant.sdpa_decode_gqa_cascade] threadgroup of " +
@@ -638,7 +639,7 @@ void KQuantSDPACascade::eval_gpu(
     const int pzero = 0;
     ce.set_input_array(sums1, 22);
     ce.set_bytes(pzero, 23);
-    MTL::Size group_dims(32, gqa_factor, 1);
+    MTL::Size group_dims(32, gqa_factor, qL > 1 ? (qL + 1) / 2 : 1);
     MTL::Size grid_dims(n_kv_heads, B, s_pr);
     ce.dispatch_threadgroups(grid_dims, group_dims);
   }
@@ -681,7 +682,8 @@ void KQuantSDPACascade::eval_gpu(
     ce.set_bytes(v_head_stride, 9);
     ce.set_bytes(v_seq_stride, 10);
     ce.set_bytes(scale, 11);
-    ce.set_bytes(qL, 12);
+    const int q_len_shared = 1; // unclamped: verify rows see the whole prefix
+    ce.set_bytes(q_len_shared, 12);
     ce.set_bytes(n_rows, 13);
     MTL::Size group_dims(32, tg / 32, 1);
     MTL::Size grid_dims(n_kv_heads, 1, s_sh);
@@ -1262,8 +1264,9 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     throw std::invalid_argument(
         std::string(op) + "head_dim must be 64, 128, 256 or 512.");
   }
-  if (q.shape(2) != 1) {
-    throw std::invalid_argument(std::string(op) + "q_len must be 1.");
+  int qL = q.shape(2);
+  if (qL < 1 || qL > 8) {
+    throw std::invalid_argument(std::string(op) + "q_len must be in [1, 8].");
   }
   if (k_shared.shape(0) != 1 || v_shared.shape(0) != 1 ||
       v_shared.shape(1) != n_kv_heads ||
@@ -1291,15 +1294,20 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
   if (gqa_factor > 16) {
     throw std::invalid_argument(std::string(op) + "gqa factor must be <= 16.");
   }
-  int n_rows = B * gqa_factor;
+  int n_rows = B * gqa_factor * qL;
   int max_rows = D == 512 ? 32 : 64;
   if (n_rows > max_rows) {
     throw std::invalid_argument(
-        std::string(op) + "B * gqa must be <= " + std::to_string(max_rows) +
-        " at head_dim " + std::to_string(D) + ".");
+        std::string(op) + "B * gqa * q_len must be <= " +
+        std::to_string(max_rows) + " at head_dim " + std::to_string(D) + ".");
+  }
+  if (qL > 1 && gqa_factor * ((qL + 1) / 2) > 32) {
+    throw std::invalid_argument(
+        std::string(op) + "gqa * ceil(q_len/2) must be <= 32.");
   }
+  const int tile_default = D <= 128 ? 32 : D == 256 ? 16 : 8;
   if (tile_c == 0) {
-    tile_c = D <= 128 ? 32 : D == 256 ? 16 : 8;
+    tile_c = tile_default;
   }
   const bool tile_ok = (D <= 128 && (tile_c == 32 || tile_c == 16)) ||
       (D == 256 && (tile_c == 16 || tile_c == 8)) || (D == 512 && tile_c == 8);
@@ -1307,6 +1315,11 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     throw std::invalid_argument(
         std::string(op) + "tile_c not instantiated for this head_dim.");
   }
+  if (qL > 1 && tile_c != tile_default) {
+    // the verify-width (_p2) private kernel exists at the default tile only
+    throw std::invalid_argument(
+        std::string(op) + "q_len > 1 requires the default tile_c.");
+  }
   if (splits_shared < 0 || splits_shared > 128 || splits_priv < 0 ||
       splits_priv > 128) {
     throw std::invalid_argument(
@@ -1314,12 +1327,13 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
   }
 
   auto q_c = mx::contiguous(q, false, s);
-  // kv-head-major fold for the shared row-tile pass.
+  // kv-head-major fold for the shared row-tile pass, query axis innermost:
+  // row = (b*gqa + g)*qL + t.
   auto q_folded = mx::contiguous(
       mx::reshape(
           mx::transpose(
-              mx::reshape(q_c, {B, n_kv_heads, gqa_factor, D}, s),
-              {1, 0, 2, 3},
+              mx::reshape(q_c, {B, n_kv_heads, gqa_factor, qL, D}, s),
+              {1, 0, 2, 3, 4},
               s),
           {1, n_kv_heads, n_rows, D},
           s),
@@ -1359,7 +1373,7 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
       return_lse);
   auto out_shape = q.shape();
   if (return_lse) {
-    mx::Shape lse_shape = {B, n_q_heads, 1};
+    mx::Shape lse_shape = {B, n_q_heads, qL};
     return mx::array::make_arrays(
         {std::move(out_shape), std::move(lse_shape)},
         {dt, mx::float32},
diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py
index 9988932..256eaa0 100644
--- a/tests/test_sdpa.py
+++ b/tests/test_sdpa.py
@@ -663,6 +663,82 @@ def test_sdpa_cascade_fused_validation():
         kq.sdpa_decode_gqa_cascade(
             q, k_sh, v_sh, k_pr[:, :, :0, :], v_pr[:, :, :0, :], scale
         )  # empty private region
+    q9 = mx.concatenate([q] * 9, axis=2)
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_cascade(q9, k_sh, v_sh, k_pr, v_pr, scale)  # qL > 8
+    # over the folded-row cap: B2 * gqa8 * qL5 = 80 > 64
+    qw, k_w, v_w = _make(B, 64, Hkv, 5, 64, D, mx.bfloat16, seed=29, strided=False)
+    _, ksw, vsw = _make(1, 64, Hkv, 1, 512, D, mx.bfloat16, seed=30, strided=False)
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_cascade(qw, ksw, vsw, k_w, v_w, scale)
+
+
+def _cascade_verify_ref(q, k, v, pads, scale, qL):
+    # f32 masked reference: per-row left pad + end-aligned causal block
+    B, Hq, _, D = q.shape
+    Hkv, L = k.shape[1], k.shape[2]
+    kr = mx.repeat(k, Hq // Hkv, axis=1).astype(mx.float32)
+    vr = mx.repeat(v, Hq // Hkv, axis=1).astype(mx.float32)
+    s = (q.astype(mx.float32) * scale) @ kr.swapaxes(-1, -2)
+    pos = mx.arange(L)[None, None, None, :]
+    end = (L - qL) + mx.arange(qL)[None, None, :, None]
+    pad = mx.array(pads)[:, None, None, None]
+    keep = (pos >= pad) & (pos <= end)
+    s = mx.where(keep, s, mx.array(-mx.inf))
+    return mx.softmax(s, axis=-1) @ vr
+
+
+@pytest.mark.parametrize(
+    "B,Hq,Hkv,D,qL,pads",
+    [
+        (2, 4, 2, 256, 8, [0, 64]),  # gemma-31b assistant geometry
+        (2, 12, 2, 256, 5, [0, 64]),  # qwen nextn geometry at cap-2
+        (2, 8, 4, 512, 3, [0, 16]),  # hd512 d-split walk
+        (4, 16, 8, 128, 8, [0, 3, 511, 64]),  # 64 folded rows exactly
+        (1, 8, 8, 64, 5, [0]),
+    ],
+)
+def test_sdpa_cascade_fused_verify_width(B, Hq, Hkv, D, qL, pads):
+    # qL > 1: end-aligned causal on the private suffix, full shared
+    # visibility, per-row starts honored
+    P, sp = 2048, 96
+    scale = 1.0 / (D**0.5)
+    mx.random.seed(41)
+    L = max(pads) + P + sp
+    kb = mx.random.normal((B, Hkv, L, D)).astype(mx.float16)
+    vb = mx.random.normal((B, Hkv, L, D)).astype(mx.float16)
+    pk = kb[0:1, :, pads[0] : pads[0] + P]
+    pv = vb[0:1, :, pads[0] : pads[0] + P]
+    rk, rv = [], []
+    for b in range(B):
+        rk.append(
+            mx.concatenate(
+                [kb[b : b + 1, :, : pads[b]], pk, kb[b : b + 1, :, pads[b] + P :]],
+                axis=2,
+            )
+        )
+        rv.append(
+            mx.concatenate(
+                [vb[b : b + 1, :, : pads[b]], pv, vb[b : b + 1, :, pads[b] + P :]],
+                axis=2,
+            )
+        )
+    k = mx.concatenate(rk, axis=0)
+    v = mx.concatenate(rv, axis=0)
+    q = mx.random.normal((B, Hq, qL, D)).astype(mx.float16)
+    c0 = min(pads) + P
+    starts = None
+    if any(pads):
+        starts = mx.array([p + P - c0 for p in pads], dtype=mx.int32)
+    got, lse = kq.sdpa_decode_gqa_cascade(
+        q, pk, pv, k[:, :, c0:], v[:, :, c0:], scale, starts=starts, return_lse=True
+    )
+    ref = _cascade_verify_ref(q, k, v, pads, scale, qL)
+    _eval_or_skip(got, ref)
+    assert lse.shape == (B, Hq, qL)
+    rel = _rel(got, ref)
+    print(f"  [cascade] verify D={D} qL={qL} B={B}: rel={rel:.3e}")
+    assert rel < REL_BOUND[mx.float16], f"cascade verify rel {rel:.3e}"
 
 
 @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16])

From 8af1cd0858b5633c521eed656784ecfb77eecd7f Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Tue, 28 Jul 2026 11:39:55 -0700
Subject: [PATCH 4/5] feat(sdpa): q8-kv operands for the cascade op (both
 passes dequant on tile stage)

---
 bindings.cpp                              |  29 ++++
 metal/mlx/backend/metal/kernels/kq_sdpa.h | 129 +++++++++++++---
 src/kquant.h                              |  24 ++-
 src/kquant_sdpa.cpp                       | 171 ++++++++++++++++++++--
 tests/test_sdpa.py                        | 109 ++++++++++++++
 5 files changed, 427 insertions(+), 35 deletions(-)

diff --git a/bindings.cpp b/bindings.cpp
index 89e7f4f..7479374 100644
--- a/bindings.cpp
+++ b/bindings.cpp
@@ -410,6 +410,14 @@ NB_MODULE(_ext, m) {
          int splits_priv,
          int tile_c,
          bool return_lse,
+         const std::optional<mx::array>& k_shared_scales,
+         const std::optional<mx::array>& k_shared_biases,
+         const std::optional<mx::array>& v_shared_scales,
+         const std::optional<mx::array>& v_shared_biases,
+         const std::optional<mx::array>& k_priv_scales,
+         const std::optional<mx::array>& k_priv_biases,
+         const std::optional<mx::array>& v_priv_scales,
+         const std::optional<mx::array>& v_priv_biases,
          mx::StreamOrDevice s) -> nb::object {
         auto outs = mlx_kquant::sdpa_decode_gqa_cascade(
             std::move(q),
@@ -423,6 +431,14 @@ NB_MODULE(_ext, m) {
             splits_priv,
             tile_c,
             return_lse,
+            k_shared_scales,
+            k_shared_biases,
+            v_shared_scales,
+            v_shared_biases,
+            k_priv_scales,
+            k_priv_biases,
+            v_priv_scales,
+            v_priv_biases,
             s);
         if (return_lse) {
           return nb::make_tuple(outs[0], outs[1]);
@@ -441,6 +457,14 @@ NB_MODULE(_ext, m) {
       "tile_c"_a = 0,
       nb::kw_only(),
       "return_lse"_a = false,
+      "k_shared_scales"_a = nb::none(),
+      "k_shared_biases"_a = nb::none(),
+      "v_shared_scales"_a = nb::none(),
+      "v_shared_biases"_a = nb::none(),
+      "k_priv_scales"_a = nb::none(),
+      "k_priv_biases"_a = nb::none(),
+      "v_priv_scales"_a = nb::none(),
+      "v_priv_biases"_a = nb::none(),
       "stream"_a = nb::none(),
       R"(
         Fused shared-prefix (cascade) decode attention: every batch row
@@ -469,6 +493,11 @@ NB_MODULE(_ext, m) {
             splits_priv (int): private-region split count; 0 = default.
             tile_c (int): private-pass staged tile height; 0 picks by
                 head_dim.
+            k_shared_scales ... v_priv_biases (array, optional): quantized
+                KV (mlx affine wire, bits 8 / group 64). Pass all eight and
+                both k/v slabs bind as packed uint32 words ([.., S, D/4])
+                with scales/biases [.., S, D/64] in q's dtype; dequant
+                happens at tile stage. Not supported at D=512.
 
         Returns:
             array: attention output [B, n_q_heads, 1, D]. With
diff --git a/metal/mlx/backend/metal/kernels/kq_sdpa.h b/metal/mlx/backend/metal/kernels/kq_sdpa.h
index 962e7f5..9955e20 100644
--- a/metal/mlx/backend/metal/kernels/kq_sdpa.h
+++ b/metal/mlx/backend/metal/kernels/kq_sdpa.h
@@ -566,6 +566,54 @@ METAL_FUNC void kq_fa_stage_rows(
   }
 }
 
+// q8 variant: rows arrive as packed uint32 affine wire (bits 8, group 64)
+// with per-group scale/bias; dequant lands in the staged tile so the MMA
+// consumers are unchanged. Works in uint4 units (16 codes) -- one
+// vectorized wire load and one scale/bias pair per unit, never straddling
+// a 64-element group (same staging shape as the decode kernel's q8 path).
+template <typename T, int BK, int D, int LDS, int NT>
+METAL_FUNC void kq_fa_stage_rows_q8(
+    threadgroup T* dst,
+    const device uint32_t* src,
+    const device T* scales,
+    const device T* biases,
+    size_t seq_stride_w,
+    size_t sb_seq_stride,
+    int rows_valid,
+    int flat_tid) {
+  using T4 = metal::vec<T, 4>;
+  constexpr int DU = D / 16; // uint4 units per row
+  constexpr int LDS4 = LDS / 4;
+  threadgroup T4* dst4 = (threadgroup T4*)dst;
+  for (int u = flat_tid; u < BK * DU; u += NT) {
+    const int r = u / DU;
+    const int c16 = u - r * DU;
+    const int i = r * LDS4 + c16 * 4;
+    if (r < rows_valid) {
+      const uint4 w =
+          ((const device uint4*)(src + (size_t)r * seq_stride_w))[c16];
+      const size_t g = (c16 * 16) / 64;
+      const float sc = float(scales[(size_t)r * sb_seq_stride + g]);
+      const float bi = float(biases[(size_t)r * sb_seq_stride + g]);
+#pragma unroll
+      for (short j = 0; j < 4; j++) {
+        const uint32_t x = j == 0 ? w.x : j == 1 ? w.y : j == 2 ? w.z : w.w;
+        const float4 q4 = float4(
+            float(x & 0xff),
+            float((x >> 8) & 0xff),
+            float((x >> 16) & 0xff),
+            float(x >> 24));
+        dst4[i + j] = T4(q4 * sc + bi);
+      }
+    } else {
+#pragma unroll
+      for (short j = 0; j < 4; j++) {
+        dst4[i + j] = T4(T(0));
+      }
+    }
+  }
+}
+
 // Simdgroup-matrix (steel MMA) speculative-verify attention, pass 1. The
 // caller folds the GQA group into the query rows -- q [B, Hq, qL, D] becomes
 // [B, Hkv, G*qL, D] with kv-major heads -- so the kernel sees an MHA problem
@@ -605,6 +653,14 @@ template <typename T, int D, int BQ = 32>
     const constant float& scale [[buffer(11)]],
     const constant int& q_len [[buffer(12)]],
     const constant int& n_rows [[buffer(13)]],
+    const device T* k_scales [[buffer(14)]],
+    const device T* k_biases [[buffer(15)]],
+    const device T* v_scales [[buffer(16)]],
+    const device T* v_biases [[buffer(17)]],
+    const constant size_t& ks_head_stride [[buffer(18)]],
+    const constant size_t& ks_seq_stride [[buffer(19)]],
+    const constant size_t& vs_head_stride [[buffer(20)]],
+    const constant size_t& vs_seq_stride [[buffer(21)]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]],
     uint3 tid [[threadgroup_position_in_grid]],
@@ -634,12 +690,25 @@ template <typename T, int D, int BQ = 32>
   const int k0 = split_idx * chunk;
   const int k1 = min(k0 + chunk, N);
 
-  const device T* kbase = keys +
-      (size_t)(batch_idx * num_kv_heads + kv_head_idx) * k_head_stride +
-      (size_t)k0 * k_seq_stride;
-  const device T* vbase = values +
-      (size_t)(batch_idx * num_kv_heads + kv_head_idx) * v_head_stride +
-      (size_t)k0 * v_seq_stride;
+  // With gqa_kv_q8 the k/v buffers hold packed uint32 wire and the strides
+  // arrive in WORDS (D/4 per row); otherwise they hold T elements.
+  const size_t kv_hb = (size_t)(batch_idx * num_kv_heads + kv_head_idx);
+  const device T* kbase =
+      keys + kv_hb * k_head_stride + (size_t)k0 * k_seq_stride;
+  const device T* vbase =
+      values + kv_hb * v_head_stride + (size_t)k0 * v_seq_stride;
+  const device uint32_t* kwbase = (const device uint32_t*)keys +
+      kv_hb * k_head_stride + (size_t)k0 * k_seq_stride;
+  const device uint32_t* vwbase = (const device uint32_t*)values +
+      kv_hb * v_head_stride + (size_t)k0 * v_seq_stride;
+  const device T* ksb =
+      k_scales + kv_hb * ks_head_stride + (size_t)k0 * ks_seq_stride;
+  const device T* kbb =
+      k_biases + kv_hb * ks_head_stride + (size_t)k0 * ks_seq_stride;
+  const device T* vsb =
+      v_scales + kv_hb * vs_head_stride + (size_t)k0 * vs_seq_stride;
+  const device T* vbb =
+      v_biases + kv_hb * vs_head_stride + (size_t)k0 * vs_seq_stride;
 
   // Fragment coordinates: this thread owns row (row0 + sm) and the column
   // pair at sn of every 8x8 fragment.
@@ -676,12 +745,24 @@ template <typename T, int D, int BQ = 32>
   for (int kt = k0; kt < k1; kt += BK) {
     const int krem = min(k1 - kt, BK);
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    kq_fa_stage_rows<T, BK, D, LDS, kNWarps * 32>(
-        KV_smem,
-        kbase + (size_t)(kt - k0) * k_seq_stride,
-        k_seq_stride,
-        krem,
-        flat_tid);
+    if (gqa_kv_q8) {
+      kq_fa_stage_rows_q8<T, BK, D, LDS, kNWarps * 32>(
+          KV_smem,
+          kwbase + (size_t)(kt - k0) * k_seq_stride,
+          ksb + (size_t)(kt - k0) * ks_seq_stride,
+          kbb + (size_t)(kt - k0) * ks_seq_stride,
+          k_seq_stride,
+          ks_seq_stride,
+          krem,
+          flat_tid);
+    } else {
+      kq_fa_stage_rows<T, BK, D, LDS, kNWarps * 32>(
+          KV_smem,
+          kbase + (size_t)(kt - k0) * k_seq_stride,
+          k_seq_stride,
+          krem,
+          flat_tid);
+    }
     Stile.clear();
     threadgroup_barrier(mem_flags::mem_threadgroup);
 
@@ -723,12 +804,24 @@ template <typename T, int D, int BQ = 32>
     }
 
     threadgroup_barrier(mem_flags::mem_threadgroup);
-    kq_fa_stage_rows<T, BK, D, LDS, kNWarps * 32>(
-        KV_smem,
-        vbase + (size_t)(kt - k0) * v_seq_stride,
-        v_seq_stride,
-        krem,
-        flat_tid);
+    if (gqa_kv_q8) {
+      kq_fa_stage_rows_q8<T, BK, D, LDS, kNWarps * 32>(
+          KV_smem,
+          vwbase + (size_t)(kt - k0) * v_seq_stride,
+          vsb + (size_t)(kt - k0) * vs_seq_stride,
+          vbb + (size_t)(kt - k0) * vs_seq_stride,
+          v_seq_stride,
+          vs_seq_stride,
+          krem,
+          flat_tid);
+    } else {
+      kq_fa_stage_rows<T, BK, D, LDS, kNWarps * 32>(
+          KV_smem,
+          vbase + (size_t)(kt - k0) * v_seq_stride,
+          v_seq_stride,
+          krem,
+          flat_tid);
+    }
 
     // Online softmax on this thread's row (registers only, overlapping the
     // V load). A row with no valid key yet keeps max at finite_min and
diff --git a/src/kquant.h b/src/kquant.h
index c51c4fe..1315050 100644
--- a/src/kquant.h
+++ b/src/kquant.h
@@ -229,7 +229,14 @@ mx::array sdpa_fa_verify(
 // (verify width: per-row end-aligned causal on the private suffix, full
 // visibility of the shared prefix); B*gqa*qL <= 64 (32 at head_dim 512);
 // gqa <= 16; gqa*ceil(qL/2) <= 32 at qL > 1; Sp >= qL. Returns {out} or
-// {out, lse} (lse [B, Hq, qL]) with return_lse. Metal-only.
+// {out, lse} (lse [B, Hq, qL]) with return_lse.
+//
+// Quantized KV (mlx affine wire, bits 8 / group 64): pass all eight of
+// k/v x scales/biases x shared/priv and both slabs bind as packed uint32
+// words ([.., S, D/4]) with per-group scale/bias arrays ([.., S, D/64],
+// q's dtype); dequant happens at tile stage, downstream math unchanged.
+// All-eight-or-none; head_dim 512 not supported with quantized KV.
+// Metal-only.
 std::vector<mx::array> sdpa_decode_gqa_cascade(
     mx::array q,
     mx::array k_shared,
@@ -242,6 +249,14 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     int splits_priv = 0,
     int tile_c = 0,
     bool return_lse = false,
+    const std::optional<mx::array>& k_shared_scales = std::nullopt,
+    const std::optional<mx::array>& k_shared_biases = std::nullopt,
+    const std::optional<mx::array>& v_shared_scales = std::nullopt,
+    const std::optional<mx::array>& v_shared_biases = std::nullopt,
+    const std::optional<mx::array>& k_priv_scales = std::nullopt,
+    const std::optional<mx::array>& k_priv_biases = std::nullopt,
+    const std::optional<mx::array>& v_priv_scales = std::nullopt,
+    const std::optional<mx::array>& v_priv_biases = std::nullopt,
     mx::StreamOrDevice s = {});
 
 // Sparse page-gather decode: attend only the C-row pages listed per
@@ -865,14 +880,16 @@ class KQuantSDPACascade : public mx::Primitive {
       int splits_priv,
       int tile_c,
       bool has_starts,
-      bool return_lse)
+      bool return_lse,
+      bool has_kv_q8 = false)
       : mx::Primitive(stream),
         scale_(scale),
         splits_shared_(splits_shared),
         splits_priv_(splits_priv),
         tile_c_(tile_c),
         has_starts_(has_starts),
-        return_lse_(return_lse) {}
+        return_lse_(return_lse),
+        has_kv_q8_(has_kv_q8) {}
 
   void eval_cpu(
       const std::vector<mx::array>& inputs,
@@ -896,6 +913,7 @@ class KQuantSDPACascade : public mx::Primitive {
   int tile_c_;
   bool has_starts_;
   bool return_lse_;
+  bool has_kv_q8_;
 };
 
 // Fused MoE GLU gather (see moe_glu_gather). Inference-only.
diff --git a/src/kquant_sdpa.cpp b/src/kquant_sdpa.cpp
index 5c489fe..ae5e7ff 100644
--- a/src/kquant_sdpa.cpp
+++ b/src/kquant_sdpa.cpp
@@ -442,9 +442,11 @@ void KQuantSDPAFAVerify::eval_gpu(
   bool has_sinks = false;
   bool has_lse = write_lse;
   bool has_cascade = false;
+  bool no_q8 = false;
   mx::metal::MTLFCList fc = {
       {&splits, MTL::DataType::DataTypeInt, 2},
       {&has_sinks, MTL::DataType::DataTypeBool, 3},
+      {&no_q8, MTL::DataType::DataTypeBool, 5},
       {&has_lse, MTL::DataType::DataTypeBool, 6},
       {&has_cascade, MTL::DataType::DataTypeBool, 7},
   };
@@ -483,6 +485,15 @@ void KQuantSDPAFAVerify::eval_gpu(
     ce.set_bytes(scale, 11);
     ce.set_bytes(q_len, 12);
     ce.set_bytes(n_rows, 13);
+    // q8 operand slots: compiled out (fc 5 false); bind dummies.
+    size_t zero = 0;
+    for (int i = 0; i < 4; i++) {
+      ce.set_input_array(sums, 14 + i);
+    }
+    ce.set_bytes(zero, 18);
+    ce.set_bytes(zero, 19);
+    ce.set_bytes(zero, 20);
+    ce.set_bytes(zero, 21);
     MTL::Size group_dims(32, tg / 32, 1);
     MTL::Size grid_dims(n_kv_heads, B, splits);
     ce.dispatch_threadgroups(grid_dims, group_dims);
@@ -540,6 +551,10 @@ void KQuantSDPACascade::eval_gpu(
   const auto& k_pr = inputs[4];
   const auto& v_pr = inputs[5];
   const bool starts = has_starts_;
+  const bool kv_q8 = has_kv_q8_;
+  // Optional trailing inputs: [starts], then the eight q8 scale/bias
+  // arrays (shared k/v pairs first, then private).
+  const int q8_base = 6 + int(starts);
   kq_sdpa_check_layout("sdpa_decode_gqa_cascade", q, k_pr, v_pr);
   kq_sdpa_check_layout("sdpa_decode_gqa_cascade", qf, k_sh, v_sh);
 
@@ -588,17 +603,18 @@ void KQuantSDPACascade::eval_gpu(
   {
     bool f = false;
     bool has_starts = starts;
+    bool q8 = kv_q8;
     mx::metal::MTLFCList fc = {
         {&s_pr, MTL::DataType::DataTypeInt, 2},
         {&f, MTL::DataType::DataTypeBool, 3},
         {&has_starts, MTL::DataType::DataTypeBool, 4},
-        {&f, MTL::DataType::DataTypeBool, 5},
+        {&q8, MTL::DataType::DataTypeBool, 5},
         {&f, MTL::DataType::DataTypeBool, 8},
     };
     std::string kname = "kq_sdpa_gqa_2pass_1_" + ts + "_" + std::to_string(D) +
         "_c" + std::to_string(tile_c_) + (qL > 1 ? "_p2" : "");
     std::string hash = kname + "_s" + std::to_string(s_pr) +
-        (has_starts ? "_st1" : "_st0") + "_casc";
+        (has_starts ? "_st1" : "_st0") + (q8 ? "_q8" : "") + "_casc";
     auto kernel = kq_get_kernel(d, kname, hash, fc);
     const size_t tg =
         size_t(32) * gqa_factor * (qL > 1 ? size_t((qL + 1) / 2) : 1);
@@ -629,13 +645,34 @@ void KQuantSDPACascade::eval_gpu(
     ce.set_bytes(qL, 12);
     ce.set_input_array(starts ? inputs[6] : sums1, 13);
     size_t zero = 0;
-    for (int i = 0; i < 4; i++) {
-      ce.set_input_array(sums1, 14 + i);
+    if (kv_q8) {
+      const auto& ksc = inputs[q8_base + 4];
+      const auto& kbi = inputs[q8_base + 5];
+      const auto& vsc = inputs[q8_base + 6];
+      const auto& vbi = inputs[q8_base + 7];
+      ce.set_input_array(ksc, 14);
+      ce.set_input_array(kbi, 15);
+      ce.set_input_array(vsc, 16);
+      ce.set_input_array(vbi, 17);
+      size_t ks_head = static_cast<size_t>(
+          ksc.shape(1) == 1 ? ksc.strides(0) : ksc.strides(1));
+      size_t ks_seq = static_cast<size_t>(ksc.strides(2));
+      size_t vs_head = static_cast<size_t>(
+          vsc.shape(1) == 1 ? vsc.strides(0) : vsc.strides(1));
+      size_t vs_seq = static_cast<size_t>(vsc.strides(2));
+      ce.set_bytes(ks_head, 18);
+      ce.set_bytes(ks_seq, 19);
+      ce.set_bytes(vs_head, 20);
+      ce.set_bytes(vs_seq, 21);
+    } else {
+      for (int i = 0; i < 4; i++) {
+        ce.set_input_array(sums1, 14 + i);
+      }
+      ce.set_bytes(zero, 18);
+      ce.set_bytes(zero, 19);
+      ce.set_bytes(zero, 20);
+      ce.set_bytes(zero, 21);
     }
-    ce.set_bytes(zero, 18);
-    ce.set_bytes(zero, 19);
-    ce.set_bytes(zero, 20);
-    ce.set_bytes(zero, 21);
     const int pzero = 0;
     ce.set_input_array(sums1, 22);
     ce.set_bytes(pzero, 23);
@@ -648,14 +685,17 @@ void KQuantSDPACascade::eval_gpu(
   // serves all B*gqa folded rows.
   {
     bool f = false;
+    bool q8 = kv_q8;
     mx::metal::MTLFCList fc = {
         {&s_sh, MTL::DataType::DataTypeInt, 2},
         {&f, MTL::DataType::DataTypeBool, 3},
+        {&q8, MTL::DataType::DataTypeBool, 5},
     };
     const int bq = n_rows <= 32 ? 32 : 64;
     std::string kname = "kq_sdpa_fa_verify_2pass_1_" + ts + "_" +
         std::to_string(D) + "_bq" + std::to_string(bq);
-    std::string hash = kname + "_s" + std::to_string(s_sh) + "_casc";
+    std::string hash =
+        kname + "_s" + std::to_string(s_sh) + (q8 ? "_q8" : "") + "_casc";
     auto kernel = kq_get_kernel(d, kname, hash, fc);
     const size_t tg = D == 512 ? 256 : (bq / 8) * 32;
     if (tg > kernel->maxTotalThreadsPerThreadgroup()) {
@@ -685,6 +725,35 @@ void KQuantSDPACascade::eval_gpu(
     const int q_len_shared = 1; // unclamped: verify rows see the whole prefix
     ce.set_bytes(q_len_shared, 12);
     ce.set_bytes(n_rows, 13);
+    size_t zero = 0;
+    if (kv_q8) {
+      const auto& ksc = inputs[q8_base + 0];
+      const auto& kbi = inputs[q8_base + 1];
+      const auto& vsc = inputs[q8_base + 2];
+      const auto& vbi = inputs[q8_base + 3];
+      ce.set_input_array(ksc, 14);
+      ce.set_input_array(kbi, 15);
+      ce.set_input_array(vsc, 16);
+      ce.set_input_array(vbi, 17);
+      size_t ks_head = static_cast<size_t>(
+          ksc.shape(1) == 1 ? ksc.strides(0) : ksc.strides(1));
+      size_t ks_seq = static_cast<size_t>(ksc.strides(2));
+      size_t vs_head = static_cast<size_t>(
+          vsc.shape(1) == 1 ? vsc.strides(0) : vsc.strides(1));
+      size_t vs_seq = static_cast<size_t>(vsc.strides(2));
+      ce.set_bytes(ks_head, 18);
+      ce.set_bytes(ks_seq, 19);
+      ce.set_bytes(vs_head, 20);
+      ce.set_bytes(vs_seq, 21);
+    } else {
+      for (int i = 0; i < 4; i++) {
+        ce.set_input_array(sums2, 14 + i);
+      }
+      ce.set_bytes(zero, 18);
+      ce.set_bytes(zero, 19);
+      ce.set_bytes(zero, 20);
+      ce.set_bytes(zero, 21);
+    }
     MTL::Size group_dims(32, tg / 32, 1);
     MTL::Size grid_dims(n_kv_heads, 1, s_sh);
     ce.dispatch_threadgroups(grid_dims, group_dims);
@@ -1221,7 +1290,8 @@ bool KQuantSDPACascade::is_equivalent(const mx::Primitive& other) const {
   const auto& o = static_cast<const KQuantSDPACascade&>(other);
   return scale_ == o.scale_ && splits_shared_ == o.splits_shared_ &&
       splits_priv_ == o.splits_priv_ && tile_c_ == o.tile_c_ &&
-      has_starts_ == o.has_starts_ && return_lse_ == o.return_lse_;
+      has_starts_ == o.has_starts_ && return_lse_ == o.return_lse_ &&
+      has_kv_q8_ == o.has_kv_q8_;
 }
 
 std::vector<mx::array> sdpa_decode_gqa_cascade(
@@ -1236,10 +1306,31 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     int splits_priv,
     int tile_c,
     bool return_lse,
+    const std::optional<mx::array>& k_shared_scales,
+    const std::optional<mx::array>& k_shared_biases,
+    const std::optional<mx::array>& v_shared_scales,
+    const std::optional<mx::array>& v_shared_biases,
+    const std::optional<mx::array>& k_priv_scales,
+    const std::optional<mx::array>& k_priv_biases,
+    const std::optional<mx::array>& v_priv_scales,
+    const std::optional<mx::array>& v_priv_biases,
     mx::StreamOrDevice s_) {
   auto s = mx::to_stream(s_);
   const char* op = "[mlx_kquant.sdpa_decode_gqa_cascade] ";
 
+  const int n_q8 = int(k_shared_scales.has_value()) +
+      int(k_shared_biases.has_value()) + int(v_shared_scales.has_value()) +
+      int(v_shared_biases.has_value()) + int(k_priv_scales.has_value()) +
+      int(k_priv_biases.has_value()) + int(v_priv_scales.has_value()) +
+      int(v_priv_biases.has_value());
+  const bool kv_q8 = n_q8 == 8;
+  if (n_q8 != 0 && n_q8 != 8) {
+    throw std::invalid_argument(
+        std::string(op) +
+        "quantized KV needs all eight scale/bias arrays (shared and "
+        "private, k and v).");
+  }
+
   if (q.ndim() != 4 || k_shared.ndim() != 4 || v_shared.ndim() != 4 ||
       k_priv.ndim() != 4 || v_priv.ndim() != 4) {
     throw std::invalid_argument(
@@ -1254,10 +1345,48 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     throw std::invalid_argument(
         std::string(op) + "q must be float16 or bfloat16.");
   }
-  for (const auto& a : {k_shared, v_shared, k_priv, v_priv}) {
-    if (a.dtype() != dt || a.shape(3) != D) {
+  if (kv_q8) {
+    if (D == 512) {
       throw std::invalid_argument(
-          std::string(op) + "k/v must share q's dtype and head_dim.");
+          std::string(op) + "quantized KV is not supported at head_dim 512.");
+    }
+    // mlx affine wire, bits 8 / group 64: packed uint32 words, one
+    // scale/bias per 64-element group in q's dtype.
+    for (const auto& a : {k_shared, v_shared, k_priv, v_priv}) {
+      if (a.dtype() != mx::uint32 || a.shape(3) != D / 4) {
+        throw std::invalid_argument(
+            std::string(op) +
+            "quantized k/v must be uint32 wire with last "
+            "dim head_dim / 4 (bits 8).");
+      }
+    }
+    const mx::array* sb[8] = {
+        &*k_shared_scales,
+        &*k_shared_biases,
+        &*v_shared_scales,
+        &*v_shared_biases,
+        &*k_priv_scales,
+        &*k_priv_biases,
+        &*v_priv_scales,
+        &*v_priv_biases};
+    for (int i = 0; i < 8; i++) {
+      const auto& ref = i < 4 ? k_shared : k_priv;
+      if (sb[i]->dtype() != dt || sb[i]->ndim() != 4 ||
+          sb[i]->shape(0) != ref.shape(0) || sb[i]->shape(1) != n_kv_heads ||
+          sb[i]->shape(2) != ref.shape(2) || sb[i]->shape(3) != D / 64) {
+        throw std::invalid_argument(
+            std::string(op) +
+            "quantized KV scales/biases must be "
+            "[B, n_kv_heads, S, head_dim / 64] in q's dtype (group 64), "
+            "matching their region.");
+      }
+    }
+  } else {
+    for (const auto& a : {k_shared, v_shared, k_priv, v_priv}) {
+      if (a.dtype() != dt || a.shape(3) != D) {
+        throw std::invalid_argument(
+            std::string(op) + "k/v must share q's dtype and head_dim.");
+      }
     }
   }
   if (D != 64 && D != 128 && D != 256 && D != 512) {
@@ -1362,6 +1491,19 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
     st = mx::reshape(st, {B}, s);
     inputs.push_back(mx::contiguous(st, false, s));
   }
+  if (kv_q8) {
+    for (const mx::array* a :
+         {&*k_shared_scales,
+          &*k_shared_biases,
+          &*v_shared_scales,
+          &*v_shared_biases,
+          &*k_priv_scales,
+          &*k_priv_biases,
+          &*v_priv_scales,
+          &*v_priv_biases}) {
+      inputs.push_back(contig_kv(*a));
+    }
+  }
 
   auto prim = std::make_shared<KQuantSDPACascade>(
       s,
@@ -1370,7 +1512,8 @@ std::vector<mx::array> sdpa_decode_gqa_cascade(
       splits_priv,
       tile_c,
       starts.has_value(),
-      return_lse);
+      return_lse,
+      kv_q8);
   auto out_shape = q.shape();
   if (return_lse) {
     mx::Shape lse_shape = {B, n_q_heads, qL};
diff --git a/tests/test_sdpa.py b/tests/test_sdpa.py
index 256eaa0..cb13dd6 100644
--- a/tests/test_sdpa.py
+++ b/tests/test_sdpa.py
@@ -741,6 +741,115 @@ def test_sdpa_cascade_fused_verify_width(B, Hq, Hkv, D, qL, pads):
     assert rel < REL_BOUND[mx.float16], f"cascade verify rel {rel:.3e}"
 
 
+def _q8(a):
+    # mlx affine wire the batch quantized caches produce (group 64, bits 8)
+    return mx.quantize(a, group_size=64, bits=8)
+
+
+@pytest.mark.parametrize(
+    "B,Hq,Hkv,D,qL",
+    [
+        (4, 32, 8, 128, 1),  # plain batch decode
+        (4, 16, 8, 256, 1),  # hd256 decode
+        (2, 12, 2, 256, 5),  # qwen verify geometry (60 rows)
+        (2, 8, 2, 128, 8),  # 64 folded rows exactly
+    ],
+)
+def test_sdpa_cascade_fused_kv_q8(B, Hq, Hkv, D, qL):
+    # q8 operands == the fp16 cascade run on the dequantized arrays,
+    # bit-exact: both stage the same T values, the math after the stage
+    # is identical
+    P, Sp = 2047, 193 + qL
+    scale = 1.0 / (D**0.5)
+    _, k_sh, v_sh = _make(1, Hq, Hkv, 1, P, D, mx.float16, seed=51, strided=False)
+    q, k_pr, v_pr = _make(B, Hq, Hkv, qL, Sp, D, mx.float16, seed=52, strided=False)
+    starts = mx.array([(7 * b) % 64 for b in range(B)], dtype=mx.int32)
+
+    ksh_w, ksh_s, ksh_b = _q8(k_sh)
+    vsh_w, vsh_s, vsh_b = _q8(v_sh)
+    kpr_w, kpr_s, kpr_b = _q8(k_pr)
+    vpr_w, vpr_s, vpr_b = _q8(v_pr)
+
+    got = kq.sdpa_decode_gqa_cascade(
+        q,
+        ksh_w,
+        vsh_w,
+        kpr_w,
+        vpr_w,
+        scale,
+        starts=starts,
+        k_shared_scales=ksh_s,
+        k_shared_biases=ksh_b,
+        v_shared_scales=vsh_s,
+        v_shared_biases=vsh_b,
+        k_priv_scales=kpr_s,
+        k_priv_biases=kpr_b,
+        v_priv_scales=vpr_s,
+        v_priv_biases=vpr_b,
+    )
+
+    def dq(w, s, b):
+        return mx.dequantize(w, s, b, group_size=64, bits=8)
+
+    ref = kq.sdpa_decode_gqa_cascade(
+        q,
+        dq(ksh_w, ksh_s, ksh_b).astype(mx.float16),
+        dq(vsh_w, vsh_s, vsh_b).astype(mx.float16),
+        dq(kpr_w, kpr_s, kpr_b).astype(mx.float16),
+        dq(vpr_w, vpr_s, vpr_b).astype(mx.float16),
+        scale,
+        starts=starts,
+    )
+    _eval_or_skip(got, ref)
+    diff = float(mx.abs(got - ref).max())
+    print(f"  [cascade] kv_q8 D={D} qL={qL}: max|d|={diff:.3e}")
+    assert diff == 0.0, f"cascade kv_q8 not bit-exact: {diff:.3e}"
+
+
+def test_sdpa_cascade_fused_kv_q8_validation():
+    B, Hq, Hkv, D = 2, 16, 8, 128
+    scale = 1.0 / (D**0.5)
+    _, k_sh, v_sh = _make(1, Hq, Hkv, 1, 512, D, mx.float16, seed=53, strided=False)
+    q, k_pr, v_pr = _make(B, Hq, Hkv, 1, 64, D, mx.float16, seed=54, strided=False)
+    ksh_w, ksh_s, ksh_b = _q8(k_sh)
+    with pytest.raises(ValueError):
+        # partial q8 set (scales without biases)
+        kq.sdpa_decode_gqa_cascade(
+            q, ksh_w, v_sh, k_pr, v_pr, scale, k_shared_scales=ksh_s
+        )
+    # D=512 rejects q8
+    _, k5, v5 = _make(1, 8, 4, 1, 512, 512, mx.float16, seed=55, strided=False)
+    q5, kp5, vp5 = _make(2, 8, 4, 1, 64, 512, mx.float16, seed=56, strided=False)
+    args = {}
+    for name, arr in (
+        ("k_shared", k5),
+        ("v_shared", v5),
+        ("k_priv", kp5),
+        ("v_priv", vp5),
+    ):
+        w, s, b = _q8(arr)
+        args[name] = w
+        args[name + "_scales"] = s
+        args[name + "_biases"] = b
+    with pytest.raises(ValueError):
+        kq.sdpa_decode_gqa_cascade(
+            q5,
+            args["k_shared"],
+            args["v_shared"],
+            args["k_priv"],
+            args["v_priv"],
+            1.0 / (512**0.5),
+            k_shared_scales=args["k_shared_scales"],
+            k_shared_biases=args["k_shared_biases"],
+            v_shared_scales=args["v_shared_scales"],
+            v_shared_biases=args["v_shared_biases"],
+            k_priv_scales=args["k_priv_scales"],
+            k_priv_biases=args["k_priv_biases"],
+            v_priv_scales=args["v_priv_scales"],
+            v_priv_biases=args["v_priv_biases"],
+        )
+
+
 @pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float16])
 @pytest.mark.parametrize("D", [64, 128, 256])
 def test_sdpa_paged_matches_selected_reference(D, dtype):

From 2853e77783ff6197e3db37a039be8e5866568750 Mon Sep 17 00:00:00 2001
From: Asher Feldman <59994+asher@users.noreply.github.com>
Date: Wed, 29 Jul 2026 15:23:52 -0700
Subject: [PATCH 5/5] changelog: cascade-kv8 entries under unreleased

---
 CHANGELOG.md | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c44db33..b30f912 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,14 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   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`.
+- `sdpa_decode_gqa_paged`: page-gather decode over per-kv-head page lists
+  for top-k sparse attention, with `starts` for left-padded batch rows.
+- Verify width (qL 1-8) on the cascade op: end-aligned causal over each
+  row's private slab with full shared-prefix visibility; `lse` gains the
+  qL axis.
+- q8 KV operands (bits 8, group 64) on the cascade op, dequantized on the
+  staged tiles in both passes; bit-exact vs the fp16 cascade on
+  dequantized arrays (head_dim 512 declines).
 
 ## [0.3.7]