"""Paged attention decode kernel for H100 PCIe (SM90, HBM2e ~2.0 TB/s).

Single-query (decode) attention over a paged KV cache. One fused CUDA kernel:

- Grid (splits, B*Hkv): each CTA owns one (batch, kv-head, token-chunk) triple and
  streams its KV pages straight from the pool with 16B cp.async into an NSTAGES-deep
  smem pipeline (each stage = 32 tokens / two pages, smem rows padded +16B so
  ldmatrix reads are bank-conflict free).
- Attention on tensor cores: warp handles 8 tokens/stage for the whole GQA group
  (pad heads to the m16 tile). QK^T via m16n8k16 bf16 MMA, online softmax in the
  MMA fragment layout (row max/sum via 2 shfl in the 4-lane row groups), p
  re-packed in registers to an m16k8 A-fragment feeding the m16n8k8 P·V MMA.
- Warps keep independent (m, l, acc) online-softmax states, merged in smem (fp32)
  once at the end of the chunk.
- Split-K for occupancy on small grids: per-CTA partial (acc, m, l) goes to a
  scratch workspace; the last-arriving CTA of each (b, kvh) group does the final
  log-sum-exp merge in-kernel (arrival ticket; every thread fences before it).
- kv loads carry an L2::evict_first hint (streamed exactly once).

Per-row variable seq_lens, non-page-aligned tails, GQA ratios 1..8, head_dim 64/128.
"""

import math
import os

import torch
import torch.nn as nn

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["H100"]

# --- Shape knobs (same contract as reference.py; check.py overrides reference's).
BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/EmptyTensor.h>
#include <cuda_bf16.h>
#include <cstdint>

#define DEV_INLINE __device__ __forceinline__
#define FULL_MASK 0xffffffffu

DEV_INLINE unsigned long long mk_policy() {
  unsigned long long policy;
  asm volatile("createpolicy.fractional.L2::evict_first.b64 %0;" : "=l"(policy));
  return policy;
}
DEV_INLINE void cp_async16_ef(void* smem_dst, const void* gsrc, unsigned long long policy) {
  unsigned dst = (unsigned)__cvta_generic_to_shared(smem_dst);
  asm volatile("cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, 16, %2;\n" ::"r"(dst),
               "l"(gsrc), "l"(policy));
}
DEV_INLINE void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
DEV_INLINE void cp_wait_all_but() {
  asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
}
DEV_INLINE void ldm_x4(uint32_t* R, unsigned sa) {
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
               : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) : "r"(sa));
}
DEV_INLINE void ldm_x4_t(uint32_t* R, unsigned sa) {
  asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
               : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) : "r"(sa));
}
DEV_INLINE void mma_qk(float* C, const uint32_t* A, const uint32_t* B) {
  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"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3])
      : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]));
}
DEV_INLINE void mma_pv(float* C, const uint32_t* A, const uint32_t B) {
  asm volatile(
      "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 "
      "{%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};\n"
      : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3])
      : "r"(A[0]), "r"(A[1]), "r"(B));
}

struct MmaParams {
  const __nv_bfloat16* __restrict__ q;
  const __nv_bfloat16* __restrict__ kv;
  const int* __restrict__ bt;
  const int* __restrict__ seq;
  __nv_bfloat16* __restrict__ out;
  float* __restrict__ ws_acc;      // (groups, S, G, D) fp32
  float* __restrict__ ws_ml;       // (groups, S, G, 2) fp32
  unsigned* __restrict__ counters; // (groups)
  int H;
  int Hkv;
  int G;
  int S;
  int bt_stride;
  int page_stride_elems;  // P*Hkv*2D
  int tok_stride_elems;   // Hkv*2D
  int pages_per_split;
  float scale_log2e;
};

constexpr float NEG_BIG = -1e30f;

// D: head dim (64/128). 4 warps. TW=8 tokens/warp/stage -> T=32 tokens/stage.
template <int D, int NSTAGES, bool DIRECT>
__global__ void __launch_bounds__(128) dec_mma(const __grid_constant__ MmaParams p) {
  constexpr int NW = 4;
  constexpr int TW = 8;
  constexpr int T = NW * TW;
  constexpr int ROWE = 2 * D;
  constexpr int ROWEP = ROWE + 8;  // +16B pad: kills ldmatrix bank conflicts
  constexpr int SLICES = D / 16;
  constexpr int NTILES = D / 8;
  constexpr int STAGE_ELEMS = T * ROWEP;

  extern __shared__ char smem_raw[];
  __nv_bfloat16* s_kv = reinterpret_cast<__nv_bfloat16*>(smem_raw);
  int* s_bt = reinterpret_cast<int*>(s_kv + NSTAGES * STAGE_ELEMS);

  const int split = blockIdx.x;
  const int gid = blockIdx.y;
  const int b = gid / p.Hkv;
  const int kvh = gid % p.Hkv;
  const int kv_len = p.seq[b];
  const int npages = (kv_len + 15) >> 4;
  const int pg0 = split * p.pages_per_split;
  const int pg1 = min(npages, pg0 + p.pages_per_split);
  const int tok0 = pg0 << 4;
  const int tok_end = min(kv_len, pg1 << 4);

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

  const int nchunk_pages = max(0, pg1 - pg0);
  for (int i = tid; i < nchunk_pages; i += NW * 32) {
    s_bt[i] = p.bt[b * p.bt_stride + pg0 + i];
  }

  const int total_tok = tok_end - tok0;
  const int total_stages = (total_tok + T - 1) / T;

  constexpr int LPR = ROWE * 2 / 16;
  constexpr int NITER = (T * LPR) / 128;
  const int row0 = tid / LPR;
  const int chunk = tid % LPR;
  const int elem_off = chunk * 8;

  unsigned long long policy = mk_policy();
  auto issue_stage = [&](int si, int slot) {
    __nv_bfloat16* dst0 = s_kv + slot * STAGE_ELEMS;
#pragma unroll
    for (int k = 0; k < NITER; ++k) {
      int row = row0 + k * (128 / LPR);
      int pg = (si * T + row) >> 4;
      int tokinpg = row & 15;
      bool valid = (si * T + row) < total_tok;
      int page = valid ? s_bt[pg] : 0;
      const __nv_bfloat16* src = p.kv + (unsigned)page * p.page_stride_elems +
                                 (unsigned)(tokinpg * p.tok_stride_elems) +
                                 (unsigned)(kvh * ROWE) + elem_off;
      cp_async16_ef(dst0 + row * ROWEP + chunk * 8, valid ? src : p.kv, policy);
    }
    cp_commit();
  };

  // ---- Q A-frags (m16k16 per slice): rows 0..G-1 real, else zero ----
  __syncthreads();  // publish s_bt before any issue
  // ALWAYS commit exactly NSTAGES-1 groups at preload (empty commits included) so
  // wait_group<NSTAGES-2> semantics hold even when total_stages < NSTAGES-1.
  for (int si = 0; si < NSTAGES - 1; ++si) {
    if (si < total_stages) issue_stage(si, si);
    else cp_commit();
  }
  uint32_t aQ[SLICES][4];
  {
    const int r0 = lane / 4;
    const int c0 = (lane % 4) * 2;
#pragma unroll
    for (int s = 0; s < SLICES; ++s) {
      uint32_t a[4] = {0u, 0u, 0u, 0u};
      if (r0 < p.G) {
        const __nv_bfloat16* qp = p.q + ((long)b * p.H + kvh * p.G + r0) * D + s * 16 + c0;
        a[0] = *reinterpret_cast<const uint32_t*>(qp);
        a[2] = *reinterpret_cast<const uint32_t*>(qp + 8);
      }
      if (r0 + 8 < p.G) {
        const __nv_bfloat16* qp = p.q + ((long)b * p.H + kvh * p.G + r0 + 8) * D + s * 16 + c0;
        a[1] = *reinterpret_cast<const uint32_t*>(qp);
        a[3] = *reinterpret_cast<const uint32_t*>(qp + 8);
      }
      const float sc = p.scale_log2e;
#pragma unroll
      for (int i = 0; i < 4; ++i) {
        float2 f = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&a[i]));
        f.x *= sc;
        f.y *= sc;
        __nv_bfloat162 packed = __floats2bfloat162_rn(f.x, f.y);
        a[i] = *reinterpret_cast<uint32_t*>(&packed);
      }
#pragma unroll
      for (int i = 0; i < 4; ++i) aQ[s][i] = a[i];
    }
  }
  float oacc[NTILES][4];
  float m0 = NEG_BIG, m8 = NEG_BIG, l0 = 0.f, l8 = 0.f;
#pragma unroll
  for (int t = 0; t < NTILES; ++t) {
#pragma unroll
    for (int i = 0; i < 4; ++i) oacc[t][i] = 0.f;
  }

  const int r0 = lane / 4;
  const int c0 = (lane % 4) * 2;

  for (int si = 0; si < total_stages; ++si) {
    cp_wait_all_but<NSTAGES - 2>();
    __syncthreads();
    int ni0 = si + NSTAGES - 1;
    if (ni0 < total_stages) {
      issue_stage(ni0, ni0 % NSTAGES);
    } else {
      cp_commit();
    }
    const int slot = si % NSTAGES;
    const int tok_base = tok0 + si * T;
    const __nv_bfloat16* sbase = s_kv + slot * STAGE_ELEMS;
    (void)tok_base;

    // ---------- QK ----------
    float accs[4] = {0.f, 0.f, 0.f, 0.f};
    {
      const int trow = warp * TW;
      int tok = lane % 8;
      const __nv_bfloat16* tbase = sbase + (trow + tok) * ROWEP;
#pragma unroll
      for (int sp = 0; sp < SLICES / 2; ++sp) {
        uint32_t b[4];
        int j = lane / 8;
        unsigned sa = (unsigned)__cvta_generic_to_shared(tbase + (sp * 4 + j) * 8);
        ldm_x4(b, sa);
        mma_qk(accs, aQ[2 * sp], b);
        mma_qk(accs, aQ[2 * sp + 1], b + 2);
      }
    }
    // mask invalid tokens: cols c0, c0+1 of rows r0, r0+8
    {
      int tloc = si * T + warp * TW + c0;
      if (tloc >= total_tok) { accs[0] = NEG_BIG; accs[2] = NEG_BIG; }
      if (tloc + 1 >= total_tok) { accs[1] = NEG_BIG; accs[3] = NEG_BIG; }
    }
    // row max
    float m0n = fmaxf(fmaxf(accs[0], accs[1]), m0);
    float m8n = fmaxf(fmaxf(accs[2], accs[3]), m8);
#pragma unroll
    for (int off = 1; off <= 2; off <<= 1) {
      m0n = fmaxf(m0n, __shfl_xor_sync(FULL_MASK, m0n, off));
      m8n = fmaxf(m8n, __shfl_xor_sync(FULL_MASK, m8n, off));
    }
    m0n = fmaxf(m0n, m0);
    m8n = fmaxf(m8n, m8);
    float alpha0 = exp2f(m0 - m0n);
    float alpha8 = exp2f(m8 - m8n);
    m0 = m0n; m8 = m8n;
    l0 *= alpha0; l8 *= alpha8;
#pragma unroll
    for (int t = 0; t < NTILES; ++t) {
      oacc[t][0] *= alpha0; oacc[t][1] *= alpha0;
      oacc[t][2] *= alpha8; oacc[t][3] *= alpha8;
    }
    float p0 = exp2f(accs[0] - m0);
    float p1 = exp2f(accs[1] - m0);
    float p2 = exp2f(accs[2] - m8);
    float p3 = exp2f(accs[3] - m8);
    float ls0 = (p0 + p1), ls8 = (p2 + p3);
#pragma unroll
    for (int off = 1; off <= 2; off <<= 1) {
      ls0 += __shfl_xor_sync(FULL_MASK, ls0, off);
      ls8 += __shfl_xor_sync(FULL_MASK, ls8, off);
    }
    l0 += ls0; l8 += ls8;
    // pack p into A frag (m16k8)
    uint32_t aP[2];
    {
      __nv_bfloat162 pa = __floats2bfloat162_rn(p0, p1);
      __nv_bfloat162 pb = __floats2bfloat162_rn(p2, p3);
      aP[0] = *reinterpret_cast<uint32_t*>(&pa);
      aP[1] = *reinterpret_cast<uint32_t*>(&pb);
    }
    // ---------- PV ----------
    {
      int tok = lane % 8;
#pragma unroll
      for (int n4 = 0; n4 < NTILES / 4; ++n4) {
        uint32_t bv[4];
        int j = lane / 8;
        int chunk = D / 8 + n4 * 4 + j;
        unsigned sa = (unsigned)__cvta_generic_to_shared(sbase + (warp * TW + tok) * ROWEP + chunk * 8);
        ldm_x4_t(bv, sa);
#pragma unroll
        for (int nt = 0; nt < 4; ++nt) {
          mma_pv(oacc[n4 * 4 + nt], aP, bv[nt]);
        }
      }
    }
  }

  // barrier before repurposing smem for the merge buffers (lagging warps may still be
  // ldmatrix-ing the last stages' KV rows, which alias s_o/s_ml).
  cp_wait_all_but<0>();
  __syncthreads();

  // ---------- cross-warp merge ----------
  float* s_o = reinterpret_cast<float*>(smem_raw);
  float* s_ml = reinterpret_cast<float*>(smem_raw + NW * 16 * D * 4);
#pragma unroll
  for (int t = 0; t < NTILES; ++t) {
    float* dst = s_o + (warp * 16 + r0) * D + t * 8 + c0;
    dst[0] = oacc[t][0];
    dst[1] = oacc[t][1];
    dst = s_o + (warp * 16 + r0 + 8) * D + t * 8 + c0;
    dst[0] = oacc[t][2];
    dst[1] = oacc[t][3];
  }
  if ((lane & 3) == 0) {
    float* w = s_ml + (warp * 16 + r0) * 2;
    w[0] = m0; w[1] = l0;
    w[8 * 2] = m8; w[8 * 2 + 1] = l8;
  }
  __syncthreads();

  const int G = p.G;
  for (int item = tid; item < G * (D / 8); item += 128) {
    int r = item / (D / 8);
    int dc = (item % (D / 8)) * 8;
    float mw[NW], lw[NW];
    float M = NEG_BIG;
#pragma unroll
    for (int w = 0; w < NW; ++w) {
      mw[w] = s_ml[(w * 16 + r) * 2];
      lw[w] = s_ml[(w * 16 + r) * 2 + 1];
      M = fmaxf(M, mw[w]);
    }
    float ltot = 0.f;
    float acc[8];
#pragma unroll
    for (int e = 0; e < 8; ++e) acc[e] = 0.f;
    for (int w = 0; w < NW; ++w) {
      float aw = exp2f(mw[w] - M);
      ltot += lw[w] * aw;
      const float* ow = s_o + (w * 16 + r) * D + dc;
#pragma unroll
      for (int e = 0; e < 8; ++e) acc[e] += aw * ow[e];
    }
    int h = kvh * G + r;
    if constexpr (DIRECT) {
      float inv = ltot > 0.f ? 1.f / ltot : 0.f;
      __nv_bfloat162 outv[4];
#pragma unroll
      for (int e = 0; e < 4; ++e) {
        outv[e] = __floats2bfloat162_rn(acc[2 * e] * inv, acc[2 * e + 1] * inv);
      }
      *reinterpret_cast<uint4*>(p.out + ((long)b * p.H + h) * D + dc) =
          *reinterpret_cast<uint4*>(outv);
    } else {
      float* wa = p.ws_acc + (((long)gid * p.S + split) * G + r) * D + dc;
#pragma unroll
      for (int e = 0; e < 2; ++e) {
        *reinterpret_cast<float4*>(wa + 4 * e) =
            make_float4(acc[4 * e], acc[4 * e + 1], acc[4 * e + 2], acc[4 * e + 3]);
      }
      if (dc == 0) {
        float* wm = p.ws_ml + (((long)gid * p.S + split) * G + r) * 2;
        wm[0] = M;
        wm[1] = ltot;
      }
    }
  }

  if constexpr (!DIRECT) {
    __threadfence();  // every thread orders its own partial writes before the ticket
    __syncthreads();
    __shared__ unsigned s_ticket;
    if (tid == 0) {
      s_ticket = atomicAdd(&p.counters[gid], 1u);
    }
    __syncthreads();
    if (s_ticket != (unsigned)(p.S - 1)) return;
    if (tid == 0) p.counters[gid] = 0u;
    asm volatile("fence.acquire.gpu;" ::: "memory");

    // final log-sum-exp merge across splits
    for (int item = tid; item < G * (D / 8); item += 128) {
      int r = item / (D / 8);
      int dc = (item % (D / 8)) * 8;
      float M = NEG_BIG;
      for (int si = 0; si < p.S; ++si) {
        const float* wm = p.ws_ml + (((long)gid * p.S + si) * G + r) * 2;
        M = fmaxf(M, wm[0]);
      }
      float ltot = 0.f;
      float acc[8];
#pragma unroll
      for (int e = 0; e < 8; ++e) acc[e] = 0.f;
      for (int si = 0; si < p.S; ++si) {
        const float* wm = p.ws_ml + (((long)gid * p.S + si) * G + r) * 2;
        float aw = exp2f(wm[0] - M);
        ltot += wm[1] * aw;
        const float* wa = p.ws_acc + (((long)gid * p.S + si) * G + r) * D + dc;
#pragma unroll
        for (int e = 0; e < 8; ++e) acc[e] += aw * wa[e];
      }
      int h = kvh * G + r;
      float inv = ltot > 0.f ? 1.f / ltot : 0.f;
      __nv_bfloat162 outv[4];
#pragma unroll
      for (int e = 0; e < 4; ++e) {
        outv[e] = __floats2bfloat162_rn(acc[2 * e] * inv, acc[2 * e + 1] * inv);
      }
      *reinterpret_cast<uint4*>(p.out + ((long)b * p.H + h) * D + dc) =
          *reinterpret_cast<uint4*>(outv);
    }
  }
}

// ---------------- host ----------------

static inline int cdiv2(int a, int b) { return (a + b - 1) / b; }

torch::Tensor paged_decode(torch::Tensor q, torch::Tensor kv, torch::Tensor bt,
                           torch::Tensor seq) {
  const int B = q.size(0);
  const int H = q.size(1);
  const int D = q.size(2);
  const int Hkv = kv.size(2);
  const int P = kv.size(1);
  TORCH_CHECK(P == 16, "page size must be 16");
  const int G = H / Hkv;
  const int max_blocks = bt.size(1);
  const int groups = B * Hkv;

  auto out = at::empty({B, H, D}, q.options());
  if (B == 0 || H == 0) return out;
  TORCH_CHECK(G <= 8, "G>8 unsupported");
  TORCH_CHECK(D == 128 || D == 64, "D must be 64/128");

  // --- tuned (nstages, splits) per shape, heuristic fallback ---
  int nst = 3;
  int S = 1;
  auto setcfg = [&](int n, int s) { nst = n; S = s; };
  if (D == 128) {
    if (groups >= 96) setcfg(3, 1);
    else if (groups >= 48) setcfg(4, 1);
    else setcfg(3, std::max(2, std::min(96 / groups, max_blocks / 24)));
  } else {
    setcfg(4, groups >= 48 ? 1 : std::max(2, std::min(128 / groups, max_blocks / 24)));
  }
  S = std::min(std::max(S, 1), 32);
  S = std::min(S, max_blocks);
  const int pps = cdiv2(max_blocks, S);

  float* ws_acc_ptr = nullptr;
  float* ws_ml_ptr = nullptr;
  torch::Tensor ws;
  static torch::Tensor counters;
  if (S > 1) {
    long acc_elems = (long)groups * S * G * D;
    long ml_elems = (long)groups * S * G * 2;
    ws = at::empty({acc_elems + ml_elems}, q.options().dtype(at::kFloat));
    ws_acc_ptr = ws.data_ptr<float>();
    ws_ml_ptr = ws_acc_ptr + acc_elems;
    if (!counters.defined() || counters.numel() < groups) {
      counters = at::zeros({std::max(groups, 4096)}, q.options().dtype(at::kInt));
    }
  }

  MmaParams p;
  p.q = reinterpret_cast<const __nv_bfloat16*>(q.data_ptr());
  p.kv = reinterpret_cast<const __nv_bfloat16*>(kv.data_ptr());
  p.bt = bt.data_ptr<int>();
  p.seq = seq.data_ptr<int>();
  p.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr());
  p.ws_acc = ws_acc_ptr;
  p.ws_ml = ws_ml_ptr;
  p.counters = S > 1 ? reinterpret_cast<unsigned*>(counters.data_ptr<int>()) : nullptr;
  p.H = H;
  p.Hkv = Hkv;
  p.G = G;
  p.S = S;
  p.bt_stride = max_blocks;
  p.page_stride_elems = 16 * Hkv * 2 * D;
  p.tok_stride_elems = Hkv * 2 * D;
  p.pages_per_split = pps;
  p.scale_log2e = (float)(1.0 / sqrt((double)D) * 1.4426950408889634);

  dim3 grid(S, groups);
  auto stream = at::cuda::getCurrentCUDAStream();
  const int T = 32;
  int smem_stage = nst * T * (2 * D + 8) * 2;
  int smem_merge = 4 * 16 * (D + 2) * 4;
  int sm = std::max(smem_stage, smem_merge) + pps * 4;

  // per-instantiation smem-attribute flags: bit i = (D64:4|D128:0) + nst(3:0|4:1)*2 + direct(0|1)
  static int attr_done = 0;
  auto launch = [&](auto kern, int bit) {
    if (sm > 48 * 1024 && !(attr_done & (1 << bit))) {
      cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, 227 * 1024 - 1024);
      attr_done |= (1 << bit);
    }
    kern<<<grid, dim3(128), sm, stream>>>(p);
  };

  if (D == 128) {
    if (nst == 4) {
      if (S == 1) launch(dec_mma<128, 4, true>, 3);
      else launch(dec_mma<128, 4, false>, 2);
    } else {
      if (S == 1) launch(dec_mma<128, 3, true>, 1);
      else launch(dec_mma<128, 3, false>, 0);
    }
  } else {
    if (S == 1) launch(dec_mma<64, 4, true>, 7);
    else launch(dec_mma<64, 4, false>, 6);
  }
  cudaError_t err = cudaGetLastError();
  TORCH_CHECK(err == cudaSuccess, "launch failed: ", cudaGetErrorString(err));
  return out;
}

// -------- lean python binding (METH_FASTCALL) --------
#include <torch/csrc/autograd/python_variable.h>
static PyObject* fast_decode(PyObject*, PyObject* const* args, Py_ssize_t nargs) {
  try {
    const at::Tensor& q = THPVariable_Unpack(args[0]);
    const at::Tensor& kv = THPVariable_Unpack(args[1]);
    const at::Tensor& bt = THPVariable_Unpack(args[2]);
    const at::Tensor& seq = THPVariable_Unpack(args[3]);
    at::Tensor out = paged_decode(q, kv, bt, seq);
    return THPVariable_Wrap(std::move(out));
  } catch (const std::exception& e) {
    PyErr_SetString(PyExc_RuntimeError, e.what());
    return nullptr;
  }
}
static PyMethodDef fast_methods[] = {
    {"decode", (PyCFunction)(void*)fast_decode, METH_FASTCALL, nullptr},
    {nullptr, nullptr, 0, nullptr}};
void register_fast(pybind11::module_& m) { PyModule_AddFunctions(m.ptr(), fast_methods); }
"""

_CPP_SRC = r"""
torch::Tensor paged_decode(torch::Tensor q, torch::Tensor kv, torch::Tensor bt, torch::Tensor seq);
void register_fast(pybind11::module_& m);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("paged_decode", &paged_decode);
  register_fast(m);
}
"""

_ext = None


def _get_ext():
    global _ext
    if _ext is None:
        os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0a"
        from torch.utils.cpp_extension import load_inline

        _ext = load_inline(
            name="paged_attention_decode_ext",
            cpp_sources=[_CPP_SRC],
            cuda_sources=[_CUDA_SRC],
            functions=None,
            extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
            verbose=False,
        )
    return _ext


class Model(nn.Module):
    """Single-query paged attention decode (drop-in for reference.Model)."""

    def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size):
        super().__init__()
        assert num_heads % num_kv_heads == 0, "num_heads must be a multiple of num_kv_heads (GQA)"
        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.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)
        self._fn = _get_ext().decode

    def forward(self, query, kv_cache, block_table, seq_lens):
        return self._fn(query, kv_cache, block_table, seq_lens)


def get_inputs():
    B, H, Hkv, D, L, P = BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, 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]