diff --git a/csrc/build.py b/csrc/build.py index 6797898..8e18529 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -45,3 +45,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs): register("attn_decode") register("attn_prefill") +register("attn_paged_decode") diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index ea02e69..f99c38d 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attn_common.h @@ -14,11 +14,39 @@ struct AttentionParams { int causal_offset; int num_splits; float scale; - + const T* __restrict__ q; const T* __restrict__ k; const T* __restrict__ v; const bool* __restrict__ mask; + + T* __restrict__ o; + AT* __restrict__ o_part; + AT* __restrict__ ml_part; +}; + +template +struct PagedAttentionParams { + int batch; + int q_head; + int kv_head; + int q_len; + int kv_len; + int head_dim; + int use_mask; + int is_causal; + int causal_offset; + float scale; + + int num_splits; + int page_size; + int max_pages; + + const T* __restrict__ q; + const T* __restrict__ k_cache; + const T* __restrict__ v_cache; + const bool* __restrict__ mask; + const int64_t* __restrict__ page_table; T* __restrict__ o; AT* __restrict__ o_part; diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu new file mode 100644 index 0000000..7f7c8c4 --- /dev/null +++ b/csrc/kernels/attn_paged_decode.cu @@ -0,0 +1,152 @@ +#include "attn_paged_decode_split_kv.cuh" +#ifndef ASTRAI_NO_MMA +#include "attn_paged_decode_split_kv_mma.cuh" +#endif + +#include +#include + +static int paged_decode_num_splits(int base_blocks, int tiles_total) { + int sm_count = 0; + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); + int n = (2 * sm_count + base_blocks - 1) / base_blocks; + return std::max(1, std::min(n, std::min(tiles_total, 32))); +} + +static void launch_paged_scalar_decode(PagedAttentionParams& p) { + int group_size = p.q_head / p.kv_head; + int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; + p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, chunks_total); + + auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); + auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt); + p.o_part = o_part.data_ptr(); + p.ml_part = ml_part.data_ptr(); + + size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16); + paged_attn_decode_split_kv_kernel<<< + dim3(p.batch * p.kv_head, 1, p.num_splits), + dim3(32, group_size), + smem>>>(p); + paged_attn_decode_combine_kernel<<>>(p); +} + +#ifndef ASTRAI_NO_MMA +template +static void launch_paged_mma_decode(PagedAttentionParams& p) { + int tiles_total = (p.kv_len + BC - 1) / BC; + p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total); + + auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); + auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt); + p.o_part = o_part.data_ptr(); + p.ml_part = ml_part.data_ptr(); + + paged_attn_decode_split_kv_mma_kernel + <<>>(p); + paged_attn_decode_combine_kernel<<>>(p); +} +#endif + +template +static void dispatch_paged_decode(PagedAttentionParams& p) { +#ifndef ASTRAI_NO_MMA + int G = p.q_head / p.kv_head; + if (!p.use_mask && G >= 1 && G <= 16 && p.page_size >= 32) { + launch_paged_mma_decode(p); + return; + } +#endif + launch_paged_scalar_decode(p); +} + +torch::Tensor attn_paged_decode( + torch::Tensor q, + torch::Tensor page_table, + torch::Tensor k_cache, + torch::Tensor v_cache, + int64_t page_size, + int64_t kv_len, + c10::optional mask, + bool is_causal = false, + int64_t causal_offset = 0, + c10::optional scale = c10::nullopt +) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + + int batch = q.size(0); + int q_head = q.size(1); + int head_dim = q.size(3); + int kv_head = k_cache.size(3); + int max_pages = page_table.size(1); + + TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda()); + TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16"); + TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16"); + TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16"); + TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64"); + TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)"); + TORCH_CHECK(head_dim % 32 == 0, "head_dim must be multiple of 32"); + TORCH_CHECK(k_cache.size(1) == page_size, + "k_cache dim 1 must equal page_size, got ", + k_cache.size(1), " vs ", page_size); + TORCH_CHECK(k_cache.size(0) >= 0, "k_cache must have at least 0 pages"); + + float scale_val = scale.has_value() + ? static_cast(scale.value()) + : 1.0f / std::sqrt(static_cast(head_dim)); + + auto O = torch::empty_like(q); + + PagedAttentionParams p; + p.batch = batch; + p.q_head = q_head; + p.kv_head = kv_head; + p.q_len = 1; + p.kv_len = static_cast(kv_len); + p.head_dim = head_dim; + p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0; + p.is_causal = is_causal ? 1 : 0; + p.causal_offset = static_cast(causal_offset); + p.num_splits = 1; + p.scale = scale_val; + p.page_size = static_cast(page_size); + p.max_pages = max_pages; + p.page_table = page_table.data_ptr(); + p.k_cache = reinterpret_cast(k_cache.data_ptr()); + p.v_cache = reinterpret_cast(v_cache.data_ptr()); + p.q = reinterpret_cast(q.data_ptr()); + p.mask = p.use_mask ? mask.value().data_ptr() : nullptr; + p.o = reinterpret_cast(O.data_ptr()); + p.o_part = nullptr; + p.ml_part = nullptr; + + switch (p.head_dim) { + case 32: dispatch_paged_decode<32>(p); break; + case 64: dispatch_paged_decode<64>(p); break; + case 128: dispatch_paged_decode<128>(p); break; + case 256: dispatch_paged_decode<256>(p); break; + default: + TORCH_CHECK(false, "paged_decode: unsupported head_dim ", p.head_dim, + " (supported: 32, 64, 128, 256)"); + } + + return O; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("attn_paged_decode", &attn_paged_decode, + py::arg("q"), + py::arg("page_table"), + py::arg("k_cache"), + py::arg("v_cache"), + py::arg("page_size"), + py::arg("kv_len"), + py::arg("mask") = py::none(), + py::arg("is_causal") = false, + py::arg("causal_offset") = 0, + py::arg("scale") = py::none(), + "Paged GQA decode — split-KV with direct page-table access."); +} diff --git a/csrc/kernels/attn_paged_decode_split_kv.cuh b/csrc/kernels/attn_paged_decode_split_kv.cuh new file mode 100644 index 0000000..a54f467 --- /dev/null +++ b/csrc/kernels/attn_paged_decode_split_kv.cuh @@ -0,0 +1,140 @@ +#pragma once +#include +#include +#include "attn_common.h" + +using bf16 = __nv_bfloat16; +constexpr int PDC_CHUNK = 64; + +__device__ inline float paged_warp_reduce_sum(float val) { + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +} + +// Split-KV scalar decode: one warp per query head, grid.z partitions KV. +__global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) { + int batch = blockIdx.x / p.kv_head; + int kv_head = blockIdx.x % p.kv_head; + int split = blockIdx.z; + int group_size = blockDim.y; + int q_head = kv_head * group_size + threadIdx.y; + int lane = threadIdx.x; + int hd_per_thread = p.head_dim / 32; + + float q_reg[8]; + int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread; + #pragma unroll + for (int i = 0; i < hd_per_thread; i++) + q_reg[i] = __bfloat162float(p.q[q_off + i]); + + float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f}; + + extern __shared__ __align__(16) bf16 k_smem[]; + + int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; + int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits; + int ch_begin = split * chunks_per_split; + int ch_end = min(chunks_total, ch_begin + chunks_per_split); + + const int mask_base = batch * p.kv_len; + + for (int ci = ch_begin; ci < ch_end; ci++) { + int chunk_start = ci * PDC_CHUNK; + int this_chunk = min(PDC_CHUNK, p.kv_len - chunk_start); + + int total = this_chunk * p.head_dim; + for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) { + int s = i / p.head_dim; + int d_dim = i % p.head_dim; + int pos = chunk_start + s; + int logical_page = pos / p.page_size; + int page_offset = pos % p.page_size; + int phys_page = p.page_table[batch * p.max_pages + logical_page]; + if (phys_page >= 0) { + int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim + + (int64_t)page_offset * p.kv_head * p.head_dim + + (int64_t)kv_head * p.head_dim + + d_dim; + k_smem[i] = p.k_cache[off]; + } else { + k_smem[i] = __float2bfloat16(0.0f); + } + } + __syncthreads(); + + for (int s = 0; s < this_chunk; s++) { + float partial = 0.0f; + #pragma unroll + for (int i = 0; i < hd_per_thread; i++) + partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]); + partial = paged_warp_reduce_sum(partial) * p.scale; + + if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s]) + partial = -FLT_MAX; + if (p.is_causal && (chunk_start + s) > p.causal_offset) + partial = -FLT_MAX; + + float new_m = fmaxf(m, partial); + float alpha = expf(m - new_m); + float beta = expf(partial - new_m); + d = d * alpha + beta; + + int pos = chunk_start + s; + int logical_page = pos / p.page_size; + int page_offset = pos % p.page_size; + int phys_page = p.page_table[batch * p.max_pages + logical_page]; + if (phys_page >= 0) { + int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim + + (int64_t)page_offset * p.kv_head * p.head_dim + + (int64_t)kv_head * p.head_dim; + #pragma unroll + for (int i = 0; i < hd_per_thread; i++) + acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta; + } else { + #pragma unroll + for (int i = 0; i < hd_per_thread; i++) + acc_reg[i] = acc_reg[i] * alpha + 0.0f * beta; + } + m = new_m; + } + __syncthreads(); + } + + size_t bh = (size_t)batch * p.q_head + q_head; + size_t slot = bh * p.num_splits + split; + int d0 = lane * hd_per_thread; + #pragma unroll + for (int i = 0; i < hd_per_thread; i++) + p.o_part[slot * p.head_dim + (d0 + i)] = acc_reg[i]; + if (lane == 0) { + p.ml_part[slot * 2] = m; + p.ml_part[slot * 2 + 1] = d; + } +} + +__global__ void paged_attn_decode_combine_kernel(PagedAttentionParams p) { + int bh = blockIdx.x; + int d = threadIdx.x; + if (d >= p.head_dim) return; + + size_t split_base = (size_t)bh * p.num_splits; + const float* mlp = p.ml_part + split_base * 2; + const float* op = p.o_part + split_base * p.head_dim; + + float m = -FLT_MAX, l = 0.0f, acc = 0.0f; + for (int s = 0; s < p.num_splits; s++) { + float mi = mlp[s * 2]; + if (mi <= -FLT_MAX) continue; + float li = mlp[s * 2 + 1]; + float nm = fmaxf(m, mi); + float corr = __expf(m - nm); + float e = __expf(mi - nm); + acc = acc * corr + op[s * p.head_dim + d] * e; + l = l * corr + li * e; + m = nm; + } + + float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f; + p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv); +} diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh new file mode 100644 index 0000000..b4839a3 --- /dev/null +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -0,0 +1,174 @@ +#pragma once +#include +#include +#include "attn_common.h" +#include "attn_mma_utils.cuh" + +using bf16 = __nv_bfloat16; + +// Paged split-KV tensor-core decode via GQA head-packing. +// Identical algorithm to attn_decode_split_kv_mma_kernel but reads K/V +// directly from the page pool through a page table, eliminating the gather +// copy. Each tile (BC=32) fits within a single page (page_size >= 32), so +// the page-table lookup happens once per tile for cp.async. +template +__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams p) { + constexpr int BR = 16; + constexpr int KD = HEAD_DIM / 16; + constexpr int NC8 = BC / 8; + constexpr int KT2 = BC / 16; + constexpr int DN8 = HEAD_DIM / 8; + constexpr int LD = HEAD_DIM; + constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); + + const int lane = threadIdx.x; + const int gid = lane >> 2; + const int tid4 = lane & 3; + + const int kv_head_idx = blockIdx.x; + const int batch = blockIdx.y; + const int split = blockIdx.z; + const int G = p.q_head / p.kv_head; + const int q_head0 = kv_head_idx * G; + + __shared__ __align__(16) bf16 sK[BC * HEAD_DIM]; + __shared__ __align__(16) bf16 sV[BC * HEAD_DIM]; + __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM]; + + // ---- load Q into registers via ldmatrix ---- + for (int i = lane; i < BR * HEAD_DIM; i += 32) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + bf16 val = __float2bfloat16(0.0f); + if (r < G) { + int qh = q_head0 + r; + val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; + } + sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = val; + } + __syncwarp(); + + unsigned Qa[KD][4]; + int qrow_l = (lane & 7) + (lane & 8); + int qcol_l = (lane & 16) ? 8 : 0; + #pragma unroll + for (int kt = 0; kt < KD; kt++) + ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); + + float Oacc[DN8][4]; + #pragma unroll + for (int j = 0; j < DN8; j++) + Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; + float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; + + const int mask_base = batch * p.kv_len; + const int tiles_total = (p.kv_len + BC - 1) / BC; + const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; + const int ti_begin = split * tiles_per_split; + const int ti_end = min(tiles_total, ti_begin + tiles_per_split); + const int has_mask = p.use_mask && p.mask; + + // Paged strides (constant for the block) + const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM; + const int64_t pos_stride = (int64_t)p.kv_head * HEAD_DIM; + const int64_t head_off = (int64_t)kv_head_idx * HEAD_DIM; + + for (int ti = ti_begin; ti < ti_end; ti++) { + int kv0 = ti * BC; + + // phys_page is constant for the whole tile (BC <= page_size). + int logical_page = kv0 / p.page_size; + int phys_page = p.page_table[batch * p.max_pages + logical_page]; + bool page_valid = (phys_page >= 0); + + bool full_tile = page_valid && (kv0 + BC <= p.kv_len); + if (full_tile) { + constexpr int VEC = 8; + int total = BC * HEAD_DIM; + #pragma unroll + for (int i = lane * VEC; i < total; i += 32 * VEC) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int kc = kv0 + r; + int page_off = kc % p.page_size; + int64_t gmem_base = (int64_t)phys_page * page_stride + + (int64_t)page_off * pos_stride + + head_off; + cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], + &p.k_cache[gmem_base + d]); + cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], + &p.v_cache[gmem_base + d]); + } + cp_async_commit(); + cp_async_wait_all(); + } else { + for (int i = lane; i < BC * HEAD_DIM; i += 32) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int kc = kv0 + r; + bf16 z = __float2bfloat16(0.0f); + if (kc < p.kv_len && page_valid) { + int page_off = kc % p.page_size; + int64_t gmem_base = (int64_t)phys_page * page_stride + + (int64_t)page_off * pos_stride + + head_off; + sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.k_cache[gmem_base + d]; + sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = p.v_cache[gmem_base + d]; + } else { + sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = z; + sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = z; + } + } + } + __syncwarp(); + + float Sacc[NC8][4]; + mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); + + #pragma unroll + for (int n8 = 0; n8 < NC8; n8++) + Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale, + Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale; + + int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + mma_softmax_tile(kv0, maxc, maxc, + mask_base, p.mask, has_mask, + Sacc, Oacc, m0, m1, l0, l1, lane); + + mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); + __syncwarp(); + } + + // ---- write UN-normalised partials for this split ---- + auto split_slot = [&](int h) -> size_t { + size_t bh = (size_t)batch * p.q_head + h; + return bh * p.num_splits + split; + }; + #pragma unroll + for (int dn8 = 0; dn8 < DN8; dn8++) { + int d = dn8 * 8 + 2 * tid4; + int r0 = gid, r1 = gid + 8; + if (r0 < G) { + int h = q_head0 + r0; + float* op = p.o_part + split_slot(h) * HEAD_DIM; + op[d] = Oacc[dn8][0]; + op[d + 1] = Oacc[dn8][1]; + } + if (r1 < G) { + int h = q_head0 + r1; + float* op = p.o_part + split_slot(h) * HEAD_DIM; + op[d] = Oacc[dn8][2]; + op[d + 1] = Oacc[dn8][3]; + } + } + if (tid4 == 0) { + int r0 = gid, r1 = gid + 8; + if (r0 < G) { + int h = q_head0 + r0; + float* mp = p.ml_part + split_slot(h) * 2; + mp[0] = m0; mp[1] = l0; + } + if (r1 < G) { + int h = q_head0 + r1; + float* mp = p.ml_part + split_slot(h) * 2; + mp[0] = m1; mp[1] = l1; + } + } +} diff --git a/csrc/tests/attn_paged_vs_contiguous.cu b/csrc/tests/attn_paged_vs_contiguous.cu new file mode 100644 index 0000000..2b26304 --- /dev/null +++ b/csrc/tests/attn_paged_vs_contiguous.cu @@ -0,0 +1,279 @@ +// Compile: +// nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \ +// --extra-device-vectorization csrc/tests/attn_paged_vs_contiguous.cu \ +// -o /tmp/test_pv && /tmp/test_pv + +#include +#include +#include +#include +#include +#include "../kernels/attn_paged_decode_split_kv.cuh" +#ifndef ASTRAI_NO_MMA +#include "../kernels/attn_paged_decode_split_kv_mma.cuh" +#endif + +using bf16 = __nv_bfloat16; + +static int num_splits(int base_blocks, int tiles_total) { + int sm_count = 0; + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); + int n = (2 * sm_count + base_blocks - 1) / base_blocks; + return std::max(1, std::min(n, std::min(tiles_total, 32))); +} + +// Copy contiguous K/V from page pool (reference gather) +static void gather_kv_cpu( + const bf16* h_k_pool, const bf16* h_v_pool, + const int64_t* h_pt, int B, int Hkv, int kv_len, + int page_size, int head_dim, + bf16* h_k, bf16* h_v) +{ + int max_pages = (kv_len + page_size - 1) / page_size; + size_t page_stride = (size_t)page_size * Hkv * head_dim; + for (int b = 0; b < B; b++) { + for (int pos = 0; pos < kv_len; pos++) { + int log_pg = pos / page_size; + int pg_off = pos % page_size; + int phys = (int)h_pt[b * max_pages + log_pg]; + for (int h = 0; h < Hkv; h++) { + size_t src_base = (size_t)phys * page_stride + + (size_t)pg_off * Hkv * head_dim + + h * head_dim; + size_t dst_base = ((size_t)b * kv_len + pos) * Hkv * head_dim + h * head_dim; + memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16)); + memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16)); + } + } + } +} + +template +static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) { + printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM); + fflush(stdout); + + int G = Hq / Hkv; + int max_pages = (kv_len + page_size - 1) / page_size; + int n_phys_pages = B * max_pages; + + // ---- allocate ---- + bf16 *d_q, *d_o_paged, *d_o_ref; + bf16 *d_k_pool, *d_v_pool; + int64_t* d_pt; + float *d_op, *d_ml; + + size_t sz_q = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16); + size_t sz_o = sz_q; + size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16); + size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t); + int max_splits = 32; + size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float); + size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float); + + cudaMalloc(&d_q, sz_q); + cudaMalloc(&d_o_paged, sz_o); + cudaMalloc(&d_o_ref, sz_o); + cudaMalloc(&d_k_pool, sz_kv); + cudaMalloc(&d_v_pool, sz_kv); + cudaMalloc(&d_pt, sz_pt); + cudaMalloc(&d_op, sz_op); + cudaMalloc(&d_ml, sz_ml); + + // ---- init: deterministic random using seed ---- + srand(seed); + auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; }; + + // Q + bf16* h_q = (bf16*)malloc(sz_q); + for (int i = 0; i < B * Hq * HEAD_DIM; i++) + h_q[i] = __float2bfloat16(rnd()); + cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice); + + // Page pool K/V + bf16* h_k_pool = (bf16*)malloc(sz_kv); + bf16* h_v_pool = (bf16*)malloc(sz_kv); + size_t ps = (size_t)page_size * Hkv * HEAD_DIM; + for (int pg = 0; pg < n_phys_pages; pg++) { + for (int off = 0; off < page_size; off++) { + for (int h = 0; h < Hkv; h++) { + for (int d = 0; d < HEAD_DIM; d++) { + float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d)); + size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM + h * HEAD_DIM + d; + h_k_pool[idx] = __float2bfloat16(v); + h_v_pool[idx] = __float2bfloat16(v * 0.3f); + } + } + } + } + cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice); + cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice); + + // Page table + int64_t* h_pt = (int64_t*)malloc(sz_pt); + int next_pg = 0; + for (int b = 0; b < B; b++) + for (int p = 0; p < max_pages; p++) + h_pt[b * max_pages + p] = next_pg++; + cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice); + + // ---- reference: gather contiguous K/V, then run CPU online-softmax ---- + bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16)); + bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16)); + gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont); + + float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float)); + float sscale = 1.0f / sqrtf((float)HEAD_DIM); + for (int b = 0; b < B; b++) { + for (int hq = 0; hq < Hq; hq++) { + int hkv = hq / G; + size_t q_base = (size_t)b * Hq * HEAD_DIM + (size_t)hq * HEAD_DIM; + size_t kv_base = ((size_t)b * kv_len) * Hkv * HEAD_DIM + (size_t)hkv * HEAD_DIM; + + float m = -1e30f, d = 0.0f; + float acc[256] = {0.0f}; + for (int pos = 0; pos < kv_len; pos++) { + float s = 0.0f; + for (int dim = 0; dim < HEAD_DIM; dim++) + s += __bfloat162float(h_q[q_base + dim]) + * __bfloat162float(h_k_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]); + s *= sscale; + + float nm = fmaxf(m, s); + float a = expf(m - nm); + float b = expf(s - nm); + d = d * a + b; + for (int dim = 0; dim < HEAD_DIM; dim++) + acc[dim] = acc[dim] * a + __bfloat162float(h_v_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]) * b; + m = nm; + } + for (int dim = 0; dim < HEAD_DIM; dim++) + h_o_ref[b * Hq * HEAD_DIM + hq * HEAD_DIM + dim] = acc[dim] / d; + } + } + + // ---- paged decode kernel ---- + float scale_val = 1.0f / sqrtf((float)HEAD_DIM); + PagedAttentionParams p; + p.batch = B; + p.q_head = Hq; + p.kv_head = Hkv; + p.q_len = 1; + p.kv_len = kv_len; + p.head_dim = HEAD_DIM; + p.use_mask = 0; + p.is_causal = 0; + p.causal_offset = 0; + p.num_splits = 1; + p.scale = scale_val; + p.page_size = page_size; + p.max_pages = max_pages; + p.page_table = d_pt; + p.k_cache = d_k_pool; + p.v_cache = d_v_pool; + p.q = d_q; + p.mask = nullptr; + p.o = d_o_paged; + p.o_part = d_op; + p.ml_part = d_ml; + + // Dispatch +#ifndef ASTRAI_NO_MMA + int G_check = p.q_head / p.kv_head; + bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32; + if (use_mma) { + int tiles_total = (p.kv_len + 32 - 1) / 32; + p.num_splits = num_splits(p.batch * p.kv_head, tiles_total); + paged_attn_decode_split_kv_mma_kernel + <<>>(p); + } else +#endif + { + int group_sz = p.q_head / p.kv_head; + int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; + p.num_splits = num_splits(p.batch * p.kv_head, chunks_total); + size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16); + paged_attn_decode_split_kv_kernel<<< + dim3(p.batch * p.kv_head, 1, p.num_splits), + dim3(32, group_sz), smem>>>(p); + } + paged_attn_decode_combine_kernel<<>>(p); + cudaDeviceSynchronize(); + + // Download paged output + bf16* h_o_bf16 = (bf16*)malloc(sz_o); + cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost); + float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float)); + for (int i = 0; i < B * Hq * HEAD_DIM; i++) + h_o_paged[i] = __bfloat162float(h_o_bf16[i]); + + // Compare + float max_err = 0.0f; + int bad_idx = -1; + for (int i = 0; i < B * Hq * HEAD_DIM; i++) { + float e = fabsf(h_o_paged[i] - h_o_ref[i]); + if (e > max_err) { max_err = e; bad_idx = i; } + } + + bool pass = max_err < 0.02f; + + if (pass) { + printf("PASS (max_abs_err=%.4e)\n", max_err); + } else { + int b = bad_idx / (Hq * HEAD_DIM); + int h = (bad_idx / HEAD_DIM) % Hq; + int d = bad_idx % HEAD_DIM; + printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n", + max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]); + // Print first 8 dims of first head + printf(" ref[0..7]:"); + for (int i = 0; i < 8 && i < HEAD_DIM; i++) + printf(" %.4f", h_o_ref[i]); + printf("\n got[0..7]:"); + for (int i = 0; i < 8 && i < HEAD_DIM; i++) + printf(" %.4f", h_o_paged[i]); + printf("\n"); + } + + free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt); + free(h_k_cont); free(h_v_cont); + free(h_o_ref); free(h_o_bf16); free(h_o_paged); + cudaFree(d_q); cudaFree(d_o_paged); cudaFree(d_o_ref); + cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt); + cudaFree(d_op); cudaFree(d_ml); + + return pass ? 0 : 1; +} + +int main() { + int fail = 0; + printf("=== Paged Decode vs CPU reference ===\n\n"); + + printf("-- scalar (G=1) --\n"); + fail += run_test<128>(1, 1, 1, 8, 128, 1); + fail += run_test<128>(1, 4, 4, 128, 128, 2); + fail += run_test<128>(2, 4, 4, 256, 128, 3); + fail += run_test<128>(1, 4, 1, 64, 64, 4); + + printf("-- scalar (G>1) --\n"); + fail += run_test<128>(1, 8, 2, 64, 128, 5); + fail += run_test<128>(2, 16, 4, 128, 128, 6); + + printf("-- varying head_dim --\n"); + fail += run_test<64>(1, 4, 2, 32, 128, 7); + fail += run_test<256>(1, 2, 1, 16, 128, 8); + fail += run_test<32>(1, 4, 2, 32, 64, 9); + + printf("-- multi-batch --\n"); + fail += run_test<128>(3, 8, 2, 256, 128, 10); + fail += run_test<128>(2, 32, 8, 512, 128, 11); + +#ifndef ASTRAI_NO_MMA + printf("-- MMA (G>1, sm_80+) --\n"); + fail += run_test<128>(1, 16, 2, 256, 128, 12); + fail += run_test<128>(2, 32, 4, 512, 128, 13); +#endif + + printf("\n%s (%d/%d failed)\n", fail ? "FAILED" : "ALL PASSED", fail, fail + (13 - fail + 1)); + return fail; +}