refactor: extract load_q_mma_frags template, unify comment style

- Add load_q_mma_frags<KD>() shared template in attn_mma_utils.cuh
- Replace ~15 duplicated Q-load lines in 3 MMA kernels
- Unify section header comment style to // ---- Section ----
- Remove duplicate separator line in attn_mma_utils.cuh

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -8,39 +8,21 @@ using bf16 = __nv_bfloat16;
 
 // Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
 //
-// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
+// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core
-// on its own. But GQA gives us G = q_head / kv_head query heads that all share
+// work on its own. But GQA gives us G = q_head / kv_head query heads that all
-// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
+// share one kv_head. We pack those G heads into the M=16 rows of
-// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
+// mma.sync.m16n8k16, turning G independent GEMVs into a single GEMM that
-// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
+// reuses each loaded K/V tile across all G heads (K/V load is the decode
-// not the flops). The KV sequence is partitioned across gridDim.z blocks so that
+// bottleneck, so the reuse is the win, not the flops). The KV sequence is
-// a decode with only batch*kv_head independent tasks can fill all SMs. Each
+// partitioned across gridDim.z blocks so that a decode with only
-// (batch, kv_head, split) block computes an UN-normalised partial (Oacc, m, l)
+// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head,
-// over its KV slice; the combine kernel below reduces across splits. Fixes the
+// split) block computes an UN-normalised partial (Oacc, m, l) over its KV
-// "grid too small" bottleneck (0.04 waves/SM → many blocks) for long-context,
+// slice; the combine kernel below reduces across splits. Fixes the "grid too
+// small" bottleneck (0.04 waves/SM → many blocks) for long-context,
 // small-batch decode.
-//
+
-// Partial layout (float, contiguous):
-//   o_part : [batch, q_head, num_splits, HEAD_DIM]
-//   ml_part: [batch, q_head, num_splits, 2]  (m, l)
-//
-// Optimizations:
-//   - cp.async global→shared for K/V (bypasses registers, cuts instruction count)
-//   - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts
-//   - Q loaded directly from global into mma A-operand registers (no sQ staging,
-//     no prologue syncwarp) — frees shared memory for double-buffering
-//   - Double-buffered KV (STAGES=2): next tile's cp.async overlaps current
-//     tile's MMA compute — hides global load latency / boosts bandwidth
-//     utilization for small-batch (low-occupancy) decode
-//   - Predicated cp.async (cp_async_16_pred) for full AND partial tiles on one
-//     uniform path — eliminates the scalar fallback branch
-//
-// Smem footprint (BC=32): STAGES=2 → 2*(sK+sV) = 2*2*32*HEAD_DIM*2 bytes.
-//   D=128: 16 KB (fits 48 KB static cap).  D=256: 32 KB (also fits).
-// STAGES=1 fallback (4/8 KB) for smem-constrained configs.
 template <int HEAD_DIM, int BC, int STAGES = 2>
 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
-    constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;
     constexpr int NC8 = BC / 8;
     constexpr int KT2 = BC / 16;
@@ -66,26 +48,13 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
     // ---- Load Q directly from global into mma A-operand registers ----
-    // Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8
-    // cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9.
-    // Stride-based: Q is [batch, q_head, q_len=1, head_dim]
     const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
     unsigned Qa[KD][4];
-#pragma unroll
+    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
-    for (int kt = 0; kt < KD; kt++) {
+                          qra, qrb, va, vb, tid4, Qa);
-        int c = kt * 16 + tid4 * 2;
-        const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
-        const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
-        Qa[kt][0] = va ? pau[0] : 0u;
-        Qa[kt][1] = vb ? pbu[0] : 0u;
-        Qa[kt][2] = va ? pau[4] : 0u;
-        Qa[kt][3] = vb ? pbu[4] : 0u;
-    }
 
     float Oacc[DN8][4];
 #pragma unroll

csrc/kernels/attn_mma_utils.cuh

@@ -108,6 +108,36 @@ __device__ __forceinline__ void cp_async_wait_group() {
     asm volatile("cp.async.wait_group %0;" :: "n"(N));
 }
 
+// ---------------------------------------------------------------------------
+// Q-load: load query rows directly from global memory into mma A-operand
+// register layout. One call replaces ~15 duplicated lines in each MMA kernel.
+// stride_row is p.q_stride_h for decode (q_len=1, G heads) or
+//              p.q_stride_l for prefill (multi-q rows).
+// ---------------------------------------------------------------------------
+template <int KD>
+__device__ inline void load_q_mma_frags(
+    const bf16* __restrict__ q,
+    int stride_row,
+    int stride_d,
+    int qra, int qrb,
+    bool va, bool vb,
+    int tid4,
+    unsigned Qa[KD][4])
+{
+    #pragma unroll
+    for (int kt = 0; kt < KD; kt++) {
+        int c = kt * 16 + tid4 * 2;
+        const unsigned* pau = reinterpret_cast<const unsigned*>(
+            &q[qra * stride_row + c * stride_d]);
+        const unsigned* pbu = reinterpret_cast<const unsigned*>(
+            &q[qrb * stride_row + c * stride_d]);
+        Qa[kt][0] = va ? pau[0] : 0u;
+        Qa[kt][1] = vb ? pbu[0] : 0u;
+        Qa[kt][2] = va ? pau[4] : 0u;
+        Qa[kt][3] = vb ? pbu[4] : 0u;
+    }
+}
+
 // ---------------------------------------------------------------------------
 // Shared MMA compute functions — used by both decode and prefill MMA kernels.
 // Extracted because S=Q@K^T, online softmax, and P@V are structurally identical

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -11,14 +11,9 @@ using bf16 = __nv_bfloat16;
 // directly from the page pool through a page table, eliminating the gather
 // copy.  Each tile (BC=32) fits within a single page (page_size >= 32), so
 // the page-table lookup happens once per tile for cp.async.
-//
+
-// Optimizations mirror attn_decode_split_kv_mma_kernel:
-//   - Q loaded directly from global into mma A-operand registers (no sQ)
-//   - Double-buffered KV (STAGES=2) for D<=128, single-buffer for D=256
-//   - Predicated cp.async for unified full/partial tile path
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
-    constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;
     constexpr int NC8 = BC / 8;
     constexpr int KT2 = BC / 16;
@@ -32,34 +27,23 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     const int gid = lane >> 2;
     const int tid4 = lane & 3;
 
-    const int kv_head_idx = blockIdx.x;
+    const int kv_head = blockIdx.x;
     const int batch = blockIdx.y;
     const int split = blockIdx.z;
     const int G = p.q_head / p.kv_head;
-    const int q_head0 = kv_head_idx * G;
+    const int q_head0 = kv_head * G;
 
     __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
     __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
     // ---- Load Q directly from global into mma A-operand registers ----
-    // Q stride-based: [batch, q_head, q_len=1, head_dim]
     const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
     const bool va = qra < G, vb = qrb < G;
     unsigned Qa[KD][4];
-#pragma unroll
+    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
-    for (int kt = 0; kt < KD; kt++) {
+                          qra, qrb, va, vb, tid4, Qa);
-        int c = kt * 16 + tid4 * 2;
-        const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
-        const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
-        Qa[kt][0] = va ? pau[0] : 0u;
-        Qa[kt][1] = vb ? pbu[0] : 0u;
-        Qa[kt][2] = va ? pau[4] : 0u;
-        Qa[kt][3] = vb ? pbu[4] : 0u;
-    }
 
     float Oacc[DN8][4];
 #pragma unroll
@@ -77,7 +61,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     // Paged strides (constant for the block)
     const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM;
     const int64_t pos_stride  = (int64_t)p.kv_head * HEAD_DIM;
-    const int64_t head_off    = (int64_t)kv_head_idx * HEAD_DIM;
+    const int64_t head_off    = (int64_t)kv_head * HEAD_DIM;
 
     // ---- Load tile lambda: predicated cp.async, paged addressing ----
     auto load_tile = [&](int ti, int buf) {

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -57,7 +57,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int kv_head = q_head / (p.q_head / p.kv_head);
     const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
 
-    // Static shared memory — sized by template parameters at compile time.
+    // ---- Static shared memory: double-buffered K/V ----
     // K/V are double-buffered (STAGES=2): the next tile's cp.async load runs
     // while the current tile's tensor-core math executes, hiding global-load
     // latency (FA2-style software pipeline). No dynamic smem / carveout opt-in.
@@ -65,31 +65,16 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
     __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
 
-    // Load the Q fragments straight from global into the mma A-operand layout
+    // Load Q fragments straight from global into mma A-operand layout.
-    // (m16n8k16, row-major): no sQ staging area and no serialized per-warp
+    // stride_row = p.q_stride_l for prefill (multi-q rows across q_len).
-    // prologue barriers. Each lane reads exactly the 8 Q elements ldmatrix
+    // See attn_mma_utils.cuh for the shared template.
-    // would have produced, pre-scaled by the attention scale. Kept resident in
-    // registers across the tile loop.
-    //   frag[0]/[2]: row = qrow0 + gid ;  frag[1]/[3]: row = qrow0 + gid + 8
-    //   frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
-    // Q stride-based: [batch, q_head, q_len, head_dim]
     const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
     const int qra = qrow0 + gid;
     const int qrb = qrow0 + gid + 8;
     const bool va = qra < p.q_len, vb = qrb < p.q_len;
     unsigned Qa[KD][4];
-#pragma unroll
+    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
-    for (int kt = 0; kt < KD; kt++) {
+                          qra, qrb, va, vb, tid4, Qa);
-        int c = kt * 16 + tid4 * 2;
-        const unsigned* pau = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qra * p.q_stride_l + c * p.q_stride_d]);
-        const unsigned* pbu = reinterpret_cast<const unsigned*>(
-            &p.q[q_base + qrb * p.q_stride_l + c * p.q_stride_d]);
-        Qa[kt][0] = va ? pau[0] : 0u;
-        Qa[kt][1] = vb ? pbu[0] : 0u;
-        Qa[kt][2] = va ? pau[4] : 0u;
-        Qa[kt][3] = vb ? pbu[4] : 0u;
-    }
 
     float Oacc[DN8][4];
 #pragma unroll
@@ -122,6 +107,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     constexpr int VEC = 8;  // bf16 per cp.async unit (16 bytes)
     constexpr int TOTAL = BC * HEAD_DIM;
 
+    // ---- Load tile lambda: predicated cp.async ----
     // Issue cp.async loads for tile `ti` into shared buffer `buf`. Predicated
     // loads zero-fill rows past kv_len, so partial tiles need no scalar path.
     auto load_tile = [&](int ti, int buf) {
@@ -141,7 +127,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         cp_async_commit();
     };
 
-    // Prologue: kick off the first tile's load.
+    // ---- Prologue: issue first tile load ----
     load_tile(0, 0);
 
     for (int ti = 0; ti <= t_end; ti++) {