KernelBench hard · RTX PRO 6000
Paged Attention Kimi K3 (256k)
48.5%geomean peak fraction across shapes
manually audited: clean
Genuine custom CUDA flash-decoding paged attention. Canonical inputs take a fused kernel that reads live query, KV pages, block_table, and seq_lens; gathers logical pages from the randomized physical cache; computes QK and PV with bf16 mma.sync tensor-core instructions; maintains fp32 online softmax; and merges sequence-chunk partials through a self-resetting atomic ticket before writing a fresh bf16 output. A CUDA-core fused variant and a live PyTorch fallback cover other supported layouts. Persistent tensors are non-returned partial/counter workspaces keyed only by public dimensions, not cached outputs. None of the forbidden attention libraries or SDPA forms occurs. The strong 0.4855 score is physically plausible and exactly matches the five logged bandwidth fractions.
harnesskinetic-claudeagent session3h 2mtotal wall3h 3mcheck58sbenchmark6soutput tokens341,685cost$288.47gpu-lock wait6mgpu-lock held69sregimememory
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
8×32×8×128×1024×160.039 ms47.5%0.85 TB/s · 47% of 1.8 TB/s HBM · also 3 TFLOPS (1% of compute)
32×32×8×128×2048×160.205 ms72.8%1.31 TB/s · 73% of 1.8 TB/s HBM · also 5 TFLOPS (1% of compute)
4×64×8×128×4096×160.097 ms38.5%0.69 TB/s · 39% of 1.8 TB/s HBM · also 6 TFLOPS (1% of compute)
16×32×8×128×1535×160.094 ms59.5%1.07 TB/s · 60% of 1.8 TB/s HBM · also 4 TFLOPS (1% of compute)
8×16×4×64×2000×160.027 ms34.0%0.61 TB/s · 34% of 1.8 TB/s HBM · also 2 TFLOPS (0% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(47.5% · 72.8% · 38.5% · 59.5% · 34.0%) = 48.6%
Kernel source (redacted)
"""Paged-attention decode: custom fused CUDA flash-decoding kernel (SM120).
Single fused kernel launch, cp.async pipelined:
- Grid = one CTA per (batch, chunk, kv_head) work item (kv_head fastest so
concurrent CTAs read whole pages densely). Each warp streams its token
range through an SS-stage cp.async pipeline into shared memory.
- For G=8, pairs of warps cover the same tokens; each warp handles only
G/2 query heads (KV fetched once, L2 dedups), halving register pressure.
- Online flash softmax in exp2 domain, fp32 accumulators.
- Last CTA to finish a (batch, kv_head) pair (atomic ticket) reduces all
chunk partials and writes the final bf16 output; the ticket self-resets.
"""
import math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
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 <cuda_bf16.h>
#include <cstdlib>
using bf16 = __nv_bfloat16;
using bf162 = __nv_bfloat162;
#define SS 6
#define NEG_INF (-1e30f)
__device__ __forceinline__ float2 bf2_to_f2(uint32_t u) {
bf162 h = *reinterpret_cast<bf162*>(&u);
return __bfloat1622float2(h);
}
__device__ __forceinline__ void cp_async16(uint32_t smem, const void* gptr, int src_bytes, uint64_t pol) {
asm volatile("cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2, %3;\n"
::"r"(smem), "l"(gptr), "r"(src_bytes), "l"(pol));
}
__device__ __forceinline__ uint64_t mkpolicy(int mode) {
uint64_t pol = 0;
if (mode == 0) asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(pol));
else if (mode == 1) asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" : "=l"(pol));
else asm volatile("createpolicy.fractional.L2::evict_normal.b64 %0, 1.0;" : "=l"(pol));
return pol;
}
// ---------- mma helpers ----------
__device__ __forceinline__ void ldsm_x4(uint32_t* r, uint32_t a) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(a));
}
__device__ __forceinline__ void ldsm_x4_t(uint32_t* r, uint32_t a) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];"
: "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(a));
}
__device__ __forceinline__ void mma16816(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};"
: "+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]));
}
__device__ __forceinline__ void mma1688(float* c, const uint32_t* a, 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};"
: "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
: "r"(a[0]), "r"(a[1]), "r"(b));
}
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }
template <int D, int G, int CT, int NW>
__global__ void __launch_bounds__(NW * 32, 5)
paged_fused(const bf16* __restrict__ qp,
const bf16* __restrict__ kvp,
const int* __restrict__ btp,
const int* __restrict__ slp,
float* __restrict__ part,
int* __restrict__ counters,
bf16* __restrict__ out,
const int Hkv, const int max_pages, const int nch,
const float k_scale, const int pshift, const int pol_mode)
{
constexpr int QK_LANES = D / 8;
constexpr int TPG = 2 * QK_LANES;
constexpr int TPI = 32 / TPG;
constexpr int RANGES = (G == 8) ? 2 : 1; // token-range replicas
constexpr int GW = G / RANGES; // query heads per warp
constexpr int TSET = NW / RANGES; // distinct token ranges
constexpr int W_TOK = CT / TSET;
constexpr int W_IT = W_TOK / TPI;
constexpr int ROW = G * (D + 2);
constexpr bool QF32 = false; // keep q as bf16, unpack in the loop (saves 16 regs)
const int item = blockIdx.x;
const int b = item / (Hkv * nch);
const int r = item % (Hkv * nch);
const int chunk = r / Hkv;
const int hkv = r % Hkv;
const int L = slp[b];
const int t0 = chunk * CT;
if (L <= 0) {
if (chunk == 0 && hkv == 0) {
for (int i = threadIdx.x; i < G * Hkv * D; i += blockDim.x) {
const int h = i / D, d = i % D;
out[((size_t)b * (Hkv * G) + h) * D + d] = __float2bfloat16(0.f);
}
}
return;
}
if (t0 >= L) return;
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
const int tg = lane / TPG;
const int lpos = lane % TPG;
const bool isK = lpos < QK_LANES;
const int lv = isK ? lpos : (lpos - QK_LANES);
const int range_id = warp % TSET;
const int role = warp / TSET;
const int page_mask = (1 << pshift) - 1;
const size_t row_bytes = ((size_t)1 << pshift) * Hkv * (2 * D) * 2;
const size_t tok_stride = (size_t)Hkv * (2 * D) * 2;
const size_t lane_off = (size_t)(isK ? (lv * 8) : (D + lv * 8)) * 2;
__shared__ __align__(16) char kv_stages[NW][SS][512];
// merge buffers alias the (fully drained) stage memory: all real cp.async
// ops have completed by the time the token loop exits
float* acc_flat = reinterpret_cast<float*>(kv_stages);
#define ACC_SH(vg_, g_, d_) acc_flat[((size_t)(vg_) * GW + (g_)) * D + (d_)]
__shared__ float msh[NW * TPI][GW];
__shared__ float lsh[NW * TPI][GW];
__shared__ int flag_sh;
const char* kvb = reinterpret_cast<const char*>(kvp)
+ (size_t)hkv * (2 * D) * 2 + lane_off;
const int* bt_row = btp + (size_t)b * max_pages;
const int tw = t0 + range_id * W_TOK;
// page-id cache (warp token range spans a bounded number of pages)
const int pg0 = tw >> pshift;
const int npg = min(((tw + W_TOK - 1) >> pshift) - pg0, 4);
int pid[5];
#pragma unroll
for (int j = 0; j < 5; ++j) {
const int idx = min(pg0 + min(j, npg), max_pages - 1);
pid[j] = __ldg(bt_row + idx);
}
const bool pid_ok = (((tw + W_TOK - 1) >> pshift) - pg0) <= 4;
const uint32_t stage_base = (uint32_t)__cvta_generic_to_shared(&kv_stages[warp][0][0]) + lane * 16;
uint64_t l2pol = mkpolicy(pol_mode);
#define PGID(t_) \
({ \
int pgv_; \
if (pid_ok) { \
const int pi_ = ((t_) >> pshift) - pg0; \
pgv_ = pi_ == 0 ? pid[0] : (pi_ == 1 ? pid[1] : (pi_ == 2 ? pid[2] : (pi_ == 3 ? pid[3] : pid[4]))); \
} else { \
pgv_ = __ldg(bt_row + min((t_) >> pshift, max_pages - 1)); \
} \
pgv_; \
})
#define ISSUE(s_) \
{ \
const int t_ = tw + (s_) * TPI + tg; \
const int pg_ = PGID(t_); \
const char* src_ = kvb + (size_t)pg_ * row_bytes \
+ (size_t)(t_ & page_mask) * tok_stride; \
cp_async16(stage_base + ((s_) % SS) * 512, src_, (t_ < L) ? 16 : 0, l2pol); \
}
#pragma unroll
for (int s = 0; s < SS - 1; ++s) { if (s < W_IT) ISSUE(s); cp_commit(); }
// query fragments (K lanes only)
float qf[QF32 ? GW : 1][8];
uint4 qreg[QF32 ? 1 : GW];
if (isK) {
const bf16* qb = qp + ((size_t)b * (Hkv * G) + (size_t)(hkv * G + role * GW)) * D + lv * 8;
if (QF32) {
#pragma unroll
for (int g = 0; g < GW; ++g) {
const uint4 qr = __ldg(reinterpret_cast<const uint4*>(qb + (size_t)g * D));
const uint32_t* qx = reinterpret_cast<const uint32_t*>(&qr);
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float2 f = bf2_to_f2(qx[i]);
qf[g][2 * i] = f.x; qf[g][2 * i + 1] = f.y;
}
}
} else {
#pragma unroll
for (int g = 0; g < GW; ++g)
qreg[g] = __ldg(reinterpret_cast<const uint4*>(qb + (size_t)g * D));
}
}
float m[GW], l[GW], acc[GW][8];
#pragma unroll
for (int g = 0; g < GW; ++g) {
m[g] = NEG_INF; l[g] = 0.f;
#pragma unroll
for (int i = 0; i < 8; ++i) acc[g][i] = 0.f;
}
for (int it = 0; it < W_IT; ++it) {
if (it + SS - 1 < W_IT) ISSUE(it + SS - 1);
cp_commit();
cp_wait<SS - 1>();
__syncwarp();
const uint4 kvv = *reinterpret_cast<const uint4*>(&kv_stages[warp][it % SS][lane * 16]);
const uint32_t* kx = reinterpret_cast<const uint32_t*>(&kvv);
float kf8[8];
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float2 f = bf2_to_f2(kx[i]);
kf8[2 * i] = f.x; kf8[2 * i + 1] = f.y;
}
const int t = tw + it * TPI + tg;
const bool valid = t < L;
#pragma unroll
for (int g = 0; g < GW; ++g) {
float s = 0.f;
if (isK) {
if (QF32) {
#pragma unroll
for (int i = 0; i < 8; ++i) s = fmaf(kf8[i], qf[g][i], s);
} else {
const uint32_t* qx = reinterpret_cast<const uint32_t*>(&qreg[g]);
#pragma unroll
for (int i = 0; i < 4; ++i) {
const float2 f = bf2_to_f2(qx[i]);
s = fmaf(kf8[2 * i], f.x, s);
s = fmaf(kf8[2 * i + 1], f.y, s);
}
}
}
#pragma unroll
for (int o = QK_LANES / 2; o > 0; o >>= 1)
s += __shfl_xor_sync(0xffffffffu, s, o, TPG);
s = __shfl_xor_sync(0xffffffffu, s, QK_LANES, TPG);
s = valid ? (s * k_scale) : NEG_INF;
const float mn = fmaxf(m[g], s);
const bool dead = mn == NEG_INF;
const float p = dead ? 0.f : exp2f(s - mn);
const float c = dead ? 1.f : exp2f(m[g] - mn);
l[g] = l[g] * c + p;
m[g] = mn;
if (!isK) {
if (c != 1.f) {
#pragma unroll
for (int i = 0; i < 8; ++i) acc[g][i] *= c;
}
#pragma unroll
for (int i = 0; i < 8; ++i) acc[g][i] = fmaf(p, kf8[i], acc[g][i]);
}
}
}
// ---- merge token-groups in shared (each V-half tracks its own m/l/acc;
// K lanes carry clobbered scores and are never read for state) ----
__syncthreads(); // stages are dead past this point; safe to alias
const int vg = warp * TPI + tg;
if (lane % TPG == QK_LANES) {
#pragma unroll
for (int g = 0; g < GW; ++g) { msh[vg][g] = m[g]; lsh[vg][g] = l[g]; }
}
if (!isK) {
#pragma unroll
for (int g = 0; g < GW; ++g) {
#pragma unroll
for (int i = 0; i < 8; ++i) ACC_SH(vg, g, lv * 8 + i) = acc[g][i];
}
}
__syncthreads();
float* row = part + (((size_t)b * Hkv + hkv) * nch + chunk) * (size_t)ROW;
for (int dd = threadIdx.x; dd < D; dd += NW * 32) {
#pragma unroll
for (int g = 0; g < G; ++g) {
const int rl = g / GW; // which token-replica role
const int gl = g % GW;
float M = NEG_INF;
#pragma unroll
for (int rr = 0; rr < TSET * TPI; ++rr)
M = fmaxf(M, msh[(rl * TSET + rr % TSET) * TPI + rr / TSET][gl]);
float lt = 0.f, o = 0.f;
#pragma unroll
for (int rr = 0; rr < TSET * TPI; ++rr) {
const int w = (rl * TSET + rr % TSET) * TPI + rr / TSET;
const float e = exp2f(msh[w][gl] - M);
lt += lsh[w][gl] * e;
o += ACC_SH(w, gl, dd) * e;
}
float* rp = row + g * (D + 2);
rp[dd] = o;
if (dd == 0) { rp[D] = M; rp[D + 1] = lt; }
}
}
__threadfence();
__syncthreads();
if (threadIdx.x == 0) {
const int nch_b = (L + CT - 1) / CT;
const int old = atomicAdd(counters + (size_t)b * Hkv + hkv, 1);
flag_sh = (old == nch_b - 1) ? 1 : 0;
}
__syncthreads();
if (flag_sh) {
__threadfence();
const size_t ctr = (size_t)b * Hkv + hkv;
const int nch_b = (L + CT - 1) / CT;
const float* seg = part + ctr * nch * (size_t)ROW;
constexpr int ND = (D + NW * 32 - 1) / (NW * 32);
#pragma unroll 1
for (int g = 0; g < G; ++g) {
float M[ND], lt[ND], o[ND];
#pragma unroll
for (int j = 0; j < ND; ++j) { M[j] = NEG_INF; lt[j] = 0.f; o[j] = 0.f; }
const float* r0 = seg + (size_t)g * (D + 2);
for (int i = 0; i < nch_b; ++i) {
const float* r = r0 + (size_t)i * ROW;
const float mi = r[D];
const float li = r[D + 1];
#pragma unroll
for (int j = 0; j < ND; ++j) {
const int dd = threadIdx.x + j * (NW * 32);
const float ai = (dd < D) ? r[dd] : 0.f;
const float Mn = fmaxf(M[j], mi);
const float c = exp2f(M[j] - Mn);
const float e = exp2f(mi - Mn);
o[j] = o[j] * c + ai * e;
lt[j] = lt[j] * c + li * e;
M[j] = Mn;
}
}
#pragma unroll
for (int j = 0; j < ND; ++j) {
const int dd = threadIdx.x + j * (NW * 32);
if (dd < D) {
const float v = (lt[j] > 0.f) ? (o[j] / lt[j]) : 0.f;
out[((size_t)b * (Hkv * G) + (size_t)hkv * G + g) * D + dd] = __float2bfloat16(v);
}
}
}
if (threadIdx.x == 0) counters[ctr] = 0;
}
}
// Tensor-core flash-decode item kernel: 8-token batches, cp.async pipelined,
// QK via m16n8k16 (x4 ldmatrix covers 2 k-chunks), softmax in C-fragments,
// PV via m16n8k8 (P fragments pass straight through).
template <int D, int G, int CT, int NW>
__global__ void __launch_bounds__(NW * 32, 3)
paged_mma(const bf16* __restrict__ qp,
const bf16* __restrict__ kvp,
const int* __restrict__ btp,
const int* __restrict__ slp,
float* __restrict__ part,
int* __restrict__ counters,
bf16* __restrict__ out,
const int Hkv, const int max_pages, const int nch,
const float k_scale, const int pshift, const int pol_mode)
{
constexpr int TKS = 8;
constexpr int ROWB = 4 * D + 32; // padded smem token-row bytes
constexpr int STB = TKS * ROWB;
constexpr int SSM = 2;
constexpr int DC = D / 16;
constexpr int ND = D / 8;
constexpr int W_TOK = CT / NW;
constexpr int W_IT = W_TOK / TKS;
constexpr int ROW = G * (D + 2);
constexpr int TPC = (D == 128) ? 1 : 2; // tokens per cp.async instr
constexpr int VOFF = 2 * D; // V byte offset in token row
static_assert(W_TOK % TKS == 0, "W_TOK must be a multiple of 8");
const int item = blockIdx.x;
const int b = item / (Hkv * nch);
const int r = item % (Hkv * nch);
const int chunk = r / Hkv;
const int hkv = r % Hkv;
const int L = slp[b];
const int t0 = chunk * CT;
if (L <= 0) {
if (chunk == 0 && hkv == 0) {
for (int i = threadIdx.x; i < G * Hkv * D; i += blockDim.x) {
const int h = i / D, d = i % D;
out[((size_t)b * (Hkv * G) + h) * D + d] = __float2bfloat16(0.f);
}
}
return;
}
if (t0 >= L) return;
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
const int page_mask = (1 << pshift) - 1;
const size_t row_bytes = ((size_t)1 << pshift) * Hkv * (2 * D) * 2;
const size_t tok_stride = (size_t)Hkv * (2 * D) * 2;
__shared__ __align__(16) char kv_stages[NW][SSM][STB];
float* acc_flat = reinterpret_cast<float*>(kv_stages);
#define ACC_SH(vg_, g_, d_) acc_flat[((size_t)(vg_) * G + (g_)) * D + (d_)]
__shared__ float msh[NW][G];
__shared__ float lsh[NW][G];
__shared__ int flag_sh;
const uint64_t l2pol = mkpolicy(pol_mode);
const char* kvb = reinterpret_cast<const char*>(kvp) + (size_t)hkv * (2 * D) * 2;
const int* bt_row = btp + (size_t)b * max_pages;
const int tw = t0 + warp * W_TOK;
const int pg0 = tw >> pshift;
const int npg = min(((tw + W_TOK - 1) >> pshift) - pg0, 4);
int pid[5];
#pragma unroll
for (int j = 0; j < 5; ++j) {
const int idx = min(pg0 + min(j, npg), max_pages - 1);
pid[j] = __ldg(bt_row + idx);
}
const bool pid_ok = (((tw + W_TOK - 1) >> pshift) - pg0) <= 4;
#define PGID(t_) \
({ \
int pgv_; \
if (pid_ok) { \
const int pi_ = ((t_) >> pshift) - pg0; \
pgv_ = pi_ == 0 ? pid[0] : (pi_ == 1 ? pid[1] : (pi_ == 2 ? pid[2] : (pi_ == 3 ? pid[3] : pid[4]))); \
} else { \
pgv_ = __ldg(bt_row + min((t_) >> pshift, max_pages - 1)); \
} \
pgv_; \
})
// one cp.async instr covers TPC tokens (32 lanes x 16B = 512B); lane's token
// and byte offset within the token row depend on TPC
#define ISSUE(s_) \
{ \
const int idt0_ = tw + (s_) * TKS; \
char* slot_ = &kv_stages[warp][(s_) % SSM][0]; \
_Pragma("unroll") \
for (int k_ = 0; k_ < TKS / TPC; ++k_) { \
const int tk_ = idt0_ + k_ * TPC + ((TPC == 2) ? (lane >> 4) : 0); \
const int off_ = ((TPC == 2) ? (lane & 15) : lane) * 16; \
const int pg_ = PGID(tk_); \
const char* src_ = kvb + (size_t)pg_ * row_bytes \
+ (size_t)(tk_ & page_mask) * tok_stride + off_; \
char* dst_ = slot_ + (TPC == 2 ? (k_ * 2 + (lane >> 4)) * ROWB : k_ * ROWB) + off_; \
cp_async16((uint32_t)__cvta_generic_to_shared(dst_), src_, (tk_ < L) ? 16 : 0, l2pol); \
} \
}
#pragma unroll
for (int s = 0; s < SSM - 1; ++s) { if (s < W_IT) ISSUE(s); cp_commit(); }
// ---- A fragments (Q) loaded straight to registers ----
// a0 = (row lane/4, cols 2t+2dc*8/..); A tile mapping:
// r0: rows 0-7 ?? lane holds (g, 2t,2t+1),(g+8,...),(g,2t+8..),(g+8,2t+8..)=
uint32_t afrag[DC][4];
{
const bf16* qh = qp + ((size_t)b * (Hkv * G) + (size_t)hkv * G) * D;
const int rq = lane >> 2; // row (head) this lane's first half
const int cq = (lane & 3) * 2; // col pair base within 16-dim chunk
const int r0v = rq, r1v = rq + 8;
const bf16* q0 = qh + (size_t)r0v * D + cq;
const bf16* q1 = qh + (size_t)r1v * D + cq;
#pragma unroll
for (int dc = 0; dc < DC; ++dc) {
uint32_t lo0 = 0, hi0 = 0, lo1 = 0, hi1 = 0;
if (r0v < G) {
lo0 = __ldg(reinterpret_cast<const uint32_t*>(q0 + dc * 16));
hi0 = __ldg(reinterpret_cast<const uint32_t*>(q0 + dc * 16 + 8));
}
if (r1v < G) {
lo1 = __ldg(reinterpret_cast<const uint32_t*>(q1 + dc * 16));
hi1 = __ldg(reinterpret_cast<const uint32_t*>(q1 + dc * 16 + 8));
}
afrag[dc][0] = lo0; afrag[dc][1] = lo1;
afrag[dc][2] = hi0; afrag[dc][3] = hi1;
}
}
float acc[ND][4];
float m[2], l[2];
#pragma unroll
for (int j = 0; j < ND; ++j) {
#pragma unroll
for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;
}
m[0] = m[1] = -1e20f; l[0] = l[1] = 0.f;
for (int it = 0; it < W_IT; ++it) {
if (it + SSM - 1 < W_IT) ISSUE(it + SSM - 1);
cp_commit();
cp_wait<SSM - 1>();
__syncwarp();
const char* slot = &kv_stages[warp][it % SSM][0];
const uint32_t slot_b = (uint32_t)__cvta_generic_to_shared(slot);
const int bt0 = tw + it * TKS;
// ---- QK ----
float sfrag[4];
#pragma unroll
for (int i = 0; i < 4; ++i) sfrag[i] = 0.f;
#pragma unroll
for (int dc2 = 0; dc2 < DC / 2; ++dc2) {
uint32_t bfr[4];
// x4 non-trans: lanes 0-7: tokens 0-7 @dc0; 8-15: tokens 0-7 @dc0+16B;
// 16-23: tokens 0-7 @dc1 (=+32B); 24-31: tokens 0-7 @dc1+16B
const uint32_t ad = slot_b + (lane & 7) * ROWB + dc2 * 64 + (lane >> 3) * 16;
ldsm_x4(bfr, ad);
mma16816(sfrag, afrag[dc2 * 2], bfr);
mma16816(sfrag, afrag[dc2 * 2 + 1], bfr + 2);
}
// ---- softmax in fragments ----
// scale + column masking
#pragma unroll
for (int h = 0; h < 2; ++h) {
#pragma unroll
for (int j = 0; j < 2; ++j) {
const int tok = bt0 + 2 * (lane & 3) + j;
sfrag[2 * h + j] = (tok < L) ? (sfrag[2 * h + j] * k_scale) : -1e20f;
}
}
float corr[2];
#pragma unroll
for (int h = 0; h < 2; ++h) {
float mx = fmaxf(sfrag[2 * h], sfrag[2 * h + 1]);
mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 1));
mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, 2));
const float mn = fmaxf(m[h], mx);
corr[h] = exp2f(m[h] - mn);
m[h] = mn;
}
float lsum[2] = {0.f, 0.f};
#pragma unroll
for (int h = 0; h < 2; ++h) {
#pragma unroll
for (int j = 0; j < 2; ++j) {
const int tok = bt0 + 2 * (lane & 3) + j;
const float p = (tok < L) ? exp2f(sfrag[2 * h + j] - m[h]) : 0.f;
lsum[h] += p;
sfrag[2 * h + j] = p;
}
lsum[h] += __shfl_xor_sync(0xffffffffu, lsum[h], 1);
lsum[h] += __shfl_xor_sync(0xffffffffu, lsum[h], 2);
l[h] = l[h] * corr[h] + lsum[h];
}
uint32_t pa[2];
{ const __nv_bfloat162 t_ = __floats2bfloat162_rn(sfrag[0], sfrag[1]); pa[0] = *reinterpret_cast<const uint32_t*>(&t_); }
{ const __nv_bfloat162 t_ = __floats2bfloat162_rn(sfrag[2], sfrag[3]); pa[1] = *reinterpret_cast<const uint32_t*>(&t_); }
#pragma unroll
for (int j = 0; j < ND; ++j) {
acc[j][0] *= corr[0]; acc[j][1] *= corr[0];
acc[j][2] *= corr[1]; acc[j][3] *= corr[1];
}
// ---- PV ----
#pragma unroll
for (int dq = 0; dq < ND / 4; ++dq) {
uint32_t vfr[4];
// x4.trans: lanes i*8..i*8+7: token rows 0-7 @ VOFF+dq*64+i*16
const uint32_t ad = slot_b + (lane & 7) * ROWB + VOFF + dq * 64 + (lane >> 3) * 16;
ldsm_x4_t(vfr, ad);
#pragma unroll
for (int j = 0; j < 4; ++j)
mma1688(acc[dq * 4 + j], pa, vfr[j]);
}
}
// ---- merge token-ranges (warps) in shared ----
__syncthreads();
const int row0 = lane >> 2;
const int tokq = 2 * (lane & 3);
if (row0 < G) {
#pragma unroll
for (int j = 0; j < ND; ++j) {
ACC_SH(warp, row0, j * 8 + tokq) = acc[j][0];
ACC_SH(warp, row0, j * 8 + tokq + 1) = acc[j][1];
}
}
if ((row0 + 8) < G) {
#pragma unroll
for (int j = 0; j < ND; ++j) {
ACC_SH(warp, row0 + 8, j * 8 + tokq) = acc[j][2];
ACC_SH(warp, row0 + 8, j * 8 + tokq + 1) = acc[j][3];
}
}
if ((lane & 3) == 0) {
if (row0 < G) { msh[warp][row0] = m[0]; lsh[warp][row0] = l[0]; }
if ((row0 + 8) < G) { msh[warp][row0 + 8] = m[1]; lsh[warp][row0 + 8] = l[1]; }
}
__syncthreads();
float* row = part + (((size_t)b * Hkv + hkv) * nch + chunk) * (size_t)ROW;
for (int dd = threadIdx.x; dd < D; dd += NW * 32) {
#pragma unroll
for (int g = 0; g < G; ++g) {
float M = -1e20f;
#pragma unroll
for (int w = 0; w < NW; ++w) M = fmaxf(M, msh[w][g]);
float lt = 0.f, o = 0.f;
#pragma unroll
for (int w = 0; w < NW; ++w) {
const float e = exp2f(msh[w][g] - M);
lt += lsh[w][g] * e;
o += ACC_SH(w, g, dd) * e;
}
float* rp = row + g * (D + 2);
rp[dd] = o;
if (dd == 0) { rp[D] = M; rp[D + 1] = lt; }
}
}
__threadfence();
__syncthreads();
if (threadIdx.x == 0) {
const int nch_b = (L + CT - 1) / CT;
const int old = atomicAdd(counters + (size_t)b * Hkv + hkv, 1);
flag_sh = (old == nch_b - 1) ? 1 : 0;
}
__syncthreads();
if (flag_sh) {
__threadfence();
const size_t ctr = (size_t)b * Hkv + hkv;
const int nch_b = (L + CT - 1) / CT;
const float* seg = part + ctr * nch * (size_t)ROW;
constexpr int NDD = (D + NW * 32 - 1) / (NW * 32);
#pragma unroll 1
for (int g = 0; g < G; ++g) {
float M[NDD], lt[NDD], o[NDD];
#pragma unroll
for (int j = 0; j < NDD; ++j) { M[j] = NEG_INF; lt[j] = 0.f; o[j] = 0.f; }
const float* r0 = seg + (size_t)g * (D + 2);
for (int i = 0; i < nch_b; ++i) {
const float* r = r0 + (size_t)i * ROW;
const float mi = r[D];
const float li = r[D + 1];
#pragma unroll
for (int j = 0; j < NDD; ++j) {
const int dd = threadIdx.x + j * (NW * 32);
const float ai = (dd < D) ? r[dd] : 0.f;
const float Mn = fmaxf(M[j], mi);
const float c = exp2f(M[j] - Mn);
const float e = exp2f(mi - Mn);
o[j] = o[j] * c + ai * e;
lt[j] = lt[j] * c + li * e;
M[j] = Mn;
}
}
#pragma unroll
for (int j = 0; j < NDD; ++j) {
const int dd = threadIdx.x + j * (NW * 32);
if (dd < D) {
const float v = (lt[j] > 0.f) ? (o[j] / lt[j]) : 0.f;
out[((size_t)b * (Hkv * G) + (size_t)hkv * G + g) * D + dd] = __float2bfloat16(v);
}
}
}
if (threadIdx.x == 0) counters[ctr] = 0;
}
}
torch::Tensor paged_fwd(torch::Tensor q, torch::Tensor kv, torch::Tensor bt,
torch::Tensor sl, torch::Tensor partials, torch::Tensor counters,
double scale, int64_t pshift, int64_t nch, int64_t chunk,
int64_t nw, int64_t pol)
{
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 G = H / Hkv;
const int max_pages = bt.size(1);
const int n_items = B * Hkv * (int)nch;
auto out = torch::empty_like(q);
auto stream = at::cuda::getCurrentCUDAStream();
const float k_scale = (float)(scale * 1.44269504088896340736);
const bf16* qp = reinterpret_cast<const bf16*>(q.data_ptr());
const bf16* kp = reinterpret_cast<const bf16*>(kv.data_ptr());
const int* bp = bt.data_ptr<int>();
const int* sp = sl.data_ptr<int>();
float* pp = partials.data_ptr<float>();
int* cp = counters.data_ptr<int>();
bf16* op = reinterpret_cast<bf16*>(out.data_ptr());
#define LAUNCH(D_, G_, CT_) \
paged_fused<D_, G_, CT_, 4><<<n_items, 4 * 32, 0, stream>>>( \
qp, kp, bp, sp, pp, cp, op, Hkv, max_pages, (int)nch, k_scale, (int)pshift, (int)pol)
#define LAUNCH_MMA(D_, G_, CT_) \
paged_mma<D_, G_, CT_, 2><<<n_items, 2 * 32, 0, stream>>>( \
qp, kp, bp, sp, pp, cp, op, Hkv, max_pages, (int)nch, k_scale, (int)pshift, (int)pol)
#define LAUNCH_MMA1(D_, G_, CT_) \
paged_mma<D_, G_, CT_, 1><<<n_items, 1 * 32, 0, stream>>>( \
qp, kp, bp, sp, pp, cp, op, Hkv, max_pages, (int)nch, k_scale, (int)pshift, (int)pol)
static int impl = -1;
if (impl < 0) { const char* e = getenv("PA_IMPL"); impl = e ? atoi(e) : 1; }
if (impl == 1 && (chunk == 64 || chunk == 128) &&
((D == 128 && (G == 4 || G == 8)) || (D == 64 && G == 4))) {
if (D == 128 && G == 4 && chunk == 64) { LAUNCH_MMA(128, 4, 64); return out; }
if (D == 128 && G == 4 && chunk == 128) { LAUNCH_MMA(128, 4, 128); return out; }
if (D == 64 && G == 4 && chunk == 64) { LAUNCH_MMA(64, 4, 64); return out; }
if (D == 64 && G == 4 && chunk == 128) { LAUNCH_MMA(64, 4, 128); return out; }
}
if (D == 128 && G == 4) {
if (chunk == 64) LAUNCH(128, 4, 64);
else if (chunk == 128) LAUNCH(128, 4, 128);
else if (chunk == 192) LAUNCH(128, 4, 192);
else LAUNCH(128, 4, 256);
} else if (D == 128 && G == 8) {
static int nw1 = -1;
if (nw1 < 0) { const char* e = getenv("PA_MMA_NW"); nw1 = e ? atoi(e) : 2; }
if (chunk == 64 && nw1 == 1) LAUNCH_MMA1(128, 8, 64);
else if (chunk == 128 && nw1 == 1) LAUNCH_MMA1(128, 8, 128);
else if (chunk == 64) LAUNCH_MMA(128, 8, 64);
else if (chunk == 128) LAUNCH_MMA(128, 8, 128);
else TORCH_CHECK(false, "unsupported chunk ", chunk, " for G=8");
} else if (D == 64 && G == 4) {
if (chunk == 64) LAUNCH(64, 4, 64);
else if (chunk == 128) LAUNCH(64, 4, 128);
else if (chunk == 192) LAUNCH(64, 4, 192);
else LAUNCH(64, 4, 256);
} else {
TORCH_CHECK(false, "unsupported (D,G): ", D, " ", G);
}
return out;
}
"""
_CPP_SRC = """
torch::Tensor paged_fwd(torch::Tensor q, torch::Tensor kv, torch::Tensor bt,
torch::Tensor sl, torch::Tensor partials, torch::Tensor counters,
double scale, int64_t pshift, int64_t nch, int64_t chunk,
int64_t nw, int64_t pol);
"""
_ext = None
import os as _os
_ENV_CHUNK = int(_os.environ.get("PA_CHUNK", "0"))
_ENV_POL_0 = int(_os.environ.get("PA_POL0", _os.environ.get("PA_POL", "0")))
_ENV_POL_1 = int(_os.environ.get("PA_POL1", _os.environ.get("PA_POL", "0")))
_EXT_FN = None
def _get_ext():
global _ext, _EXT_FN
if _ext is None:
_ext = load_inline(
name="paged_attn_decode_v4",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=["paged_fwd"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
verbose=False,
)
_EXT_FN = _ext.paged_fwd
return _ext
class Model(nn.Module):
"""Single-query paged attention decode (custom fused CUDA kernel)."""
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, "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._ws = {}
_get_ext()
# precompute shape-derived launch parameters
self._pshift = page_size.bit_length() - 1
pairs = batch * num_kv_heads
lmax = ((seq_len + page_size - 1) // page_size) * page_size
items128 = pairs * ((lmax + 127) // 128)
self._chunk = 64 if 750 <= items128 < 3000 else 128
self._pages_nominal = (seq_len + page_size - 1) // page_size
self._nch_nominal = (lmax + self._chunk - 1) // self._chunk
self._nch = self._nch_nominal
self._ws_key = (batch, num_kv_heads, self._nch, self.group_size, head_dim)
self._wsv = None
def _fallback(self, query, kv_cache, block_table, seq_lens):
B, H, D = query.shape
Hkv = self.num_kv_heads
G = self.group_size
P = self.page_size
max_pages = block_table.size(1)
kv = kv_cache.index_select(0, block_table.reshape(-1).long())
kv = kv.reshape(B, max_pages * P, Hkv, 2 * D)
k = kv[..., :D].repeat_interleave(G, dim=2).float()
v = kv[..., D:].repeat_interleave(G, dim=2).float()
qf = query.float()
scores = torch.einsum("bhd,blhd->bhl", qf, k) * self.scale
tok = torch.arange(max_pages * P, device=query.device)[None, None, :]
scores = scores.masked_fill(tok >= seq_lens[:, None, None].long(), float("-inf"))
probs = torch.softmax(scores, dim=-1)
o = torch.einsum("bhl,blhd->bhd", probs, v)
return o.to(query.dtype)
def forward(self, query, kv_cache, block_table, seq_lens):
import os
B, H, D = query.shape
Hkv = kv_cache.size(2)
G = H // Hkv
P = self.page_size
fast = (
query.dtype == torch.bfloat16
and block_table.dtype == torch.int32
and seq_lens.dtype == torch.int32
and query.is_contiguous()
and kv_cache.is_contiguous()
and block_table.is_contiguous()
and seq_lens.is_contiguous()
)
if not fast:
return self._fallback(query, kv_cache, block_table, seq_lens)
ws = self._wsv
if block_table.size(1) == self._pages_nominal and ws is not None:
return _EXT_FN(query, kv_cache, block_table, seq_lens, ws[0], ws[1],
self.scale, self._pshift, self._nch_nominal, self._chunk, 4,
_ENV_POL_1 if G == 8 else _ENV_POL_0)
# general path: derive launch parameters from this call's tensors
dev = query.device
max_pages = block_table.size(1)
pairs = B * Hkv
items128 = pairs * ((max_pages * P + 127) // 128)
chunk = 64 if 750 <= items128 < 3000 else 128
nch = (max_pages * P + chunk - 1) // chunk
key = (B, Hkv, nch, G, D)
wsc = self._ws.get(key)
if wsc is None or wsc[0].device != dev:
partials = torch.empty(B * Hkv * nch * G * (D + 2), dtype=torch.float32, device=dev)
counters = torch.zeros(B * Hkv, dtype=torch.int32, device=dev)
wsc = (partials, counters)
self._ws[key] = wsc
if max_pages == self._pages_nominal:
self._wsv = wsc
return _EXT_FN(query, kv_cache, block_table, seq_lens, wsc[0], wsc[1],
self.scale, self._pshift, nch, chunk, 4,
_ENV_POL_1 if G == 8 else _ENV_POL_0)
# Bypass nn.Module.__call__ hook machinery (saves ~2-4us/launch on the hot path).
__call__ = forward
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]
20260716_112603_kinetic-claude_kinetic-0715_03_paged_attention