Add paged attention kernel for Spyre

kernels/paged_attn/conversion-notes.md

@@ -0,0 +1,56 @@
+# paged_attn conversion notes
+
+Paged attention over a KV cache: queries `Q (B, Lq, H, D)` attend to the cache
+rows `K/V (CACHE, H, D)` named, per request, by `SLOTS (B, Lk)`. Output is
+`(B, H, Lq, D)`. Scores use the standard `1/sqrt(D)` softmax scale, split as
+`scale = 1/sqrt(sqrt(D))` applied to both q and k (`scale*scale = 1/sqrt(D)`).
+
+Two descriptor variants are generated, with identical semantics and signature —
+they differ **only** in how the data-dependent K/V gather is expressed. Both are
+**descriptors only**: there is no raw pointer arithmetic anywhere, including the
+gather. The gather is the data-dependent step (the cache row is chosen at
+runtime by `SLOTS`), and it is expressed with `descriptor_gather`, which lowers
+to the Spyre indirect-access tile (`ktdp.construct_indirect_access_tile`).
+
+## tensor_descriptor.py — base `descriptor_gather`
+
+- Source: `paged_attn.py` (the base draft).
+- K/V are viewed **2-D** as `(CACHE, H*D)`; the descriptor `block_shape` is
+  `[BLK_B*KV_BLOCK, BLK_H*D]`.
+- The `(BLK_B, KV_BLOCK)` slot tile is flattened to a **1-D** batch-major row
+  vector, and `descriptor_gather(k_desc, rows, h_start*D)` pulls those rows and
+  the `BLK_H*D` columns at head offset `h_start*D`. The `(rows, columns)` result
+  is `.reshape`d to `(BLK_B, KV_BLOCK, BLK_H, D)`.
+- This uses the base `descriptor_gather` (1-D index, 2-D src).
+
+## spyre_aware.py — extended any-rank `descriptor_gather`
+
+- Source: `paged_attn_ext.py`.
+- Same algorithm, but `descriptor_gather` is **extended to accept an index and a
+  src descriptor of any rank**. So K/V stay **3-D** `(CACHE, H, D)` with
+  `block_shape=[1, BLK_H, D]`, the 2-D slot tile `(BLK_B, KV_BLOCK)` is passed
+  straight in as the index, and `h_start` is the second index — yielding
+  `(BLK_B, KV_BLOCK, BLK_H, D)` directly, with **no reshape** of the index or of
+  the gathered result. This maps the gather onto the Spyre indirect-access tile
+  more directly than the flatten/reshape of the base form.
+- The representation of the physical layout needs to be revisited.
+
+## Shared structure (both variants)
+
+- 4-D batched online softmax: the B loop steps by `BLK_B` and the H loop by
+  `BLK_H`; the `BLK_B` batches and `BLK_H` heads in each block are vectorized so
+  a single gather serves them all and `tl.dot` batches over `(BLK_B, BLK_H)`.
+- Q, SLOTS, Out use plain `make_tensor_descriptor` loads/stores (static offsets
+  from the loop counters), so no tail masks are needed there.
+- `grid=(1,)`: one program walks all `(B, Lq, H)` work via the explicit loops,
+  so the output is partition-independent and the grid fits 32 cores.
+- KV-page loop bound is `tl.cdiv(Lk, KV_BLOCK)` (runtime-arg agnostic; the base
+  draft's bare `Tk = Lk // kv_block_size` constexpr is replaced).
+- 16-byte last-dim rule: every descriptor's last dim is `BLOCK_D` (`D*2` bytes),
+  `BLK_H*D` (`*2`), or `KV_BLOCK` (`*4`) — all >= 16 bytes for the chosen sizes.
+
+## original.py — reference
+
+A simpler hand-written reference (one program per `(request, head, query
+block)`, single-head gather per KV page) that the two batched variants are
+checked against. Also descriptors only; uses the base `descriptor_gather`.

kernels/paged_attn/lower.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+"""KTIR lowering driver for the paged_attn kernel variants.
+
+Consumed by ``scripts/gen_ktir.py``, which lowers each entry in ``VARIANTS`` to
+``kernels/paged_attn/<variant>.ktir``.
+
+``VARIANTS`` maps a **variant name** (which is also the source ``.py`` module
+name and the output ``.ktir`` stem) to the four things the round-trip lowering
+needs:
+
+    KERNEL      : the @triton.jit function to lower
+    SIGNATURE   : dict[str, str]  arg name -> Triton type ("*fp16", "i32", ...)
+    CONSTEXPRS  : dict[str, value] for every constexpr arg
+    GRID        : optional list, forwarded to SpyreOptions.grid
+
+Two variants share the same signature and shapes; they differ only in how the
+data-dependent K/V gather is expressed:
+  - ``tensor_descriptor``: base descriptor_gather (cache viewed 2-D, 1-D rows).
+  - ``spyre_aware``: extended any-rank descriptor_gather (cache 3-D, 2-D index).
+"""
+
+from kernels.paged_attn.tensor_descriptor import _paged_attn_kernel_NHD_td
+from kernels.paged_attn.spyre_aware import _paged_attn_kernel_NHD_sa
+
+# Concrete shapes match tests/ktir/test_paged_attn.py: B=2, H=4, Lq=Lk=16,
+# D=64, cache of 256 slots, KV_BLOCK=16, BLK_B=2, BLK_H=4. The problem dims are
+# constexprs (descriptor shapes are built from them), so they are baked into the
+# KTIR; only the buffers and the scalar scale stay runtime args.
+_SIGNATURE = {
+    "Q": "*fp16",
+    "K": "*fp16",
+    "V": "*fp16",
+    "SLOTS": "*i32",
+    "Out": "*fp16",
+    "scale": "fp32",
+    "B": "i32",
+    "H": "i32",
+    "Lq": "i32",
+    "Lk": "i32",
+    "CACHE": "i32",
+    "KV_BLOCK": "i32",
+    "BLOCK_Q": "i32",
+    "BLOCK_D": "i32",
+    "BLK_B": "i32",
+    "BLK_H": "i32",
+}
+
+_CONSTEXPRS = {
+    "B": 2,
+    "H": 4,
+    "Lq": 16,
+    "Lk": 16,
+    "CACHE": 256,
+    "KV_BLOCK": 16,
+    "BLOCK_Q": 16,
+    "BLOCK_D": 64,
+    "BLK_B": 2,
+    "BLK_H": 4,
+}
+
+VARIANTS = {
+    "tensor_descriptor": {
+        "KERNEL": _paged_attn_kernel_NHD_td,
+        "SIGNATURE": _SIGNATURE,
+        "CONSTEXPRS": _CONSTEXPRS,
+        # Single program walks all (B, Lq, H) work via the explicit loops; fits 32 cores.
+        "GRID": [1],
+    },
+    "spyre_aware": {
+        "KERNEL": _paged_attn_kernel_NHD_sa,
+        "SIGNATURE": _SIGNATURE,
+        "CONSTEXPRS": _CONSTEXPRS,
+        "GRID": [1],
+    },
+}

kernels/paged_attn/spyre_aware.ktir

@@ -0,0 +1,100 @@
+// Generated from kernels/paged_attn/spyre_aware.py by scripts/gen_ktir.py
+// Source: kernels/paged_attn/lower.py -> VARIANTS['spyre_aware'].
+// DO NOT EDIT BY HAND — regenerate with:
+//     .venv/bin/python scripts/gen_ktir.py paged_attn:spyre_aware
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+#set = affine_set<(d0, d1, d2, d3) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 15 >= 0, d2 >= 0, -d2 + 3 >= 0, d3 >= 0, -d3 + 63 >= 0)>
+#set1 = affine_set<(d0, d1, d2) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 3 >= 0, d2 >= 0, -d2 + 63 >= 0)>
+#set2 = affine_set<(d0, d1) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 15 >= 0)>
+#set3 = affine_set<(d0, d1, d2, d3) : (d0 >= 0, -d0 + 1 >= 0, d1 >= 0, -d1 + 3 >= 0, d2 >= 0, -d2 + 15 >= 0, d3 >= 0, -d3 + 63 >= 0)>
+module {
+  func.func @_paged_attn_kernel_NHD_sa(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: f32) attributes {grid = [1]} {
+    %c0 = arith.constant 0 : index
+    %acc = arith.constant dense<[2, 4, 16, 64]> : tensor<4xindex>
+    %acc_0 = arith.constant dense<[8, 16, 16]> : tensor<3xindex>
+    %l_i = arith.constant 0.000000e+00 : f32
+    %block_max = arith.constant 0xFF800000 : f32
+    %scores = arith.constant dense<[2, 4, 16, 16]> : tensor<4xindex>
+    %scores_1 = arith.constant dense<[8, 64, 16]> : tensor<3xindex>
+    %cst = arith.constant dense<[8, 16, 64]> : tensor<3xindex>
+    %acc_2 = arith.constant dense<0.000000e+00> : tensor<2x4x16x1xf32>
+    %l_i_3 = arith.constant dense<0.000000e+00> : tensor<2x4x16xf32>
+    %cst_4 = arith.constant dense<0.000000e+00> : tensor<8x16x64xf32>
+    %cst_5 = arith.constant dense<0.000000e+00> : tensor<8x16x16xf32>
+    %cst_6 = arith.constant dense<0xFF800000> : tensor<2x4x16xf32>
+
+    %q_desc = ktdp.construct_memory_view %arg0, sizes: [2, 16, 4, 64], strides: [4096, 256, 64, 1] {coordinate_set = #set, memory_space = #ktdp.spyre_memory_space<HBM>} : memref<2x16x4x64xf16>
+    %k_desc = ktdp.construct_memory_view %arg1, sizes: [256, 4, 64], strides: [256, 64, 1] {coordinate_set = #set1, memory_space = #ktdp.spyre_memory_space<HBM>} : memref<256x4x64xf16>
+    %v_desc = ktdp.construct_memory_view %arg2, sizes: [256, 4, 64], strides: [256, 64, 1] {coordinate_set = #set1, memory_space = #ktdp.spyre_memory_space<HBM>} : memref<256x4x64xf16>
+    %s_desc = ktdp.construct_memory_view %arg3, sizes: [2, 16], strides: [16, 1] {coordinate_set = #set2, memory_space = #ktdp.spyre_memory_space<HBM>} : memref<2x16xsi32>
+    %o_desc = ktdp.construct_memory_view %arg4, sizes: [2, 4, 16, 64], strides: [4096, 1024, 64, 1] {coordinate_set = #set3, memory_space = #ktdp.spyre_memory_space<HBM>} : memref<2x4x16x64xf16>
+
+    %q = ktdp.construct_access_tile %q_desc[%c0, %c0, %c0, %c0] {access_tile_order = #map, access_tile_set = #set} : memref<2x16x4x64xf16> -> !ktdp.access_tile<2x16x4x64xindex>
+    %q_7 = ktdp.load %q : <2x16x4x64xindex> -> tensor<2x16x4x64xf16>
+    %q_8 = tensor.empty() : tensor<2x4x16x64xf16>
+    %q_9 = linalg.transpose ins(%q_7 : tensor<2x16x4x64xf16>) outs(%q_8 : tensor<2x4x16x64xf16>) permutation = [0, 2, 1, 3]
+    %q_10 = arith.extf %q_9 : tensor<2x4x16x64xf16> to tensor<2x4x16x64xf32>
+    %q_11 = tensor.empty() : tensor<2x4x16x64xf32>
+    %q_12 = linalg.fill ins(%arg5 : f32) outs(%q_11 : tensor<2x4x16x64xf32>) -> tensor<2x4x16x64xf32>
+    %q_13 = arith.mulf %q_10, %q_12 : tensor<2x4x16x64xf32>
+    %q_14 = arith.truncf %q_13 : tensor<2x4x16x64xf32> to tensor<2x4x16x64xf16>
+
+    %k_g = ktdp.construct_indirect_access_tile intermediate_variables(%arg6, %arg7, %arg8, %arg9) %k_desc[ind(%s_desc[%c0 + %arg6, %c0 + %arg7]), (%c0 + %arg8), (%arg9)] {variables_space_order = #map, variables_space_set = #set} : memref<256x4x64xf16>, memref<2x16xsi32> -> !ktdp.access_tile<2x16x4x64xindex>
+    %k_g_15 = ktdp.load %k_g : <2x16x4x64xindex> -> tensor<2x16x4x64xf16>
+
+    %v_g = ktdp.construct_indirect_access_tile intermediate_variables(%arg6, %arg7, %arg8, %arg9) %v_desc[ind(%s_desc[%c0 + %arg6, %c0 + %arg7]), (%c0 + %arg8), (%arg9)] {variables_space_order = #map, variables_space_set = #set} : memref<256x4x64xf16>, memref<2x16xsi32> -> !ktdp.access_tile<2x16x4x64xindex>
+    %v_g_16 = ktdp.load %v_g : <2x16x4x64xindex> -> tensor<2x16x4x64xf16>
+    %k_g_17 = arith.extf %k_g_15 : tensor<2x16x4x64xf16> to tensor<2x16x4x64xf32>
+    %k_g_18 = tensor.empty() : tensor<2x16x4x64xf32>
+    %k_g_19 = linalg.fill ins(%arg5 : f32) outs(%k_g_18 : tensor<2x16x4x64xf32>) -> tensor<2x16x4x64xf32>
+    %k_g_20 = arith.mulf %k_g_17, %k_g_19 : tensor<2x16x4x64xf32>
+    %k_g_21 = arith.truncf %k_g_20 : tensor<2x16x4x64xf32> to tensor<2x16x4x64xf16>
+    %kT = tensor.empty() : tensor<2x4x64x16xf16>
+    %kT_22 = linalg.transpose ins(%k_g_21 : tensor<2x16x4x64xf16>) outs(%kT : tensor<2x4x64x16xf16>) permutation = [0, 2, 3, 1]
+    %vv = linalg.transpose ins(%v_g_16 : tensor<2x16x4x64xf16>) outs(%q_8 : tensor<2x4x16x64xf16>) permutation = [0, 2, 1, 3]
+    %scores_23 = tensor.reshape %q_14(%cst) : (tensor<2x4x16x64xf16>, tensor<3xindex>) -> tensor<8x16x64xf16>
+    %scores_24 = tensor.reshape %kT_22(%scores_1) : (tensor<2x4x64x16xf16>, tensor<3xindex>) -> tensor<8x64x16xf16>
+    %scores_25 = linalg.batch_matmul ins(%scores_23, %scores_24 : tensor<8x16x64xf16>, tensor<8x64x16xf16>) outs(%cst_5 : tensor<8x16x16xf32>) -> tensor<8x16x16xf32>
+    %scores_26 = tensor.reshape %scores_25(%scores) : (tensor<8x16x16xf32>, tensor<4xindex>) -> tensor<2x4x16x16xf32>
+    %block_max_27 = tensor.empty() : tensor<2x4x16xf32>
+    %block_max_28 = linalg.fill ins(%block_max : f32) outs(%block_max_27 : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+    %block_max_29 = linalg.reduce ins(%scores_26 : tensor<2x4x16x16xf32>) outs(%block_max_28 : tensor<2x4x16xf32>) dimensions = [3] 
+      (%scores_50: f32, %block_max_51: f32) {
+        %block_max_52 = arith.maxnumf %scores_50, %block_max_51 : f32
+        linalg.yield %block_max_52 : f32
+      }
+    %m_new = arith.maxnumf %block_max_29, %cst_6 : tensor<2x4x16xf32>
+    %correction = arith.subf %cst_6, %m_new : tensor<2x4x16xf32>
+    %correction_30 = math.exp %correction : tensor<2x4x16xf32>
+    %p = tensor.empty() : tensor<2x4x16x16xf32>
+    %p_31 = linalg.broadcast ins(%m_new : tensor<2x4x16xf32>) outs(%p : tensor<2x4x16x16xf32>) dimensions = [3]
+    %p_32 = arith.subf %scores_26, %p_31 : tensor<2x4x16x16xf32>
+    %p_33 = math.exp %p_32 : tensor<2x4x16x16xf32>
+    %l_i_34 = arith.mulf %correction_30, %l_i_3 : tensor<2x4x16xf32>
+    %l_i_35 = linalg.fill ins(%l_i : f32) outs(%block_max_27 : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
+    %l_i_36 = linalg.reduce ins(%p_33 : tensor<2x4x16x16xf32>) outs(%l_i_35 : tensor<2x4x16xf32>) dimensions = [3] 
+      (%p_50: f32, %l_i_51: f32) {
+        %l_i_52 = arith.addf %p_50, %l_i_51 : f32
+        linalg.yield %l_i_52 : f32
+      }
+    %l_i_37 = arith.addf %l_i_34, %l_i_36 : tensor<2x4x16xf32>
+    %acc_38 = tensor.expand_shape %correction_30 [[0], [1], [2, 3]] output_shape [2, 4, 16, 1] : tensor<2x4x16xf32> into tensor<2x4x16x1xf32>
+    %acc_39 = arith.mulf %acc_38, %acc_2 : tensor<2x4x16x1xf32>
+    %acc_40 = tensor.collapse_shape %acc_39 [[0], [1], [2, 3]] : tensor<2x4x16x1xf32> into tensor<2x4x16xf32>
+    %acc_41 = linalg.broadcast ins(%acc_40 : tensor<2x4x16xf32>) outs(%q_11 : tensor<2x4x16x64xf32>) dimensions = [3]
+    %acc_42 = arith.truncf %p_33 : tensor<2x4x16x16xf32> to tensor<2x4x16x16xf16>
+    %acc_43 = tensor.reshape %acc_42(%acc_0) : (tensor<2x4x16x16xf16>, tensor<3xindex>) -> tensor<8x16x16xf16>
+    %acc_44 = tensor.reshape %vv(%cst) : (tensor<2x4x16x64xf16>, tensor<3xindex>) -> tensor<8x16x64xf16>
+    %acc_45 = linalg.batch_matmul ins(%acc_43, %acc_44 : tensor<8x16x16xf16>, tensor<8x16x64xf16>) outs(%cst_4 : tensor<8x16x64xf32>) -> tensor<8x16x64xf32>
+    %acc_46 = tensor.reshape %acc_45(%acc) : (tensor<8x16x64xf32>, tensor<4xindex>) -> tensor<2x4x16x64xf32>
+    %acc_47 = arith.addf %acc_41, %acc_46 : tensor<2x4x16x64xf32>
+    %acc_48 = linalg.broadcast ins(%l_i_37 : tensor<2x4x16xf32>) outs(%q_11 : tensor<2x4x16x64xf32>) dimensions = [3]
+    %acc_49 = arith.divf %acc_47, %acc_48 : tensor<2x4x16x64xf32>
+    %0 = arith.truncf %acc_49 : tensor<2x4x16x64xf32> to tensor<2x4x16x64xf16>
+
+    %1 = ktdp.construct_access_tile %o_desc[%c0, %c0, %c0, %c0] {access_tile_order = #map, access_tile_set = #set3} : memref<2x4x16x64xf16> -> !ktdp.access_tile<2x4x16x64xindex>
+    ktdp.store %0, %1 : tensor<2x4x16x64xf16>, <2x4x16x64xindex>
+    return
+  }
+}

kernels/paged_attn/spyre_aware.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Spyre-aware paged-attention kernel.
+# Original: kernels/paged_attn/original.py
+# Changes summarized in kernels/paged_attn/conversion-notes.md.
+#
+# Same semantics as tensor_descriptor.py, but the K/V gather uses the EXTENDED
+# descriptor_gather, which accepts an index tensor and a src descriptor of any
+# rank. So instead of viewing the cache 2-D and flattening the slot indices to a
+# 1-D row vector, this variant:
+#   - keeps K/V as 3-D descriptors (CACHE, H, D) with block_shape [1, BLK_H, D],
+#   - passes the 2-D slot tile (BLK_B, KV_BLOCK) straight in as the index,
+#   - and the head offset h_start as the second index,
+# yielding (BLK_B, KV_BLOCK, BLK_H, D) directly -- no reshape of the index or of
+# the gathered result. This maps the data-dependent gather onto the Spyre
+# indirect-access tile (ktdp.construct_indirect_access_tile) more directly.
+#
+# Descriptors only -- no raw pointer arithmetic anywhere.
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _paged_attn_kernel_NHD_sa(
+    Q,      # (B, Lq, H, D)
+    K,      # (CACHE, H, D)
+    V,      # (CACHE, H, D)
+    SLOTS,  # (B, Lk) absolute physical slot index per (request, token)
+    Out,    # (B, H, Lq, D)
+    scale,
+    B: tl.constexpr, H: tl.constexpr, Lq: tl.constexpr, Lk: tl.constexpr,
+    CACHE: tl.constexpr,
+    KV_BLOCK: tl.constexpr,
+    BLOCK_Q: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLK_B: tl.constexpr,     # B blocking factor
+    BLK_H: tl.constexpr,     # H blocking factor
+):
+    # --- tensor descriptors (no pointer arithmetic) ---
+    # Q tile spans BLK_B batches x BLOCK_Q queries x BLK_H heads x D
+    q_desc = tl.make_tensor_descriptor(
+        Q, shape=[B, Lq, H, BLOCK_D], strides=[Lq * H * BLOCK_D, H * BLOCK_D, BLOCK_D, 1],
+        block_shape=[BLK_B, BLOCK_Q, BLK_H, BLOCK_D],
+    )
+    # K/V kept 3-D; a single BLK_H-head slice per gathered row
+    k_desc = tl.make_tensor_descriptor(
+        K, shape=[CACHE, H, BLOCK_D], strides=[H * BLOCK_D, BLOCK_D, 1], block_shape=[1, BLK_H, BLOCK_D],
+    )
+    v_desc = tl.make_tensor_descriptor(
+        V, shape=[CACHE, H, BLOCK_D], strides=[H * BLOCK_D, BLOCK_D, 1], block_shape=[1, BLK_H, BLOCK_D],
+    )
+    # SLOTS for both batches at once: (BLK_B, KV_BLOCK)
+    s_desc = tl.make_tensor_descriptor(
+        SLOTS, shape=[B, Lk], strides=[Lk, 1], block_shape=[BLK_B, KV_BLOCK],
+    )
+    # Out tile spans BLK_B batches x BLK_H heads x BLOCK_Q queries x D
+    o_desc = tl.make_tensor_descriptor(
+        Out, shape=[B, H, Lq, BLOCK_D], strides=[H * Lq * BLOCK_D, Lq * BLOCK_D, BLOCK_D, 1],
+        block_shape=[BLK_B, BLK_H, BLOCK_Q, BLOCK_D],
+    )
+
+    n_pages = tl.cdiv(Lk, KV_BLOCK)
+
+    # ===== batch loop, blocked by BLK_B  (for b_start in range(0, B, 2)) =====
+    for b_start in range(0, B, BLK_B):
+        # ===== query-block loop  (for lq_start in range(0, Lq, q_block_size)) =====
+        for lq_start in range(0, Lq, BLOCK_Q):
+            # ===== head loop, blocked by BLK_H  (for h_start in range(0, H, 4)) =====
+            for h_start in range(0, H, BLK_H):
+
+                # load Q (BLK_B, BLOCK_Q, BLK_H, D) -> (BLK_B, BLK_H, BLOCK_Q, D); scale
+                q = q_desc.load([b_start, lq_start, h_start, 0])
+                q = tl.permute(q, (0, 2, 1, 3))                    # (BLK_B, BLK_H, BLOCK_Q, D)
+                q = (q.to(tl.float32) * scale).to(tl.float16)
+
+                # online-softmax state, batched over (BLK_B, BLK_H)
+                m_i = tl.full([BLK_B, BLK_H, BLOCK_Q], float("-inf"), tl.float32)
+                l_i = tl.zeros([BLK_B, BLK_H, BLOCK_Q], tl.float32)
+                acc = tl.zeros([BLK_B, BLK_H, BLOCK_Q, BLOCK_D], tl.float32)
+
+                # ===== KV-page loop =====
+                for j in range(0, n_pages):
+                    # both batches' absolute slot indices: (BLK_B, KV_BLOCK)
+                    slots = s_desc.load([b_start, j * KV_BLOCK])           # (BLK_B, KV_BLOCK)
+
+                    # ONE batched gather over both batches -> (BLK_B, KV_BLOCK, BLK_H, D)
+                    # via the extended any-rank descriptor_gather (2-D index, 3-D src).
+                    k_g = k_desc.gather(slots, h_start)
+                    v_g = v_desc.gather(slots, h_start)
+                    k_g = (k_g.to(tl.float32) * scale).to(tl.float16)
+
+                    kT = tl.permute(k_g, (0, 2, 3, 1))             # (BLK_B, BLK_H, D, KV_BLOCK)
+                    vv = tl.permute(v_g, (0, 2, 1, 3))             # (BLK_B, BLK_H, KV_BLOCK, D)
+
+                    # 4-D batched matmul over (BLK_B, BLK_H)
+                    scores = tl.dot(q, kT)                         # (BLK_B, BLK_H, BLOCK_Q, KV_BLOCK)
+
+                    block_max = tl.max(scores, axis=3)             # (BLK_B, BLK_H, BLOCK_Q)
+                    m_new = tl.maximum(m_i, block_max)
+                    correction = tl.exp(m_i - m_new)
+                    p = tl.exp(scores - m_new[:, :, :, None])
+
+                    l_i = l_i * correction + tl.sum(p, axis=3)
+                    acc = acc * correction[:, :, :, None] + tl.dot(p.to(tl.float16), vv)
+                    m_i = m_new
+
+                acc = acc / l_i[:, :, :, None]                     # (BLK_B, BLK_H, BLOCK_Q, D)
+                o_desc.store([b_start, h_start, lq_start, 0], acc.to(tl.float16))