Paged Attention GLM-5.3

"""Paged attention decode (single-query, GQA) for RTX PRO 6000 (SM120).

Strategy: split-K flash decoding on tensor cores, with a persistent
cross-item pipelined kernel.
  attn_pipe_kernel: work items are (batch, kv_head, split) and are handed
  out by STATIC round-robin so each CTA always knows its next item. One
  continuous page counter drives the cp.async stage ring ACROSS item
  boundaries -- the DRAM stream never drains between items (measured wave
  overhead of a drain-then-refill transition was ~1 us per item, which
  dominates the short sequences). Q tiles are double-buffered and
  prefetched with the next item's first page; P's dead rows are zeroed
  once at kernel start.
    Phase 1 (warp 0): S = Q K^T via mma.m16n8k16 with bf16 fragments
  (ldmatrix x4 on Q and K), online softmax in fp32, P written to smem.
    Phase 2 (all warps): acc = acc * rescale + P V, each warp owning a
  d-slice; V consumed via ldmatrix x2.trans + mma.
  Partials (m, l, acc) are reduced either inline by the last-arriving CTA
  per (batch, kv_head) (FUSE=1) or by a wide attn_combine_kernel (FUSE=0),
  whichever measured faster per shape. attn_split_kernel (work-stealing
  variant) and attn_split2_kernel (two kv heads per CTA) are kept as
  alternates but not dispatched.

Host side: each unique (pointers, dims, config) is captured into a CUDA
graph once and replayed; repeated calls with the same tensors take a
single-int replay path, which saves ~8 us of CPU submit time that the GPU
would otherwise sit idle for on short kernels.

Numerics match the fp32-softmax reference: softmax and accumulation fp32,
KV streamed in bf16 exactly as stored.

Note: torch headers live only in the binding TU (gcc 15 chokes on an ATen
header through nvcc's host pass); the kernel TU is pure CUDA.
"""
import hashlib
import math
import os

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

_CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_pipeline.h>
#include <cstdint>

#define DEVFN __device__ __forceinline__

DEVFN float2 bf2f(unsigned u) {
    // packed bf16 pair -> two floats (bits trick: hi half is already f32 bits)
    float2 r;
    r.x = __uint_as_float(u << 16);
    r.y = __uint_as_float(u & 0xFFFF0000u);
    return r;
}

// ---------------------------------------------------------------------------
// Fused cross-split reduction. Every split CTA of a (b, kvh) pair publishes
// its partials, then bumps ctr[pair]; the last one to arrive reduces all S
// splits for its G heads and writes the final output. ctr is left back at 0
// so the buffer stays valid across calls without a separate memset.
// Canonical write-fence-atomic pattern (cf. CUDA threadFenceReduction).
// ---------------------------------------------------------------------------
// ---------------------------------------------------------------------------
// Cross-split reduction, phase 2: one thread per output float4. Wide launch
// (B*H*D/4 threads) replaces the fused last-arriver reduce -- the fused
// version serialized ~11 us of latency-bound work into 32 CTAs at the very
// end of the kernel; this one spreads the same 1 MB over the whole GPU.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(128) attn_combine_kernel(
    const float* __restrict__ pacc, const float2* __restrict__ pml,
    __nv_bfloat16* __restrict__ out, int S, int B, int H, int D)
{
    const int D4 = D >> 2;
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= B * H * D4) return;
    const int d4 = i % D4;   // D4 = D/4 is a power of two
    const int r = i / D4;
    const int h = r % H;
    const int b = r / H;
    float M = -INFINITY;
    const float2* mlh = pml + (size_t)(b * H + h) * S;
    for (int s2 = 0; s2 < S; s2++) M = fmaxf(M, mlh[s2].x);
    float lo = 0.f;
    float4 ao = make_float4(0.f, 0.f, 0.f, 0.f);
    const float4* pacc4 = (const float4*)pacc + (size_t)(b * H + h) * S * D4 + d4;
    for (int s2 = 0; s2 < S; s2++) {
        float2 ml = mlh[s2];
        float w = __expf(ml.x - M);
        lo = fmaf(w, ml.y, lo);
        float4 a = pacc4[(size_t)s2 * D4];
        ao.x = fmaf(w, a.x, ao.x); ao.y = fmaf(w, a.y, ao.y);
        ao.z = fmaf(w, a.z, ao.z); ao.w = fmaf(w, a.w, ao.w);
    }
    __nv_bfloat16* o = out + ((size_t)b * H + h) * D + d4 * 4;
    const float il = 1.f / lo;
    o[0] = __float2bfloat16(ao.x * il);
    o[1] = __float2bfloat16(ao.y * il);
    o[2] = __float2bfloat16(ao.z * il);
    o[3] = __float2bfloat16(ao.w * il);
}

//---------------------------------------------------------------------------
// Split kernel (tensor-core). One CTA = one (batch, kv_head, sequence chunk)
// item, 128 threads / 4 warps. K and V stay bf16 in shared memory (no
// convert pass); mma.m16n8k16 does both GEMMs:
//   phase 1 (warp 0): S = Q K^T for all G query rows packed as the m16 rows
//                     of one warp's mma (rows G..15 zero).
//   phase 2 (all 4 warps): O += P V, warp w owns the d-slice [32w, 32w+32).
// Online softmax runs on the S fragments; P is published through shared
// memory so every warp can build its A fragments with ldmatrix.
// Pages stream through a 3-stage cp.async pipeline; persistent CTAs steal
// (pair, split) items from a work queue so the tail never goes idle.
// ---------------------------------------------------------------------------
#define PA_LDM_X4(r0, r1, r2, r3, addr)                                         \
    asm volatile(                                                               \
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"       \
        : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)                                \
        : "r"(addr))
#define PA_LDM_X2T(r0, r1, addr)                                                \
    asm volatile(                                                               \
        "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];\n"       \
        : "=r"(r0), "=r"(r1)                                                    \
        : "r"(addr))
#define PA_MMA(c0, c1, c2, c3, a0, a1, a2, a3, b0, b1)                          \
    asm volatile(                                                               \
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "                  \
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"               \
        : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)                                \
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1))

DEVFN unsigned pa_smem_u32(const void* p) {
    return (unsigned)__cvta_generic_to_shared(p);
}

template <int D, int G, int STG>
__global__ void __launch_bounds__(128, STG >= 5 ? 6 : 3) attn_split_kernel(
    const __nv_bfloat16* __restrict__ q,     // (B, H, D)
    const __nv_bfloat16* __restrict__ kv,    // (NB, 16, Hkv, 2D)
    const int* __restrict__ bt,              // (B, mb) page ids
    const int* __restrict__ sl,              // (B,)
    float* __restrict__ pacc,                // (S, B, H, D) unnormalized
    float2* __restrict__ pml,                // (S, B*H) {m, l}
    __nv_bfloat16* __restrict__ out,         // (B, H, D)
    int* __restrict__ ctr,                   // (B*Hkv,) split-arrival counters
    int* __restrict__ wq,                    // (2,) work queue: next item, done CTAs
    int S, int mb, int Hkv, int B, int items, float scale)
{
    constexpr int THREADS = 128;
    constexpr int KP = D + 8;         // bf16 smem row pad: odd # of 16B chunks
    constexpr int KS = D / 16;        // phase-1 k-steps
    constexpr int WPT = D / 32;       // phase-2 n-tiles per warp (8 dims each)
    constexpr int PP = 24;            // P16 row stride, bf16 (3 chunks)

    const int H = Hkv * G;
    const int npairs = B * Hkv;

    __shared__ __nv_bfloat16 K16[STG][16][KP];
    __shared__ __nv_bfloat16 V16[STG][16][KP];
    __shared__ __nv_bfloat16 Q16[16][KP];
    __shared__ __nv_bfloat16 P16[16][PP];
    __shared__ float rss[G];          // per-head acc rescale, warp0 -> all
    __shared__ float mls[G][2];       // per-head {M, l}, warp0 -> all
    __shared__ int s_item;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int g = lane >> 2;          // mma row of this lane = query head idx

    // ---- persistent CTAs: steal (pair, split) items until done ----------
    for (;;) {
    if (tid == 0) s_item = atomicAdd(&wq[0], 1);
    __syncthreads();
    const int item = s_item;
    if (item >= items) break;

    const int pair = item % npairs;   // pair-fastest: neighbours share pages
    const int s = item / npairs;
    const int b = pair / Hkv;
    const int kvh = pair - b * Hkv;
    const int h0 = kvh * G;

    const int Lb = sl[b];
    const int pages_b = (Lb + 15) >> 4;
    const int ppc = (pages_b + S - 1) / S;
    const int lo = s * ppc;
    const int hi = min(lo + ppc, pages_b);

    if (lo >= hi) {  // empty split -> neutral partials
        float* dst = pacc + ((size_t)b * H + h0) * S * D + (size_t)s * G * D;
        for (int i = tid; i < G * D; i += THREADS) dst[i] = 0.f;
        if (tid < G) pml[((size_t)b * H + h0 + tid) * S + s] = make_float2(-INFINITY, 0.f);
        continue;
    }

    // ---- Q16: rows 0..G-1 live, rows G..15 zero (mma padding) -----------
    {
        const __nv_bfloat16* qp = q + ((size_t)b * H + h0) * D;
        for (int i = tid; i < 16 * KP; i += THREADS) {
            int r = i / KP, c = i - r * KP;
            Q16[r][c] = (r < G && c < D) ? qp[r * D + c]
                                         : __float2bfloat16(0.f);
        }
        for (int i = tid; i < 16 * PP; i += THREADS)
            P16[i / PP][i - (i / PP) * PP] = __float2bfloat16(0.f);
    }
    __syncthreads();

    const size_t slab = (size_t)Hkv * 2 * D;   // elems between token rows
    const int* btrow = bt + (size_t)b * mb;

    // ---- cp.async one page: K and V of this kv head into bf16 tiles -----
    auto cp_page = [&](int pi, int stg) {
        const __nv_bfloat16* src =
            kv + ((size_t)btrow[pi] * 16) * slab + (size_t)kvh * 2 * D;
        constexpr int RC = D / 8;     // 16B chunks per row per K (or V)
        for (int c = tid; c < 16 * RC; c += THREADS) {
            int t = c / RC;
            int o = (c - t * RC) * 8;
            __pipeline_memcpy_async(&K16[stg][t][o],
                                    src + (size_t)t * slab + o, 16);
            __pipeline_memcpy_async(&V16[stg][t][o],
                                    src + (size_t)t * slab + D + o, 16);
        }
        __pipeline_commit();
    };

    const int np = hi - lo;
    const int npre = np < STG - 1 ? np : STG - 1;   // prefetch depth
#pragma unroll
    for (int j = 0; j < STG - 1; j++)
        if (j < npre) cp_page(lo + j, j % STG);
    __pipeline_wait_prior(npre > 1 ? npre - 1 : 0);  // page 0 landed
    __syncthreads();

    float m_r = -INFINITY, l_r = 0.f;  // warp0: stats of row g (head g)
    float acc[WPT][4];                 // this warp's d-slice, rows g and g+8
#pragma unroll
    for (int w = 0; w < WPT; w++)
#pragma unroll
        for (int j = 0; j < 4; j++) acc[w][j] = 0.f;

    for (int i = 0; i < np; i++) {
        const int page = lo + i;
        const int rem = min(16, Lb - page * 16);   // valid tokens this page
        const __nv_bfloat16(*Kp)[KP] = K16[i % 3];
        const __nv_bfloat16(*Vp)[KP] = V16[i % 3];

        if (warp == 0) {
            // ---- phase 1: S = Q K^T, all G rows, both n-tiles ----------
            float s0[4], s1[4];
#pragma unroll
            for (int j = 0; j < 4; j++) { s0[j] = 0.f; s1[j] = 0.f; }
#pragma unroll
            for (int ks = 0; ks < KS; ks++) {
                unsigned a0, a1, a2, a3, k0r, k1r, k2r, k3r;
                unsigned qa = pa_smem_u32(&Q16[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(a0, a1, a2, a3, qa);
                unsigned ka = pa_smem_u32(&Kp[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(k0r, k1r, k2r, k3r, ka);
                PA_MMA(s0[0], s0[1], s0[2], s0[3], a0, a1, a2, a3, k0r, k2r);
                PA_MMA(s1[0], s1[1], s1[2], s1[3], a0, a1, a2, a3, k1r, k3r);
            }
            // ---- online softmax over this lane's 4 tokens --------------
            // D-frag: row g holds cols 2*(lane%4)+{0,1} of each n-tile.
            const int t0 = (lane % 4) * 2;
            const int tk[4] = {t0, t0 + 1, t0 + 8, t0 + 9};
            float sc[4] = {s0[0], s0[1], s1[0], s1[1]};
            float mx = -INFINITY;
#pragma unroll
            for (int j = 0; j < 4; j++) {
                sc[j] = tk[j] < rem ? sc[j] * scale : -INFINITY;
                mx = fmaxf(mx, sc[j]);
            }
            mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 1));
            mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 2));
            // e in the QUAD-max basis (not the lane-local max): P = e*rs_pg
            // must equal exp(sc - M); a lane-local basis silently drops the
            // exp(mxl - mx) factor, which only shows once logits spread.
            float e[4], sm = 0.f;
#pragma unroll
            for (int j = 0; j < 4; j++) { e[j] = __expf(sc[j] - mx); sm += e[j]; }
            sm += __shfl_xor_sync(0xffffffffu, sm, 1);
            sm += __shfl_xor_sync(0xffffffffu, sm, 2);

            const float M = fmaxf(m_r, mx);
            const float rs_old = __expf(m_r - M);
            const float rs_pg = __expf(mx - M);
            l_r = l_r * rs_old + sm * rs_pg;
            m_r = M;
            if (g < G) {
                // P in global-max basis; masked lanes give exact zeros
#pragma unroll
                for (int j = 0; j < 4; j++)
                    P16[g][tk[j]] = __float2bfloat16(e[j] * rs_pg);
                if ((lane & 3) == 0) rss[g] = rs_old;
            }
        }
        __syncthreads();

        // ---- phase 2: acc = acc * rs + P V  (warp w owns its d-slice) --
        {
            const float rs = g < G ? rss[g] : 0.f;
            unsigned a0, a1, a2, a3;
            unsigned pa = pa_smem_u32(&P16[lane % 16][(lane / 16) * 8]);
            PA_LDM_X4(a0, a1, a2, a3, pa);
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                unsigned b0, b1;
                unsigned va = pa_smem_u32(&Vp[lane % 16][(warp * WPT + w) * 8]);
                PA_LDM_X2T(b0, b1, va);
#pragma unroll
                for (int j = 0; j < 4; j++) acc[w][j] *= rs;
                PA_MMA(acc[w][0], acc[w][1], acc[w][2], acc[w][3],
                       a0, a1, a2, a3, b0, b1);
            }
        }

        // ---- pipeline: keep STG-2 pages in flight, wait for i+1 --------
        if (i + 1 < np) {
            const int P = min(i + STG - 1, np - 1);   // last page to issue
            if (P == i + STG - 1) cp_page(lo + P, P % STG);
            __pipeline_wait_prior(P - i - 1);
            __syncthreads();  // every warp done reading stage i%3
        }
    }

    // ---- publish final stats (all warps need l for the epilogue) -------
    if (warp == 0 && (lane & 3) == 0 && g < G) {
        mls[g][0] = m_r;
        mls[g][1] = l_r;
    }
    __syncthreads();

    if (S == 1) {
        if (g < G) {
            const float il = 1.f / mls[g][1];
            __nv_bfloat16* o = out + ((size_t)b * H + h0 + g) * D;
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                o[d0] = __float2bfloat16(acc[w][0] * il);
                o[d0 + 1] = __float2bfloat16(acc[w][1] * il);
            }
        }
    } else {
        if (warp == 0 && (lane & 3) == 0 && g < G)
            pml[((size_t)b * H + h0 + g) * S + s] = make_float2(m_r, l_r);
        // partials only: a wide combine kernel follows (kernel boundary
        // orders the writes, so no fence/counter spin here). Layout
        // [B][H][S][D] keeps each thread's S loads contiguous so the
        // combine is bandwidth-bound, not latency-bound.
        if (g < G) {
            float* dst = pacc + (((size_t)b * H + h0 + g) * S + s) * D;
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                dst[d0] = acc[w][0];
                dst[d0 + 1] = acc[w][1];
            }
        }
    }
    __syncthreads();  // smem free before the next item reuses it
    }

    // last CTA to leave resets the queue for the next launch
    if (tid == 0) {
        if (atomicAdd(&wq[1], 1) == (int)gridDim.x - 1) {
            wq[0] = 0;
            wq[1] = 0;
        }
    }
}


// ---------------------------------------------------------------------------
// KV-merged split kernel: one CTA owns TWO adjacent kv heads (kvh, kvh+1)
// and all their GQA query heads -- G*2 rows of the m16 tile (rows 0..G-1 =
// kv head 0's heads, rows 8..8+G-1 = kv head 1's). Each KV row is read as
// one contiguous 2*KVH*D bf16 run (K0|V0|K1|V1), doubling the DRAM burst
// length and halving per-page loop overheads for the same bytes. Phase 2
// runs two mmas per d-tile (B = V0 into accA, B = V1 into accB); the
// cross terms land in accumulator rows the epilogue never reads.
// ---------------------------------------------------------------------------
template <int D, int G, int STG>
__global__ void __launch_bounds__(128, D == 128 ? 2 : 3) attn_split2_kernel(
    const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kv,
    const int* __restrict__ bt, const int* __restrict__ sl,
    float* __restrict__ pacc, float2* __restrict__ pml,
    __nv_bfloat16* __restrict__ out, int* __restrict__ ctr,
    int* __restrict__ wq, int S, int mb, int Hkv, int B, int items,
    float scale)
{
    constexpr int THREADS = 128;
    constexpr int KP = D + 8;
    constexpr int KS = D / 16;
    constexpr int WPT = D / 32;
    constexpr int PP = 24;

    const int H = Hkv * G;
    const int npairs = B * (Hkv / 2);   // merged pairs

    __shared__ __nv_bfloat16 K16[STG][2][16][KP];
    __shared__ __nv_bfloat16 V16[STG][2][16][KP];
    __shared__ __nv_bfloat16 Q16[16][KP];
    __shared__ __nv_bfloat16 P16[16][PP];
    __shared__ float rss[2][G];
    __shared__ float mls[2][G][2];
    __shared__ int s_item;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int g = lane >> 2;

    for (;;) {
    if (tid == 0) s_item = atomicAdd(&wq[0], 1);
    __syncthreads();
    const int item = s_item;
    if (item >= items) break;

    const int pair = item % npairs;
    const int s = item / npairs;
    const int b = pair / (Hkv / 2);
    const int kvh = (pair - b * (Hkv / 2)) * 2;
    const int h0 = kvh * G;

    const int Lb = sl[b];
    const int pages_b = (Lb + 15) >> 4;
    const int ppc = (pages_b + S - 1) / S;
    const int lo = s * ppc;
    const int hi = min(lo + ppc, pages_b);

    if (lo >= hi) {  // empty split -> neutral partials, head-row correct
        for (int i = tid; i < 2 * G * D; i += THREADS) {
            int h = h0 + i / D, d = i - (i / D) * D;
            pacc[((size_t)(b * H + h) * S + s) * D + d] = 0.f;
        }
        for (int i = tid; i < 2 * G; i += THREADS)
            pml[((size_t)b * H + h0 + i) * S + s] = make_float2(-INFINITY, 0.f);
        continue;
    }

    // ---- Q16: rows 0..G-1 = kv head 0's heads, rows 8..8+G-1 = head 1's -
    {
        const __nv_bfloat16* qp = q + ((size_t)b * H + h0) * D;
        for (int i = tid; i < 16 * KP; i += THREADS) {
            int r = i / KP, c = i - r * KP;
            const bool live = (r < G || (r >= 8 && r < 8 + G)) && c < D;
            Q16[r][c] = live ? qp[(r < 8 ? r : G + r - 8) * D + c]
                             : __float2bfloat16(0.f);
        }
        for (int i = tid; i < 16 * PP; i += THREADS)
            P16[i / PP][i - (i / PP) * PP] = __float2bfloat16(0.f);
    }
    __syncthreads();

    const size_t slab = (size_t)Hkv * 2 * D;
    const int* btrow = bt + (size_t)b * mb;

    auto cp_page = [&](int pi, int stg) {
        const __nv_bfloat16* src =
            kv + ((size_t)btrow[pi] * 16) * slab + (size_t)kvh * 2 * D;
        constexpr int RC = D / 8;
#pragma unroll
        for (int kv2 = 0; kv2 < 2; kv2++) {
            for (int c = tid; c < 16 * RC; c += THREADS) {
                int t = c / RC;
                int o = (c - t * RC) * 8;
                __pipeline_memcpy_async(&K16[stg][kv2][t][o],
                                        src + (size_t)t * slab + kv2 * 2 * D + o, 16);
                __pipeline_memcpy_async(&V16[stg][kv2][t][o],
                                        src + (size_t)t * slab + kv2 * 2 * D + D + o, 16);
            }
        }
        __pipeline_commit();
    };

    const int np = hi - lo;
    const int npre = np < STG - 1 ? np : STG - 1;
#pragma unroll
    for (int j = 0; j < STG - 1; j++)
        if (j < npre) cp_page(lo + j, j % STG);
    __pipeline_wait_prior(npre > 1 ? npre - 1 : 0);
    __syncthreads();

    float m0_r = -INFINITY, l0_r = 0.f;   // head h0+g   (kv head 0)
    float m1_r = -INFINITY, l1_r = 0.f;   // head h0+G+g (kv head 1)
    float accA[WPT][4], accB[WPT][4];
#pragma unroll
    for (int w = 0; w < WPT; w++)
#pragma unroll
        for (int j = 0; j < 4; j++) { accA[w][j] = 0.f; accB[w][j] = 0.f; }

    for (int i = 0; i < np; i++) {
        const int page = lo + i;
        const int rem = min(16, Lb - page * 16);
        const __nv_bfloat16(*Kp)[KP] = K16[i % STG][0];
        const __nv_bfloat16(*Kq)[KP] = K16[i % STG][1];
        const __nv_bfloat16(*Vp)[KP] = V16[i % STG][0];
        const __nv_bfloat16(*Vq)[KP] = V16[i % STG][1];

        if (warp == 0) {
            float sA0[4], sA1[4], sB0[4], sB1[4];
#pragma unroll
            for (int j = 0; j < 4; j++) {
                sA0[j] = 0.f; sA1[j] = 0.f; sB0[j] = 0.f; sB1[j] = 0.f;
            }
#pragma unroll
            for (int ks = 0; ks < KS; ks++) {
                unsigned a0, a1, a2, a3, ka0, ka1, ka2, ka3, kb0, kb1, kb2, kb3;
                unsigned qa = pa_smem_u32(&Q16[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(a0, a1, a2, a3, qa);
                unsigned k0a = pa_smem_u32(&Kp[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(ka0, ka1, ka2, ka3, k0a);
                unsigned k0b = pa_smem_u32(&Kq[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(kb0, kb1, kb2, kb3, k0b);
                PA_MMA(sA0[0], sA0[1], sA0[2], sA0[3], a0, a1, a2, a3, ka0, ka2);
                PA_MMA(sA1[0], sA1[1], sA1[2], sA1[3], a0, a1, a2, a3, ka1, ka3);
                PA_MMA(sB0[0], sB0[1], sB0[2], sB0[3], a0, a1, a2, a3, kb0, kb2);
                PA_MMA(sB1[0], sB1[1], sB1[2], sB1[3], a0, a1, a2, a3, kb1, kb3);
            }
            // kv head 0's head g lives in row g (frag regs 0,1 of each tile)
            // kv head 1's head g lives in row g+8 (frag regs 2,3)
            const int t0 = (lane % 4) * 2;
            const int tk[4] = {t0, t0 + 1, t0 + 8, t0 + 9};
            float sc0[4] = {sA0[0], sA0[1], sA1[0], sA1[1]};
            float sc1[4] = {sB0[2], sB0[3], sB1[2], sB1[3]};
            float mx0 = -INFINITY, mx1 = -INFINITY;
#pragma unroll
            for (int j = 0; j < 4; j++) {
                sc0[j] = tk[j] < rem ? sc0[j] * scale : -INFINITY;
                sc1[j] = tk[j] < rem ? sc1[j] * scale : -INFINITY;
                mx0 = fmaxf(mx0, sc0[j]);
                mx1 = fmaxf(mx1, sc1[j]);
            }
            mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 1));
            mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, 2));
            mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 1));
            mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, 2));
            // e in the QUAD-max basis; see attn_split_kernel note
            float e0[4], e1[4], sm0 = 0.f, sm1 = 0.f;
#pragma unroll
            for (int j = 0; j < 4; j++) {
                e0[j] = __expf(sc0[j] - mx0); sm0 += e0[j];
                e1[j] = __expf(sc1[j] - mx1); sm1 += e1[j];
            }
            sm0 += __shfl_xor_sync(0xffffffffu, sm0, 1);
            sm0 += __shfl_xor_sync(0xffffffffu, sm0, 2);
            sm1 += __shfl_xor_sync(0xffffffffu, sm1, 1);
            sm1 += __shfl_xor_sync(0xffffffffu, sm1, 2);

            const float M0 = fmaxf(m0_r, mx0), M1 = fmaxf(m1_r, mx1);
            const float rs0o = __expf(m0_r - M0), rs0p = __expf(mx0 - M0);
            const float rs1o = __expf(m1_r - M1), rs1p = __expf(mx1 - M1);
            l0_r = l0_r * rs0o + sm0 * rs0p;
            l1_r = l1_r * rs1o + sm1 * rs1p;
            m0_r = M0;
            m1_r = M1;
            if (g < G) {
                P16[g][tk[0]]     = __float2bfloat16(e0[0] * rs0p);
                P16[g][tk[1]]     = __float2bfloat16(e0[1] * rs0p);
                P16[g][tk[2]]     = __float2bfloat16(e0[2] * rs0p);
                P16[g][tk[3]]     = __float2bfloat16(e0[3] * rs0p);
                P16[8 + g][tk[0]] = __float2bfloat16(e1[0] * rs1p);
                P16[8 + g][tk[1]] = __float2bfloat16(e1[1] * rs1p);
                P16[8 + g][tk[2]] = __float2bfloat16(e1[2] * rs1p);
                P16[8 + g][tk[3]] = __float2bfloat16(e1[3] * rs1p);
                if ((lane & 3) == 0) { rss[0][g] = rs0o; rss[1][g] = rs1o; }
            }
        }
        __syncthreads();

        {
            const float rs0 = g < G ? rss[0][g] : 0.f;
            const float rs1 = g < G ? rss[1][g] : 0.f;
            unsigned a0, a1, a2, a3;
            unsigned pa = pa_smem_u32(&P16[lane % 16][(lane / 16) * 8]);
            PA_LDM_X4(a0, a1, a2, a3, pa);
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                unsigned b0, b1, c0, c1;
                unsigned va = pa_smem_u32(&Vp[lane % 16][(warp * WPT + w) * 8]);
                PA_LDM_X2T(b0, b1, va);
                unsigned vb = pa_smem_u32(&Vq[lane % 16][(warp * WPT + w) * 8]);
                PA_LDM_X2T(c0, c1, vb);
#pragma unroll
                for (int j = 0; j < 4; j++) {
                    accA[w][j] *= rs0;
                    accB[w][j] *= rs1;
                }
                PA_MMA(accA[w][0], accA[w][1], accA[w][2], accA[w][3],
                       a0, a1, a2, a3, b0, b1);
                PA_MMA(accB[w][0], accB[w][1], accB[w][2], accB[w][3],
                       a0, a1, a2, a3, c0, c1);
            }
        }

        if (i + 1 < np) {
            const int P = min(i + STG - 1, np - 1);
            if (P == i + STG - 1) cp_page(lo + P, P % STG);
            __pipeline_wait_prior(P - i - 1);
            __syncthreads();
        }
    }

    if (warp == 0 && (lane & 3) == 0 && g < G) {
        mls[0][g][0] = m0_r; mls[0][g][1] = l0_r;
        mls[1][g][0] = m1_r; mls[1][g][1] = l1_r;
    }
    __syncthreads();

    if (S == 1) {
        if (g < G) {
            const float il0 = 1.f / mls[0][g][1];
            const float il1 = 1.f / mls[1][g][1];
            __nv_bfloat16* o0 = out + ((size_t)b * H + h0 + g) * D;
            __nv_bfloat16* o1 = out + ((size_t)b * H + h0 + G + g) * D;
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                o0[d0]     = __float2bfloat16(accA[w][0] * il0);
                o0[d0 + 1] = __float2bfloat16(accA[w][1] * il0);
                o1[d0]     = __float2bfloat16(accB[w][2] * il1);
                o1[d0 + 1] = __float2bfloat16(accB[w][3] * il1);
            }
        }
    } else {
        if (warp == 0 && (lane & 3) == 0 && g < G) {
            pml[((size_t)b * H + h0 + g) * S + s] = make_float2(m0_r, l0_r);
            pml[((size_t)b * H + h0 + G + g) * S + s] = make_float2(m1_r, l1_r);
        }
        if (g < G) {
            float* d0p = pacc + (((size_t)b * H + h0 + g) * S + s) * D;
            float* d1p = pacc + (((size_t)b * H + h0 + G + g) * S + s) * D;
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                d0p[d0] = accA[w][0]; d0p[d0 + 1] = accA[w][1];
                d1p[d0] = accB[w][2]; d1p[d0 + 1] = accB[w][3];
            }
        }
    }
    __syncthreads();
    }

    if (tid == 0) {
        if (atomicAdd(&wq[1], 1) == (int)gridDim.x - 1) {
            wq[0] = 0;
            wq[1] = 0;
        }
    }
}


// ---------------------------------------------------------------------------
// Pipelined persistent split kernel. Items are taken by STATIC round-robin
// (blockIdx.x, +gridDim.x, ...) so the next item is always known, and ONE
// continuous page counter drives the stage ring ACROSS item boundaries: the
// cp.async stream never drains between items. Q is double-buffered (the next
// item's Q is prefetched in the same commit group as its first page); P16's
// dead rows and Q16's padding are zeroed once. Prefetch may run at most ONE
// item ahead of the consumer so Q16[v] is never overwritten while live.
// Rationale: measured wave overhead of drain-then-refill in the steal kernel
// is ~1 us per item transition, which dominates short items (np < 16).
// ---------------------------------------------------------------------------
template <int D, int G, int STG>
__global__ void __launch_bounds__(128, 3) attn_pipe_kernel(
    const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ kv,
    const int* __restrict__ bt, const int* __restrict__ sl,
    float* __restrict__ pacc, float2* __restrict__ pml,
    __nv_bfloat16* __restrict__ out, int* __restrict__ ctr,
    int* __restrict__ wq, int S, int mb, int Hkv, int B, int items,
    int fuse, float scale)
{
    constexpr int THREADS = 128;
    constexpr int KP = D + 8;
    constexpr int KS = D / 16;
    constexpr int WPT = D / 32;
    constexpr int PP = 24;

    const int H = Hkv * G;
    const int npairs = B * Hkv;

    __shared__ __nv_bfloat16 K16[STG][16][KP];
    __shared__ __nv_bfloat16 V16[STG][16][KP];
    __shared__ __nv_bfloat16 Q16[2][16][KP];
    __shared__ __nv_bfloat16 P16[16][PP];
    __shared__ float rss[G];
    __shared__ float mls[G][2];
    __shared__ int mine[128];     // pairs this CTA must reduce (fused)

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int g = lane >> 2;

    const size_t slab = (size_t)Hkv * 2 * D;

    struct It { int b, kvh, h0, s, lo, hi; };   // page range [lo, hi)
    auto describe = [&](int id, It& t) {
        const int pair = id % npairs;           // pair-fastest: neighbours share pages
        t.s = id / npairs;
        t.b = pair / Hkv;
        t.kvh = pair - t.b * Hkv;
        t.h0 = t.kvh * G;
        const int pages_b = (sl[t.b] + 15) >> 4;
        const int ppc = (pages_b + S - 1) / S;
        t.lo = t.s * ppc;
        t.hi = min(t.lo + ppc, pages_b);
    };
    auto neutral = [&](const It& t) {          // empty split: identity partials
        for (int i = tid; i < G * D; i += THREADS) {
            int h = i / D;
            pacc[((size_t)(t.b * H + t.h0 + h) * S + t.s) * D + i - h * D] = 0.f;
        }
        if (tid < G)
            pml[((size_t)t.b * H + t.h0 + tid) * S + t.s] =
                make_float2(-INFINITY, 0.f);
    };

    // ---- constant zero-fill once: dead rows / padding of Q16 (both bufs),
    //      and all of P16 (live rows are fully rewritten every page) --------
    for (int i = tid; i < 16 * KP; i += THREADS) {
        int r = i / KP, c = i - r * KP;
        if (r >= G || c >= D) {
            const __nv_bfloat16 z = __float2bfloat16(0.f);
            Q16[0][r][c] = z;
            Q16[1][r][c] = z;
        }
    }
    for (int i = tid; i < 16 * PP; i += THREADS)
        P16[i / PP][i - (i / PP) * PP] = __float2bfloat16(0.f);
    __syncthreads();

    // ---- consumer / prefetch stream pointers -------------------------------
    It cur, nxt;
    int cur_id = blockIdx.x;
    bool have = cur_id < items;
    if (have) {
        describe(cur_id, cur);
        while (cur.lo >= cur.hi) {
            neutral(cur);
            cur_id += gridDim.x;
            if (!(have = cur_id < items)) break;
            describe(cur_id, cur);
        }
    }
    if (!have) return;

    nxt = cur;
    int nxt_id = cur_id;
    int pg_n = nxt.lo;    // next page to issue inside nxt
    int P = -1;           // last issued page counter (stream-wide)
    int nvis = 0;         // prefetch-side visit index (Q parity)

    auto cp_page = [&](const It& t, int pi, int stg) {
        const int* btrow = bt + (size_t)t.b * mb;
        const __nv_bfloat16* src =
            kv + ((size_t)btrow[pi] * 16) * slab + (size_t)t.kvh * 2 * D;
        constexpr int RC = D / 8;
        for (int c = tid; c < 16 * RC; c += THREADS) {
            int r = c / RC;
            int o = (c - r * RC) * 8;
            __pipeline_memcpy_async(&K16[stg][r][o],
                                    src + (size_t)r * slab + o, 16);
            __pipeline_memcpy_async(&V16[stg][r][o],
                                    src + (size_t)r * slab + D + o, 16);
        }
    };
    auto cp_q = [&](const It& t, int qb) {
        const __nv_bfloat16* qp = q + ((size_t)t.b * H + t.h0) * D;
        constexpr int CH = G * D / 8;
        for (int c = tid; c < CH; c += THREADS) {
            int r = (c * 8) / D, o = (c * 8) - r * D;
            __pipeline_memcpy_async(&Q16[qb][r][o], qp + (size_t)r * D + o, 16);
        }
    };
    // raise P toward `target`; enter the next item only when the prefetch
    // pointer is still on the consumer's item (one-item-lookahead cap)
    auto try_fill = [&](int target) {
        while (P < target) {
            if (pg_n >= nxt.hi) {
                if (nxt_id != cur_id) return;      // one item ahead already
                int cand = nxt_id + gridDim.x;
                It t;
                for (;;) {
                    if (cand >= items) { nxt_id = items; return; }
                    describe(cand, t);
                    if (t.lo < t.hi) break;
                    cand += gridDim.x;
                }
                nxt_id = cand;
                nxt = t;
                pg_n = nxt.lo;
                nvis++;
                continue;
            }
            cp_page(nxt, pg_n, (P + 1) % STG);
            if (pg_n == nxt.lo) cp_q(nxt, nvis & 1);
            __pipeline_commit();
            P++;
            pg_n++;
        }
    };

    try_fill(STG - 1);   // prologue

    float m_r = -INFINITY, l_r = 0.f;
    float acc[WPT][4];
#pragma unroll
    for (int w = 0; w < WPT; w++)
#pragma unroll
        for (int j = 0; j < 4; j++) acc[w][j] = 0.f;

    int pc = 0;          // consumed page counter (stream-wide)
    int base = 0;        // pc of cur's first page
    int cur_par = 0;     // Q16 buffer of the consuming item
    for (;;) {
        __pipeline_wait_prior(P - pc);
        __syncthreads();
        // issue AFTER the barrier: the stage being refilled was last read at
        // phase 2 of page pc-1, which the barrier above covers
        try_fill(pc + STG - 1);

        const int stg = pc % STG;
        const __nv_bfloat16(*Kp)[KP] = K16[stg];
        const __nv_bfloat16(*Vp)[KP] = V16[stg];
        const int page = cur.lo + (pc - base);
        const int rem = min(16, sl[cur.b] - page * 16);

        if (warp == 0) {
            float s0[4], s1[4];
#pragma unroll
            for (int j = 0; j < 4; j++) { s0[j] = 0.f; s1[j] = 0.f; }
#pragma unroll
            for (int ks = 0; ks < KS; ks++) {
                unsigned a0, a1, a2, a3, k0r, k1r, k2r, k3r;
                unsigned qa = pa_smem_u32(&Q16[cur_par][lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(a0, a1, a2, a3, qa);
                unsigned ka = pa_smem_u32(&Kp[lane % 16][ks * 16 + (lane / 16) * 8]);
                PA_LDM_X4(k0r, k1r, k2r, k3r, ka);
                PA_MMA(s0[0], s0[1], s0[2], s0[3], a0, a1, a2, a3, k0r, k2r);
                PA_MMA(s1[0], s1[1], s1[2], s1[3], a0, a1, a2, a3, k1r, k3r);
            }
            const int t0 = (lane % 4) * 2;
            const int tk[4] = {t0, t0 + 1, t0 + 8, t0 + 9};
            float sc[4] = {s0[0], s0[1], s1[0], s1[1]};
            float mx = -INFINITY;
#pragma unroll
            for (int j = 0; j < 4; j++) {
                sc[j] = tk[j] < rem ? sc[j] * scale : -INFINITY;
                mx = fmaxf(mx, sc[j]);
            }
            mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 1));
            mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 2));
            // e in the QUAD-max basis (not the lane-local max): P = e*rs_pg
            // must equal exp(sc - M); a lane-local basis silently drops the
            // exp(mxl - mx) factor, which only shows once logits spread.
            float e[4], sm = 0.f;
#pragma unroll
            for (int j = 0; j < 4; j++) { e[j] = __expf(sc[j] - mx); sm += e[j]; }
            sm += __shfl_xor_sync(0xffffffffu, sm, 1);
            sm += __shfl_xor_sync(0xffffffffu, sm, 2);

            const float M = fmaxf(m_r, mx);
            const float rs_old = __expf(m_r - M);
            const float rs_pg = __expf(mx - M);
            l_r = l_r * rs_old + sm * rs_pg;
            m_r = M;
            if (g < G) {
#pragma unroll
                for (int j = 0; j < 4; j++)
                    P16[g][tk[j]] = __float2bfloat16(e[j] * rs_pg);
                if ((lane & 3) == 0) rss[g] = rs_old;
            }
        }
        __syncthreads();

        {
            const float rs = g < G ? rss[g] : 0.f;
            unsigned a0, a1, a2, a3;
            unsigned pa = pa_smem_u32(&P16[lane % 16][(lane / 16) * 8]);
            PA_LDM_X4(a0, a1, a2, a3, pa);
#pragma unroll
            for (int w = 0; w < WPT; w++) {
                unsigned b0, b1;
                unsigned va = pa_smem_u32(&Vp[lane % 16][(warp * WPT + w) * 8]);
                PA_LDM_X2T(b0, b1, va);
#pragma unroll
                for (int j = 0; j < 4; j++) acc[w][j] *= rs;
                PA_MMA(acc[w][0], acc[w][1], acc[w][2], acc[w][3],
                       a0, a1, a2, a3, b0, b1);
            }
        }

        if (page + 1 >= cur.hi) {           // item boundary: publish + advance
            if (S == 1) {
                if (warp == 0 && (lane & 3) == 0 && g < G) {
                    mls[g][0] = m_r;
                    mls[g][1] = l_r;
                }
                __syncthreads();
                if (g < G) {
                    const float il = 1.f / mls[g][1];
                    __nv_bfloat16* o = out + ((size_t)cur.b * H + cur.h0 + g) * D;
#pragma unroll
                    for (int w = 0; w < WPT; w++) {
                        const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                        o[d0] = __float2bfloat16(acc[w][0] * il);
                        o[d0 + 1] = __float2bfloat16(acc[w][1] * il);
                    }
                }
                __syncthreads();
            } else {
                if (warp == 0 && (lane & 3) == 0 && g < G)
                    pml[((size_t)cur.b * H + cur.h0 + g) * S + cur.s] =
                        make_float2(m_r, l_r);
                if (g < G) {
                    float* dst =
                        pacc + (((size_t)cur.b * H + cur.h0 + g) * S + cur.s) * D;
#pragma unroll
                    for (int w = 0; w < WPT; w++) {
                        const int d0 = (warp * WPT + w) * 8 + (lane % 4) * 2;
                        dst[d0] = acc[w][0];
                        dst[d0 + 1] = acc[w][1];
                    }
                }
            }

            m_r = -INFINITY;
            l_r = 0.f;
#pragma unroll
            for (int w = 0; w < WPT; w++)
#pragma unroll
                for (int j = 0; j < 4; j++) acc[w][j] = 0.f;

            base = pc + 1;
            for (;;) {
                cur_id += gridDim.x;
                if (cur_id >= items) break;
                describe(cur_id, cur);
                if (cur.lo < cur.hi) break;
                neutral(cur);
            }
            cur_par ^= 1;
        }

        pc++;
        if (cur_id >= items && P < pc) break;
    }
    __pipeline_wait_prior(0);

    if (fuse && S > 1) {
        // fused cross-split reduction: after ALL of this CTA's partial
        // stores, fence once, bump ctr[pair] for every item it owned; the
        // last arriver per pair reduces all S splits inline (ctr self-resets
        // so the buffer stays valid across launches).
        __threadfence();
        int nmine = 0;
        if (tid == 0) {
            for (int id = blockIdx.x; id < items; id += gridDim.x) {
                It t;
                describe(id, t);
                const int pair = id % npairs;
                if (atomicAdd(&ctr[pair], 1) == S - 1) {
                    ctr[pair] = 0;                 // self-reset
                    mine[nmine & 127] = pair;
                    nmine++;
                }
            }
        }
        __shared__ int nmine_sh;
        __syncthreads();
        if (tid == 0) nmine_sh = nmine;
        __syncthreads();
        for (int k = 0; k < nmine_sh; k++) {
            const int pair = mine[k];
            const int b = pair / Hkv, kvh = pair - b * Hkv, h0 = kvh * G;
            for (int i = tid; i < G * D; i += THREADS) {
                const int h = i / D, d = i - h * D;
                const float2* mlh = pml + (size_t)(b * H + h0 + h) * S;
                float M = -INFINITY;
                for (int s2 = 0; s2 < S; s2++) M = fmaxf(M, mlh[s2].x);
                float lo = 0.f, ao = 0.f;
                const float* ah = pacc + ((size_t)(b * H + h0 + h) * S) * D + d;
                for (int s2 = 0; s2 < S; s2++) {
                    const float2 ml = mlh[s2];
                    const float w = __expf(ml.x - M);
                    lo = fmaf(w, ml.y, lo);
                    ao = fmaf(w, ah[(size_t)s2 * D], ao);
                }
                out[((size_t)b * H + h0 + h) * D + d] = __float2bfloat16(ao / lo);
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Generic fallback for unsupported (D, G, page_size): one block per (b, h),
// two passes over the sequence, fp32 math. Correct, not fast.
// ---------------------------------------------------------------------------
__global__ void attn_naive_kernel(const __nv_bfloat16* __restrict__ q,
                                  const __nv_bfloat16* __restrict__ kv,
                                  const int* __restrict__ bt,
                                  const int* __restrict__ sl,
                                  __nv_bfloat16* __restrict__ out, int Hkv,
                                  int H, int D, int mb, int P, float scale)
{
    extern __shared__ float sc[];  // seq scores
    const int b = blockIdx.x / H;
    const int h = blockIdx.x - b * H;
    const int kvh = h / (H / Hkv);
    const int Lb = sl[b];
    const int pages_b = (Lb + P - 1) / P;
    const size_t slab = (size_t)Hkv * 2 * D;
    const __nv_bfloat16* qp = q + ((size_t)b * H + h) * D;
    const int* btrow = bt + (size_t)b * mb;
    const int tid = threadIdx.x;

    for (int t = tid; t < Lb; t += blockDim.x) {
        int pg = t / P;
        const __nv_bfloat16* row =
            kv + ((size_t)btrow[pg] * P + (t - pg * P)) * slab + (size_t)kvh * 2 * D;
        float sacc = 0.f;
        for (int d = 0; d < D; d++) sacc = fmaf(__bfloat162float(qp[d]), __bfloat162float(row[d]), sacc);
        sc[t] = sacc * scale;
    }
    __syncthreads();
    // block max
    float M = -INFINITY;
    for (int t = tid; t < Lb; t += blockDim.x) M = fmaxf(M, sc[t]);
    // warp reduce then smem reduce
    for (int off = 16; off; off >>= 1) M = fmaxf(M, __shfl_xor_sync(0xffffffffu, M, off));
    __shared__ float red[32];
    if ((tid & 31) == 0) red[tid >> 5] = M;
    __syncthreads();
    if (tid < 32) {
        M = (tid < (blockDim.x + 31) / 32) ? red[tid] : -INFINITY;
        for (int off = 16; off; off >>= 1) M = fmaxf(M, __shfl_xor_sync(0xffffffffu, M, off));
        if (tid == 0) red[0] = M;
    }
    __syncthreads();
    M = red[0];
    float l = 0.f;
    for (int t = tid; t < Lb; t += blockDim.x) {
        float p = __expf(sc[t] - M);
        sc[t] = p;
        l += p;
    }
    // block sum
    for (int off = 16; off; off >>= 1) l += __shfl_xor_sync(0xffffffffu, l, off);
    __shared__ float red2[32];
    if ((tid & 31) == 0) red2[tid >> 5] = l;
    __syncthreads();
    if (tid < 32) {
        l = (tid < (blockDim.x + 31) / 32) ? red2[tid] : 0.f;
        for (int off = 16; off; off >>= 1) l += __shfl_xor_sync(0xffffffffu, l, off);
        if (tid == 0) red2[0] = l;
    }
    __syncthreads();
    l = red2[0];
    for (int d = tid; d < D; d += blockDim.x) {
        float o = 0.f;
        for (int t = 0; t < Lb; t++) {
            int pg = t / P;
            const __nv_bfloat16* row =
                kv + ((size_t)btrow[pg] * P + (t - pg * P)) * slab + (size_t)kvh * 2 * D + D;
            o = fmaf(sc[t], __bfloat162float(row[d]), o);
        }
        out[((size_t)b * H + h) * D + d] = __float2bfloat16(o / l);
    }
}

typedef void (*KFn)(const __nv_bfloat16*, const __nv_bfloat16*, const int*,
                     const int*, float*, float2*, __nv_bfloat16*, int*, int*,
                     int, int, int, int, int, float);
typedef void (*KFnF)(const __nv_bfloat16*, const __nv_bfloat16*, const int*,
                       const int*, float*, float2*, __nv_bfloat16*, int*, int*,
                       int, int, int, int, int, int, float);
static inline void go_f(KFnF kfn, dim3 grid, const void* q, const void* kv,
                        const void* bt, const void* sl, void* pacc, void* pml,
                        void* out, void* ctr, void* wq, int S, int mb, int Hkv,
                        int B, int items, int fuse, float scale,
                        cudaStream_t stream) {
    kfn<<<grid, 128, 0, stream>>>(
        (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
        (const int*)sl, (float*)pacc, (float2*)pml, (__nv_bfloat16*)out,
        (int*)ctr, (int*)wq, S, mb, Hkv, B, items, fuse, scale);
}
static inline void go(KFn kfn, dim3 grid, const void* q, const void* kv,
                      const void* bt, const void* sl, void* pacc, void* pml,
                      void* out, void* ctr, void* wq, int S, int mb, int Hkv,
                      int B, int items, float scale, cudaStream_t stream) {
    kfn<<<grid, 128, 0, stream>>>(
        (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
        (const int*)sl, (float*)pacc, (float2*)pml, (__nv_bfloat16*)out,
        (int*)ctr, (int*)wq, S, mb, Hkv, B, items, scale);
}

extern "C" void pa_run(const void* q, const void* kv, const void* bt,
                       const void* sl, void* out, void* pacc, void* pml,
                       void* ctr, void* wq, int B, int H, int Hkv, int D, int S,
                       int mb, int P, int stg, int merge, int pipe,
                       int fuse, double scale, cudaStream_t stream)
{
    int G = H / Hkv;
    bool fast = P == 16 && ((D == 128 && (G == 8 || G == 4)) || (D == 64 && G == 4));
    if (fast && merge && (Hkv % 2) == 0) {
        // kv-merged path: two kv heads per CTA
        const int items = S * B * (Hkv / 2);
        static int ncta2[2][2][2];  // [D128/64][G8/4][stg 3/4]
        int id = D == 128 ? 0 : 1, ig = G == 8 ? 0 : 1, is = stg >= 4 ? 1 : 0;
        const void* kfn2 = (D == 128)
                               ? (G == 8 ? (const void*)attn_split2_kernel<128, 8, 2>
                                         : (const void*)attn_split2_kernel<128, 4, 2>)
                               : (const void*)attn_split2_kernel<64, 4, 3>;
        if (!ncta2[id][ig][is]) {
            int sm = 0, per = 1;
            cudaDeviceGetAttribute(&sm, cudaDevAttrMultiProcessorCount, 0);
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per, kfn2, 128, 0);
            ncta2[id][ig][is] = sm * (per > 0 ? per : 1);
        }
        dim3 grid(items < ncta2[id][ig][is] ? items : ncta2[id][ig][is]);
        if (D == 128 && G == 8)
            attn_split2_kernel<128, 8, 2><<<grid, 128, 0, stream>>>(
                (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
                (const int*)sl, (float*)pacc, (float2*)pml, (__nv_bfloat16*)out,
                (int*)ctr, (int*)wq, S, mb, Hkv, B, items, (float)scale);
        else if (D == 128)
            attn_split2_kernel<128, 4, 2><<<grid, 128, 0, stream>>>(
                (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
                (const int*)sl, (float*)pacc, (float2*)pml, (__nv_bfloat16*)out,
                (int*)ctr, (int*)wq, S, mb, Hkv, B, items, (float)scale);
        else
            attn_split2_kernel<64, 4, 3><<<grid, 128, 0, stream>>>(
                (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
                (const int*)sl, (float*)pacc, (float2*)pml, (__nv_bfloat16*)out,
                (int*)ctr, (int*)wq, S, mb, Hkv, B, items, (float)scale);
        if (S > 1) {
            const int nth = B * H * (D >> 2);
            attn_combine_kernel<<<(nth + 127) / 128, 128, 0, stream>>>(
                (const float*)pacc, (const float2*)pml, (__nv_bfloat16*)out,
                S, B, H, D);
        }
        return;
    }
    if (fast && pipe) {
        const int items = S * B * Hkv;
        static int nctaP[4][3][3];  // [stg 2/3/4/6][D 128/64][G 8/4]
        const int is = stg == 2 ? 0 : stg == 4 ? 1 : stg == 6 ? 2 : 3;
        const int id = D == 128 ? 0 : 1, ig = G == 8 ? 0 : 1;
        KFnF kf;
        if (D == 128 && G == 8)
            kf = stg == 2 ? (KFnF)attn_pipe_kernel<128, 8, 2>
                : stg >= 4  ? (KFnF)attn_pipe_kernel<128, 8, 4>
                            : (KFnF)attn_pipe_kernel<128, 8, 3>;
        else if (D == 128)
            kf = stg == 2 ? (KFnF)attn_pipe_kernel<128, 4, 2>
                : stg >= 4  ? (KFnF)attn_pipe_kernel<128, 4, 4>
                            : (KFnF)attn_pipe_kernel<128, 4, 3>;
        else
            kf = stg == 2 ? (KFnF)attn_pipe_kernel<64, 4, 2>
                : stg == 4  ? (KFnF)attn_pipe_kernel<64, 4, 4>
                : stg == 6  ? (KFnF)attn_pipe_kernel<64, 4, 6>
                            : (KFnF)attn_pipe_kernel<64, 4, 3>;
        if (!nctaP[is][id][ig]) {
            int sm = 0, per = 1;
            cudaDeviceGetAttribute(&sm, cudaDevAttrMultiProcessorCount, 0);
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per, (const void*)kf,
                                                          128, 0);
            nctaP[is][id][ig] = sm * (per > 0 ? per : 1);
        }
        dim3 grid(items < nctaP[is][id][ig] ? items : nctaP[is][id][ig]);
        go_f(kf, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv,
             B, items, fuse, (float)scale, stream);
        if (S > 1 && !fuse) {
            const int nth = B * H * (D >> 2);
            attn_combine_kernel<<<(nth + 127) / 128, 128, 0, stream>>>(
                (const float*)pacc, (const float2*)pml, (__nv_bfloat16*)out,
                S, B, H, D);
        }
        return;
    }
    if (fast) {
        const int items = S * B * Hkv;
        const void* kfn;
        if (stg == 4)      kfn = (D == 128) ? (G == 8 ? (const void*)attn_split_kernel<128, 8, 4>
                                                      : (const void*)attn_split_kernel<128, 4, 4>)
                                            : (const void*)attn_split_kernel<64, 4, 4>;
        else               kfn = (D == 128) ? (G == 8 ? (const void*)attn_split_kernel<128, 8, 3>
                                                      : (const void*)attn_split_kernel<128, 4, 3>)
                                            : (const void*)attn_split_kernel<64, 4, 3>;
        // persistent grid: exactly the resident-CTA count, work stolen via wq
        static int ncta[3][3][3];  // [stg 3/4/6][D 128/64][G 8/4]
        int is = stg == 4 ? 1 : stg == 6 ? 2 : 0;
        int id = D == 128 ? 0 : 1, ig = G == 8 ? 0 : 1;
        if (!ncta[is][id][ig]) {
            int sm = 0, per = 1;
            cudaDeviceGetAttribute(&sm, cudaDevAttrMultiProcessorCount, 0);
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(&per, kfn, 128, 0);
            ncta[is][id][ig] = sm * (per > 0 ? per : 1);
        }
        dim3 grid(items < ncta[is][id][ig] ? items : ncta[is][id][ig]);
        if (D == 128 && G == 8) {
            if (stg == 4)      go((KFn)attn_split_kernel<128, 8, 4>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
            else               go((KFn)attn_split_kernel<128, 8, 3>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
        } else if (D == 128) {
            if (stg == 4)      go((KFn)attn_split_kernel<128, 4, 4>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
            else               go((KFn)attn_split_kernel<128, 4, 3>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
        } else {
            if (stg == 4)      go((KFn)attn_split_kernel<64, 4, 4>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
            else if (stg == 6) go((KFn)attn_split_kernel<64, 4, 6>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
            else               go((KFn)attn_split_kernel<64, 4, 3>, grid, q, kv, bt, sl, pacc, pml, out, ctr, wq, S, mb, Hkv, B, items, (float)scale, stream);
        }
        if (S > 1) {
            const int nth = B * H * (D >> 2);
            attn_combine_kernel<<<(nth + 127) / 128, 128, 0, stream>>>(
                (const float*)pacc, (const float2*)pml, (__nv_bfloat16*)out,
                S, B, H, D);
        }
        return;
    }
    // generic fallback: dynamic smem sized by the block table's max length
    int Lmax = mb * P;
    size_t smem = (size_t)Lmax * sizeof(float);
    attn_naive_kernel<<<B * H, 128, smem, stream>>>(
        (const __nv_bfloat16*)q, (const __nv_bfloat16*)kv, (const int*)bt,
        (const int*)sl, (__nv_bfloat16*)out, Hkv, H, D, mb, P, (float)scale);
}
"""

_CPP_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <unordered_map>
extern "C" void pa_run(const void* q, const void* kv, const void* bt,
                       const void* sl, void* out, void* pacc, void* pml,
                       void* ctr, void* wq, int B, int H, int Hkv, int D,
                       int S, int mb, int P, int stg, int merge, int pipe,
                       int fuse, double scale, cudaStream_t stream);

// Launch-cache: identical pointer+dims set -> replay the captured graph.
// Pointers identify the tensors, so replayed kernels always read the data
// currently living at those addresses (content changes are fine).
struct Key {
    const void *q, *kv, *bt, *sl, *out, *pacc, *pml, *ctr, *wq;
    int B, H, Hkv, D, S, mb, P, stg, merge, pipe, fuse;
    bool operator==(const Key& o) const {
        return q == o.q && kv == o.kv && bt == o.bt && sl == o.sl &&
               out == o.out && pacc == o.pacc && pml == o.pml && ctr == o.ctr &&
               wq == o.wq &&
               B == o.B && H == o.H && Hkv == o.Hkv && D == o.D && S == o.S &&
               mb == o.mb && P == o.P && stg == o.stg && merge == o.merge &&
               pipe == o.pipe && fuse == o.fuse;
    }
};
struct KHash {
    size_t operator()(const Key& k) const {
        size_t h = 0;
        auto mix = [&](size_t v) { h ^= v + 0x9e3779b97f4a7c15ULL + (h << 6) + (h >> 2); };
        mix((size_t)k.q); mix((size_t)k.kv); mix((size_t)k.bt); mix((size_t)k.sl);
        mix((size_t)k.out); mix((size_t)k.pacc); mix((size_t)k.pml); mix((size_t)k.ctr);
        mix((size_t)k.wq);
        mix((size_t)k.B); mix((size_t)k.H); mix((size_t)k.Hkv); mix((size_t)k.D);
        mix((size_t)k.S); mix((size_t)k.mb); mix((size_t)k.P); mix((size_t)k.stg);
        mix((size_t)k.merge); mix((size_t)k.pipe); mix((size_t)k.fuse);
        return h;
    }
};
static std::unordered_map<Key, cudaGraphExec_t, KHash> g_cache;
// fast-replay registry: int id -> graph exec (avoids re-parsing 9 tensor
// args on every call; the harness calls forward with the same tensors)
static std::unordered_map<int64_t, cudaGraphExec_t> g_fast;
static std::unordered_map<Key, int64_t, KHash> g_keyid;
static int64_t g_seq = 0;

static int64_t launch(const void* q, const void* kv, const void* bt,
                      const void* sl, void* out, void* pacc, void* pml,
                      void* ctr, void* wq, int B, int H, int Hkv, int D, int S,
                      int mb, int P, int stg, int merge, int pipe, int fuse,
                      double scale, cudaStream_t stream) {
    Key key{q, kv, bt, sl, out, pacc, pml, ctr, wq, B, H, Hkv, D, S, mb, P,
            stg, merge, pipe, fuse};
    auto it = g_cache.find(key);
    if (it != g_cache.end()) {
        cudaGraphLaunch(it->second, stream);
        auto i2 = g_keyid.find(key);
        if (i2 != g_keyid.end()) return i2->second;
        const int64_t id = ++g_seq;
        g_fast.emplace(id, it->second);
        g_keyid.emplace(key, id);
        return id;
    }
    cudaStreamCaptureStatus cap;
    if (cudaStreamIsCapturing(stream, &cap) != cudaSuccess ||
        cap != cudaStreamCaptureStatusNone) {  // caller capturing: just launch
        pa_run(q, kv, bt, sl, out, pacc, pml, ctr, wq, B, H, Hkv, D, S, mb, P,
                stg, merge, pipe, fuse, scale, stream);
        return 0;
    }
    // Capture on a private side stream: BeginCapture is rejected on legacy
    // default streams (torch's default current stream is one).
    static cudaStream_t cap_stream = nullptr;
    if (!cap_stream) cudaStreamCreateWithFlags(&cap_stream, cudaStreamNonBlocking);
    cudaGraph_t graph;
    if (cudaStreamBeginCapture(cap_stream, cudaStreamCaptureModeRelaxed) != cudaSuccess) {
        cudaGetLastError();  // swallow the sticky code
        pa_run(q, kv, bt, sl, out, pacc, pml, ctr, wq, B, H, Hkv, D, S, mb,
                P, stg, merge, pipe, fuse, scale, stream);
        return 0;
    }
    pa_run(q, kv, bt, sl, out, pacc, pml, ctr, wq, B, H, Hkv, D, S, mb, P,
            stg, merge, pipe, fuse, scale, cap_stream);
    if (cudaStreamEndCapture(cap_stream, &graph) != cudaSuccess || !graph) {
        cudaGetLastError();
        return 0;
    }
    cudaGraphExec_t exec = nullptr;
    if (cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0) != cudaSuccess || !exec) {
        cudaGraphDestroy(graph);
        return 0;
    }
    cudaGraphDestroy(graph);
    if (g_cache.size() > 128) {
        for (auto& kv2 : g_cache) cudaGraphExecDestroy(kv2.second);
        g_cache.clear();
        for (auto& kv2 : g_fast) cudaGraphExecDestroy(kv2.second);
        g_fast.clear();
        g_keyid.clear();
    }
    g_cache.emplace(key, exec);
    const int64_t id = ++g_seq;
    g_fast.emplace(id, exec);
    g_keyid.emplace(key, id);
    cudaGraphLaunch(exec, stream);
    return id;
}

int64_t run(torch::Tensor q, torch::Tensor kv, torch::Tensor bt, torch::Tensor sl,
         torch::Tensor out, torch::Tensor pacc, torch::Tensor pml,
         torch::Tensor ctr, torch::Tensor wq,
         int64_t B, int64_t H, int64_t Hkv, int64_t D, int64_t S, int64_t mb,
         int64_t P, int64_t stg, int64_t merge, int64_t pipe,
         int64_t fuse, double scale) {
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    return launch(q.data_ptr(), kv.data_ptr(), bt.data_ptr(), sl.data_ptr(),
           out.data_ptr(),
           pacc.defined() && pacc.numel() ? pacc.data_ptr() : (void*)q.data_ptr(),
           pml.defined() && pml.numel() ? pml.data_ptr() : (void*)q.data_ptr(),
           ctr.defined() && ctr.numel() ? ctr.data_ptr() : (void*)q.data_ptr(),
           wq.data_ptr(),
           (int)B, (int)H, (int)Hkv, (int)D, (int)S, (int)mb, (int)P,
           (int)stg, (int)merge, (int)pipe, (int)fuse, scale, stream);
}

int64_t replay(int64_t id) {
    auto it = g_fast.find(id);
    if (it == g_fast.end()) return 0;
    cudaGraphLaunch(it->second, at::cuda::getCurrentCUDAStream());
    return 1;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("run", &run);
    m.def("replay", &replay);
}
"""

_BUILD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_pa_build")
_EXT = None


def _build():
    global _EXT
    if _EXT is not None:
        return _EXT
    os.makedirs(_BUILD_DIR, exist_ok=True)
    tag = hashlib.sha256((_CUDA_SRC + _CPP_SRC).encode()).hexdigest()[:12]
    cu = os.path.join(_BUILD_DIR, "pa_kernel.cu")
    cp = os.path.join(_BUILD_DIR, "pa_bind.cpp")
    for path, src in ((cu, _CUDA_SRC), (cp, _CPP_SRC)):
        try:
            if open(path).read() == src:
                continue
        except OSError:
            pass
        with open(path, "w") as f:
            f.write(src)
    _EXT = load(
        name=f"pa_ext_m_{tag}",
        sources=[cu, cp],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
        verbose=False,
    )
    return _EXT


# --- Shape knobs (overridden by check.py / benchmark.py from shapes.py) ----
BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16


class Model(nn.Module):
    def __init__(
        self,
        batch: int,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        seq_len: int,
        page_size: int,
    ):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.batch = batch
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.seq_len = seq_len
        self.page_size = page_size
        self.group_size = num_heads // num_kv_heads
        self.scale = 1.0 / math.sqrt(head_dim)
        self.ext = _build()

        pages = (seq_len + page_size - 1) // page_size
        pairs = batch * num_kv_heads
        # (S, STG, FUSE) tuned per benchmark shape on the RTX PRO 6000 with
        # the harness timing pattern (median of 30, 128 MB dirty L2 flush).
        # S = sequence splits, STG = cp.async page stages per CTA,
        # FUSE = 1 reduces the S splits inside the kernel, 0 uses a wide
        # separate combine kernel.
        tuned = {
            (8, 32, 8, 128, 1024): (4, 4, 0),
            (32, 32, 8, 128, 2048): (4, 3, 1),
            (4, 64, 8, 128, 4096): (4, 3, 0),
            (16, 32, 8, 128, 1535): (2, 3, 1),
            (8, 16, 4, 64, 2000): (8, 4, 0),
        }
        t = tuned.get((batch, num_heads, num_kv_heads, head_dim, seq_len))
        if t is None:
            # generic: enough CTAs to fill the GPU, else split the sequence
            self.S = max(1, min(pages, round(384 / pairs)))
            if pages * pairs <= 376:
                self.S = 1
            self.STG = 3
            self.FUSE = 1
        else:
            self.S, self.STG, self.FUSE = t
        self.MERGE = 0  # unused fast-path variant (two kv heads per CTA)
        self.PIPE = 1   # cross-item pipelined persistent kernel
        self._kid = 0
        self._lq = self._lk = self._lb = self._ls = None

        self.register_buffer(
            "out_b", torch.empty(batch, num_heads, head_dim, dtype=torch.bfloat16),
            persistent=False,
        )
        if self.S > 1:
            self.register_buffer(
                "pacc",
                torch.empty(self.S, batch, num_heads, head_dim, dtype=torch.float32),
                persistent=False,
            )
            self.register_buffer(
                "pml", torch.empty(self.S, batch * num_heads, 2, dtype=torch.float32),
                persistent=False,
            )
        else:
            self.register_buffer("pacc", torch.empty(0), persistent=False)
            self.register_buffer("pml", torch.empty(0), persistent=False)
        if self.S > 1:
            self.register_buffer(
                "ctr",
                torch.zeros(batch * num_kv_heads, dtype=torch.int32),
                persistent=False,
            )
        else:
            self.register_buffer("ctr", torch.empty(0), persistent=False)

        # persistent work queue: [0] next item claim, [1] exited CTAs
        self.register_buffer(
            "wq", torch.zeros(2, dtype=torch.int32), persistent=False
        )

        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)

    def forward(self, query, kv_cache, block_table, seq_lens):
        # fast path: the harness times repeated calls with the SAME tensors,
        # and pybind re-parsing nine tensor args costs ~8 us of CPU that the
        # GPU sits idle for on short kernels. First call captures a graph and
        # returns its id; after that a single-int replay launches it.
        if (
            self._kid
            and query is self._lq
            and kv_cache is self._lk
            and block_table is self._lb
            and seq_lens is self._ls
        ):
            self.ext.replay(self._kid)
            return self.out_b
        kid = self.ext.run(
            query,
            kv_cache,
            block_table,
            seq_lens,
            self.out_b,
            self.pacc,
            self.pml,
            self.ctr,
            self.wq,
            self.batch,
            self.num_heads,
            self.num_kv_heads,
            self.head_dim,
            self.S,
            block_table.size(1),
            self.page_size,
            self.STG,
            self.MERGE,
            self.PIPE,
            self.FUSE,
            self.scale,
        )
        if kid:
            self._kid = kid
            self._lq, self._lk = query, kv_cache
            self._lb, self._ls = block_table, seq_lens
        return self.out_b


def get_inputs():
    B = BATCH
    H = NUM_HEADS
    Hkv = NUM_KV_HEADS
    D = HEAD_DIM
    L = SEQ_LEN
    P = PAGE_SIZE

    pages_per_seq = (L + P - 1) // P
    total_pages = max(B * pages_per_seq + 8, 64)

    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1

    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
    block_table = perm.contiguous()
    seq_lens = torch.full((B,), L, dtype=torch.int32)

    return [query, kv_cache, block_table, seq_lens]


def get_init_inputs():
    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]