feat: paged decode attention with split-KV (scalar + MMA)
- PagedAttentionParams merged into attn_common.h - Scalar variant: warp-per-query-head split-KV, resolves page table per-position for the K/V shared-memory tile load - MMA variant (sm_80+): tensor-core head-packing with cp.async, single page-table lookup per tile (BC=32 fits within page_size>=32) - Standalone test: 14 cases across head_dim 32/64/128/256, GQA, multi-batch, both paths verified against CPU reference
This commit is contained in:
parent
2c0b5d0b5e
commit
89ece26c25
|
|
@ -45,3 +45,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
|
|||
|
||||
register("attn_decode")
|
||||
register("attn_prefill")
|
||||
register("attn_paged_decode")
|
||||
|
|
|
|||
|
|
@ -24,3 +24,31 @@ struct AttentionParams {
|
|||
AT* __restrict__ o_part;
|
||||
AT* __restrict__ ml_part;
|
||||
};
|
||||
|
||||
template<typename T, typename AT = float>
|
||||
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;
|
||||
AT* __restrict__ ml_part;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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 <torch/extension.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
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<bf16>& 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<float>();
|
||||
p.ml_part = ml_part.data_ptr<float>();
|
||||
|
||||
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.batch * p.q_head, p.head_dim>>>(p);
|
||||
}
|
||||
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
template <int HEAD_DIM, int BC>
|
||||
static void launch_paged_mma_decode(PagedAttentionParams<bf16>& 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<float>();
|
||||
p.ml_part = ml_part.data_ptr<float>();
|
||||
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <int HEAD_DIM>
|
||||
static void dispatch_paged_decode(PagedAttentionParams<bf16>& 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<HEAD_DIM, 32>(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<torch::Tensor> mask,
|
||||
bool is_causal = false,
|
||||
int64_t causal_offset = 0,
|
||||
c10::optional<double> 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<float>(scale.value())
|
||||
: 1.0f / std::sqrt(static_cast<float>(head_dim));
|
||||
|
||||
auto O = torch::empty_like(q);
|
||||
|
||||
PagedAttentionParams<bf16, float> p;
|
||||
p.batch = batch;
|
||||
p.q_head = q_head;
|
||||
p.kv_head = kv_head;
|
||||
p.q_len = 1;
|
||||
p.kv_len = static_cast<int>(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<int>(causal_offset);
|
||||
p.num_splits = 1;
|
||||
p.scale = scale_val;
|
||||
p.page_size = static_cast<int>(page_size);
|
||||
p.max_pages = max_pages;
|
||||
p.page_table = page_table.data_ptr<int64_t>();
|
||||
p.k_cache = reinterpret_cast<const bf16*>(k_cache.data_ptr());
|
||||
p.v_cache = reinterpret_cast<const bf16*>(v_cache.data_ptr());
|
||||
p.q = reinterpret_cast<const bf16*>(q.data_ptr());
|
||||
p.mask = p.use_mask ? mask.value().data_ptr<bool>() : nullptr;
|
||||
p.o = reinterpret_cast<bf16*>(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.");
|
||||
}
|
||||
|
|
@ -0,0 +1,140 @@
|
|||
#pragma once
|
||||
#include <cuda_bf16.h>
|
||||
#include <float.h>
|
||||
#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<bf16> 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<bf16> 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);
|
||||
}
|
||||
|
|
@ -0,0 +1,174 @@
|
|||
#pragma once
|
||||
#include <cfloat>
|
||||
#include <cuda_bf16.h>
|
||||
#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 <int HEAD_DIM, int BC>
|
||||
__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16;
|
||||
constexpr int NC8 = BC / 8;
|
||||
constexpr int KT2 = BC / 16;
|
||||
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<KD, NC8>(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<NC8, DN8>(kv0, maxc, maxc,
|
||||
mask_base, p.mask, has_mask,
|
||||
Sacc, Oacc, m0, m1, l0, l1, lane);
|
||||
|
||||
mma_pv_accumulate<DN8, KT2>(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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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 <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cassert>
|
||||
#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 <int HEAD_DIM>
|
||||
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<bf16, float> 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<HEAD_DIM, 32>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(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.batch * p.q_head, p.head_dim>>>(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;
|
||||
}
|
||||
Loading…
Reference in New Issue