ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: SGLang-style paged attention kernels replace page-table path

Browse Source 
- PagedAttentionParams uses flat KV pool + req_to_token + kv_indptr/qo_indptr instead of page_table
- MMA split-KV decode and split-Q prefill kernels with indirect ragged-batch addressing
- Prefill kernel accepts 4D mask (causal-aware); decode kernel supports 2D mask
- CudaBackend is inference-only: kv_cache=None raises, no torch fallback
- benchmark.py: required --ckpt, --backend/--compare options
- Parallel build isolates build-temp/build-lib per subprocess
- Standalone test covers decode/prefill with mask, 27 cases pass
This commit is contained in:
ViperEkura
2026-08-01 15:41:25 +08:00
parent 9960f79920
commit 41dcf0feb9
17 changed files with 1683 additions and 535 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+50 -64
View File
@@ -38,8 +38,10 @@ import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.attention_ops import attn_paged_decode, attn_prefill
from astrai.extension.loader import is_available
from astrai.extension.attention_ops import (
attn_paged_decode,
attn_paged_prefill,
)
from astrai.inference.core.cache import KVCache
_current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
@@ -307,24 +309,19 @@ _default_backend = TorchNativeBackend()
class CudaBackend(AttentionBackend):
"""CUDA kernel backend with direct KV cache access.
Decode path: writes K/V to cache, then calls ``attn_paged_decode``
with ``page_size=1`` (each token slot is a single-token "page").
The ``req_to_token`` table serves directly as the page table.
Decode path: writes K/V to the flat pool, then calls
``attn_paged_decode`` with req_to_token + kv_indptr.
Prefill path: writes K/V to cache, gathers full-sequence K/V via
indirect indexing (same as TorchNativeBackend), then calls
``attn_prefill``.
Prefill path: writes K/V to the flat pool, then calls
``attn_paged_prefill`` with ragged-batch support via qo_indptr +
kv_indptr.
Training path (``kv_cache is None``): calls ``attn_prefill`` directly
on the projected q/k/v.
``kv_cache is None`` (training) is not handled — use
``TorchNativeBackend`` for training.
Falls back to ``TorchNativeBackend`` for any path where the
corresponding CUDA kernel is not available.
Raises ``RuntimeError`` if the required kernel is not available.
"""
def __init__(self):
self._fallback = TorchNativeBackend()
def fwd_decode(
self,
q: Tensor,
@@ -335,47 +332,34 @@ class CudaBackend(AttentionBackend):
attn_mask: Optional[Tensor] = None,
is_causal: bool = False,
) -> Tensor:
if kv_cache is None or not is_available("attn_paged_decode"):
return self._fallback.fwd_decode(
q, k, v, kv_cache, layer_id, attn_mask, is_causal
)
if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.max_len
b = q.size(0)
q_3d = q.squeeze(1)
if kv_cache.page_table is not None:
page_table = kv_cache.page_table
else:
page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
if q.size(0) == 1:
mask = None
elif kv_cache.decode_mask is not None:
if b > 1 and kv_cache.decode_mask is not None:
mask = kv_cache.decode_mask
else:
mask = (
torch.arange(max_len, device=q.device)[None, :]
< kv_cache.seq_lens[:, None]
)
out = attn_paged_decode(
q,
page_table,
k_cache,
v_cache,
page_size=1,
kv_len=max_len,
q_3d,
kv_cache.k_buffer[layer_id],
kv_cache.v_buffer[layer_id],
kv_cache.req_to_token,
kv_cache.req_pool_indices,
kv_indptr,
kv_cache.max_len,
mask=mask,
is_causal=is_causal,
)
out = out.flatten(2)
return out
return out.unsqueeze(1).flatten(2)
def fwd_prefill(
self,
@@ -388,32 +372,34 @@ class CudaBackend(AttentionBackend):
is_causal: bool = False,
) -> Tensor:
if kv_cache is None:
if is_available("attn_prefill"):
out = attn_prefill(q, k, v, mask=attn_mask, is_causal=is_causal)
return out.flatten(2)
return self._fallback.fwd_prefill(
q, k, v, kv_cache, layer_id, attn_mask, is_causal
)
if not is_available("attn_prefill"):
return self._fallback.fwd_prefill(
q, k, v, kv_cache, layer_id, attn_mask, is_causal
)
raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.max_len
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
k_full = kv_cache.k_buffer[layer_id, indices]
v_full = kv_cache.v_buffer[layer_id, indices]
b = q.size(0)
q_len = q.size(1)
out = attn_prefill(q, k_full, v_full, mask=attn_mask, is_causal=is_causal)
return out.flatten(2)
kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
qo_indptr = torch.arange(b + 1, dtype=torch.int32, device=q.device) * q_len
q_flat = q.reshape(b * q_len, q.size(2), q.size(3))
out = attn_paged_prefill(
q_flat,
kv_cache.k_buffer[layer_id],
kv_cache.v_buffer[layer_id],
kv_cache.req_to_token,
kv_cache.req_pool_indices,
kv_indptr,
qo_indptr,
attn_mask,
q_len,
is_causal=is_causal,
)
return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2)
_BACKEND_REGISTRY: dict[ATTN_BACKEND, type[AttentionBackend]] = {
astrai/extension/attention_ops.py
+71 -16
View File
@@ -92,39 +92,94 @@ def attn_prefill(
def attn_paged_decode(
q: torch.Tensor,
page_table: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
page_size: int,
kv_len: int,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
kv_indptr: torch.Tensor,
max_seq_len: int,
mask: Optional[torch.Tensor] = None,
is_causal: bool = False,
) -> torch.Tensor:
"""Paged GQA decode attention (q_len == 1, direct page-table access).
"""SGLang-style paged decode (q_len == 1, flat KV pool).
Reads K/V directly from a flat pool [size, kv_head, head_dim] via
req_to_token indirect indexing. Each request has its own seq_len
(from kv_indptr), eliminating padding waste.
Args:
q: [batch, 1, n_heads, head_dim] (blhd, bf16)
page_table: [batch, max_pages] (int64)
k_cache: [n_pages, page_size, n_kv_heads, head_dim] (bf16)
q: [batch, n_heads, head_dim] (bf16, 3D — no seq dim)
k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
v_cache: same as k_cache
page_size: tokens per page
kv_len: actual sequence length per request
mask: 2D [batch, kv_len] or 3D [batch, 1, kv_len] (bool, True=keep)
req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
req_pool_indices: [batch] (int64) — rows into req_to_token
kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
max_seq_len: max per-request seq_len (Python int, for split computation)
mask: 2D [batch, max_seq_len] (bool, True=keep) or None
is_causal: apply causal mask
Returns:
[batch, 1, n_heads, head_dim] (blhd, bf16)
[batch, n_heads, head_dim] (bf16, 3D)
"""
_check_available("attn_paged_decode")
causal_offset = (kv_len - 1) if is_causal else -1
causal_offset = 0 if is_causal else -1
return _modules["attn_paged_decode"].attn_paged_decode(
q,
page_table,
k_cache,
v_cache,
page_size,
kv_len,
req_to_token,
req_pool_indices,
kv_indptr,
max_seq_len,
mask=mask,
causal_offset=causal_offset,
layout=TensorLayout.BLHD,
)
def attn_paged_prefill(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
kv_indptr: torch.Tensor,
qo_indptr: torch.Tensor,
mask: Optional[torch.Tensor] = None,
max_q_len: int = 0,
is_causal: bool = False,
) -> torch.Tensor:
"""SGLang-style paged prefill (ragged batch, flat KV pool).
Reads K/V directly from a flat pool [size, kv_head, head_dim] via
req_to_token. Supports ragged batches: each request has its own
q_len and kv_len, addressed via qo_indptr and kv_indptr.
Args:
q: [total_q, n_heads, head_dim] (bf16, 3D — flattened across requests)
k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
v_cache: same as k_cache
req_to_token: [num_reqs, max_context_len] (int64)
req_pool_indices: [batch] (int64)
kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
max_q_len: max per-request q_len (Python int, for grid computation)
is_causal: apply causal mask
Returns:
[total_q, n_heads, head_dim] (bf16, 3D)
"""
_check_available("attn_paged_prefill")
causal_offset = 0 if is_causal else -1
return _modules["attn_paged_prefill"].attn_paged_prefill(
q,
k_cache,
v_cache,
req_to_token,
req_pool_indices,
kv_indptr,
qo_indptr,
mask,
max_q_len,
causal_offset=causal_offset,
)
astrai/extension/loader.py
+7 -1
View File
@@ -11,7 +11,13 @@ import logging
logger = logging.getLogger(__name__)
KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode", "rotary_emb"]
KERNEL_NAMES = [
"attn_decode",
"attn_prefill",
"attn_paged_decode",
"attn_paged_prefill",
"rotary_emb",
]
_available: dict[str, bool] = {}
_modules: dict[str, object] = {}
csrc/build.py
+1
View File
@@ -72,4 +72,5 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
register("attn_decode")
register("attn_prefill")
register("attn_paged_decode")
register("attn_paged_prefill")
register("rotary_emb")
csrc/kernels/attn_common.h
+34 -14
View File
@@ -43,35 +43,55 @@ struct AttentionParams {
AT* __restrict__ ml_part;
};
// ---- PagedAttentionParams ----
// SGLang-style indirect params over a shared KV pool.
// k_cache/v_cache: [size, kv_head, head_dim] (bare buffers, no gather).
// req_to_token: [num_reqs, max_context_len] token -> slot.
// req_pool_indices:[batch] rows of the current batch into req_to_token.
// kv_indptr: [batch+1] prefix sum of per-request seq_lens (device).
// qo_indptr: [batch+1] prefix sum of per-request q_len (prefill) or
// nullptr for decode (q_len == 1 everywhere).
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 num_splits;
int use_mask;
int causal_offset;
int causal_offset; // -1 = non-causal; >=0 = causal (per-request offset
// computed inside kernel from kv_indptr/qo_indptr)
float scale;
int num_splits;
int page_size;
int max_pages;
// Q: [total_q, q_head, head_dim] (3D flattened — no batch dim).
// For decode total_q == batch (q_len=1 per request).
// For prefill total_q == qo_indptr[batch].
int q_stride_l, q_stride_h, q_stride_d;
// Q strides (layout-agnostic)
int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
// Q: [total_q, q_head, head_dim]
const T* __restrict__ q;
// Mask strides (2D, 3D, or 4D)
// Flat KV pool: [size, kv_head, head_dim]
const T* __restrict__ k_cache;
const T* __restrict__ v_cache;
// Indexing
const int64_t* __restrict__ req_to_token; // [num_reqs, max_context_len]
const int64_t* __restrict__ req_pool_indices; // [batch]
const int* __restrict__ kv_indptr; // [batch+1]
const int* __restrict__ qo_indptr; // [batch+1] or nullptr (decode)
int max_context_len; // req_to_token stride (dim 1)
int max_seq_len; // max per-request seq_len (host-side, for split computation)
int total_q; // total Q tokens across all requests (host-side, for grid)
int max_q_len; // max per-request q_len (host-side, for prefill grid)
// Mask: [batch, max_seq_len] (decode) or [batch, 1, q_len, kv_len]
// (prefill, optional). mask_h_stride/mask_q_stride are 0 when those
// dims are size 1 (broadcast).
int mask_b_stride;
int mask_h_stride;
int mask_q_stride;
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;
csrc/kernels/attn_dispatchers.cuh
+35 -7
View File
@@ -12,6 +12,7 @@
#include "attn_prefill_split_q_mma.cuh"
#include "attn_decode_split_kv_mma.cuh"
#include "attn_paged_decode_split_kv_mma.cuh"
#include "attn_paged_prefill_split_q_mma.cuh"
#endif
// Cached SM count — cudaDeviceGetAttribute is a host-side call that was
@@ -145,18 +146,18 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
}
// ======================================================================
// Paged Decode
// Paged Decode (SGLang-style: flat pool + req_to_token + kv_indptr)
// ======================================================================
#ifndef ASTRAI_NO_MMA
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int group_size) {
static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
int G = p.q_head / p.kv_head;
constexpr int MAX_G = 16;
constexpr int BC = 16;
int num_passes = (G + MAX_G - 1) / MAX_G;
int tiles_total = (p.kv_len + BC - 1) / BC;
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total, 2);
int tiles_total = (p.max_seq_len + BC - 1) / BC;
p.num_splits = compute_num_splits(p.batch * p.kv_head * num_passes, tiles_total, 2);
constexpr int STAGES = 2;
using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
@@ -166,10 +167,10 @@ static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int gr
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size) {
int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
int chunks_total = (p.max_seq_len + PDC_CHUNK - 1) / PDC_CHUNK;
p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit
int g = min(group_size, 32);
dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
dim3 block(32, g);
paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
@@ -182,10 +183,37 @@ static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
int group_size = p.q_head / p.kv_head;
#ifndef ASTRAI_NO_MMA
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, group_size);
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, 0);
#else
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
#endif
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
}
// ======================================================================
// Paged Prefill (SGLang-style: flat pool + ragged batch)
// ======================================================================
#ifndef ASTRAI_NO_MMA
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p) {
constexpr int WARPS = 4;
constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
int max_q_tiles = (p.max_q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS);
dim3 grid(max_q_tiles, p.q_head, p.batch);
dim3 block(Traits::NUM_THREADS);
paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
}
#endif
template <int HEAD_DIM>
static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p) {
bool is_causal = (p.causal_offset >= 0);
bool has_mask = (p.use_mask && p.mask);
#ifndef ASTRAI_NO_MMA
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p);
#endif
}
csrc/kernels/attn_entry_utils.cuh
+147 -23
View File
@@ -134,54 +134,178 @@ inline void attn_pack_params(
pack_mask(mask, p);
}
// ---- attn_pack_paged_params ----
// ---- attn_pack_paged_decode_params ----
// SGLang-style: flat KV pool + req_to_token indexing + variable
// seq_lens via kv_indptr. Q is [batch, q_head, head_dim] (q_len=1 per req).
template<typename T>
inline void attn_pack_paged_params(
inline void attn_pack_paged_decode_params(
torch::Tensor q,
torch::Tensor page_table,
torch::Tensor k_cache,
torch::Tensor v_cache,
int64_t page_size,
int64_t kv_len,
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
int64_t max_seq_len,
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
int64_t layout,
PagedAttentionParams<T>& p
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda() && kv_indptr.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(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
TORCH_CHECK(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
TORCH_CHECK(q.dim() == 3, "q must be 3D [batch, q_head, head_dim]");
extract_q_dims_and_strides(q, layout, p);
p.kv_head = (int)k_cache.size(2);
p.kv_len = (int)kv_len;
p.page_size = (int)page_size;
p.max_pages = (int)page_table.size(1);
TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
p.batch = (int)q.size(0);
p.q_head = (int)q.size(1);
p.head_dim = (int)q.size(2);
p.kv_head = (int)k_cache.size(1);
TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
TORCH_CHECK(p.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(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
p.q_stride_l = (int)q.stride(0);
p.q_stride_h = (int)q.stride(1);
p.q_stride_d = (int)q.stride(2);
p.k_cache = (const T*)k_cache.data_ptr();
p.v_cache = (const T*)v_cache.data_ptr();
p.q = (const T*)q.data_ptr();
p.req_to_token = req_to_token.data_ptr<int64_t>();
p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
p.kv_indptr = kv_indptr.data_ptr<int>();
p.qo_indptr = nullptr;
p.max_context_len = (int)req_to_token.size(1);
p.max_seq_len = (int)max_seq_len;
p.total_q = p.batch; // decode: 1 Q token per request
p.max_q_len = 1;
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
p.page_table = page_table.data_ptr<int64_t>();
if (p.use_mask) {
auto m = mask.value();
TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_h_stride = 0;
p.mask_q_stride = 0;
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_h_stride = 0;
p.mask_q_stride = 0;
}
p.o = nullptr;
p.o_part = nullptr;
p.ml_part = nullptr;
}
// ---- attn_pack_paged_prefill_params ----
// SGLang-style: flat KV pool + req_to_token + ragged batch via qo_indptr.
// Q is [total_q, q_head, head_dim] (flattened across all requests).
template<typename T>
inline void attn_pack_paged_prefill_params(
torch::Tensor q,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
torch::Tensor qo_indptr,
c10::optional<torch::Tensor> mask,
int64_t max_q_len,
int64_t causal_offset,
double scale,
PagedAttentionParams<T>& p
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda());
TORCH_CHECK(kv_indptr.is_cuda() && qo_indptr.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(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be int32");
TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
TORCH_CHECK(q.dim() == 3, "q must be 3D [total_q, q_head, head_dim]");
p.q_head = (int)q.size(1);
p.head_dim = (int)q.size(2);
p.kv_head = (int)k_cache.size(1);
p.batch = (int)req_pool_indices.size(0);
TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]");
TORCH_CHECK(qo_indptr.size(0) == p.batch + 1, "qo_indptr must be [batch+1]");
p.q_stride_l = (int)q.stride(0);
p.q_stride_h = (int)q.stride(1);
p.q_stride_d = (int)q.stride(2);
p.k_cache = (const T*)k_cache.data_ptr();
p.v_cache = (const T*)v_cache.data_ptr();
p.q = (const T*)q.data_ptr();
p.req_to_token = req_to_token.data_ptr<int64_t>();
p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
p.kv_indptr = kv_indptr.data_ptr<int>();
p.qo_indptr = qo_indptr.data_ptr<int>();
p.max_context_len = (int)req_to_token.size(1);
p.total_q = (int)q.size(0); // prefill: flattened Q across all requests
p.max_q_len = (int)max_q_len;
// max_seq_len is unused by the prefill path (decode uses it for split
// computation); fill with max_q_len only to keep the POD struct defined.
p.max_seq_len = p.max_q_len;
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
if (p.use_mask) {
auto m = mask.value();
TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
if (m.dim() == 2) {
TORCH_CHECK(m.size(1) <= p.max_context_len, "mask kv_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_h_stride = 0;
p.mask_q_stride = 0;
} else if (m.dim() == 4) {
TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch");
TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch");
TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
} else {
TORCH_CHECK(false, "mask must be 2D or 4D");
}
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_h_stride = 0;
p.mask_q_stride = 0;
}
p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
p.o = nullptr;
p.o_part = nullptr;
p.ml_part = nullptr;
pack_mask(mask, p);
}
csrc/kernels/attn_paged_decode.cu
+15 -15
View File
@@ -3,23 +3,23 @@
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,
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
int64_t max_seq_len,
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
int64_t layout
double scale
) {
PagedAttentionParams<bf16> p;
attn_pack_paged_params(q, page_table, k_cache, v_cache,
page_size, kv_len, mask, causal_offset, scale, layout, p);
attn_pack_paged_decode_params(q, k_cache, v_cache,
req_to_token, req_pool_indices, kv_indptr,
max_seq_len, mask, causal_offset, scale, p);
auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
p.o = (bf16*)O.data_ptr();
alloc_split_partials(p);
DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
@@ -30,14 +30,14 @@ torch::Tensor attn_paged_decode(
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("req_to_token"),
py::arg("req_pool_indices"),
py::arg("kv_indptr"),
py::arg("max_seq_len"),
py::arg("mask") = py::none(),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
py::arg("layout") = (int64_t)BHLD,
"Paged GQA decode — split-KV with direct page-table access.");
"SGLang-style paged decode: flat KV pool + req_to_token + kv_indptr.");
}
csrc/kernels/attn_paged_decode_split_kv.cuh
+19 -20
View File
@@ -5,6 +5,8 @@
#include "attn_warp_utils.cuh"
constexpr int PDC_CHUNK = 64;
// Scalar paged decode (fallback for sm < 80, no tensor cores).
// Reads K/V from flat pool via req_to_token indexing.
template <int HEAD_DIM, bool IsCausal, bool HasMask>
__global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
int batch = blockIdx.x / p.kv_head;
@@ -15,8 +17,11 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
int lane = threadIdx.x;
int hd_per_thread = p.head_dim / 32;
const int seq_len = p.kv_indptr[batch + 1] - p.kv_indptr[batch];
const int64_t req_idx = p.req_pool_indices[batch];
float q_reg[8];
int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
int q_off = batch * p.q_stride_l + q_head * p.q_stride_h
+ lane * hd_per_thread * p.q_stride_d;
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
@@ -26,16 +31,19 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
extern __shared__ __align__(16) bf16 k_smem[];
int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
int chunks_total = (seq_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.mask_b_stride + q_head * p.mask_h_stride;
const int mask_base = batch * p.mask_b_stride;
const int64_t pool_stride = (int64_t)p.kv_head * p.head_dim;
const int64_t head_off = (int64_t)kv_head * p.head_dim;
const int64_t rtt_stride = (int64_t)p.max_context_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 this_chunk = min(PDC_CHUNK, seq_len - chunk_start);
int total = this_chunk * p.head_dim;
for (int i = threadIdx.y * 32 + lane; i < total;
@@ -43,14 +51,9 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
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;
int64_t slot = p.req_to_token[req_idx * rtt_stride + pos];
if (slot >= 0) {
int64_t off = slot * pool_stride + head_off + d_dim;
k_smem[i] = p.k_cache[off];
} else {
k_smem[i] = __float2bfloat16(0.0f);
@@ -85,17 +88,13 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
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];
int64_t slot = p.req_to_token[req_idx * rtt_stride + pos];
if (masked) {
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
} else 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;
} else if (slot >= 0) {
int64_t v_base = slot * pool_stride + head_off;
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = fmaf(acc_reg[i], alpha,
@@ -148,6 +147,6 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
}
float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
int o_off = batch * p.q_stride_l + q_head * p.q_stride_h + d * p.q_stride_d;
p.o[o_off] = __float2bfloat16(acc * inv);
}
csrc/kernels/attn_paged_decode_split_kv_mma.cuh
+31 -28
View File
@@ -5,12 +5,16 @@
#include "attn_mma_utils.cuh"
#include "attn_warp_utils.cuh"
// Paged split-KV tensor-core decode via GQA head-packing.
// Reads K/V directly from the page pool through a page table — one tile
// (BC=32) fits within a single page (page_size >= 32), so the page-table
// lookup happens once per tile for cp.async.
// SGLang-style split-KV tensor-core decode.
//
// IsCausal and HasMask are compile-time bools.
// Reads K/V directly from a flat pool [size, kv_head, head_dim] via
// req_to_token indexing — no gather, no page-table dimension.
// Each batch element has its own seq_len (from kv_indptr), eliminating
// padding waste: short sequences only process the tiles they own.
//
// For decode (q_len=1), causal masking is implicit — each request attends
// to [0, seq_len) which is exactly its valid range. The IsCausal flag
// is accepted for dispatch uniformity but does not change maxc.
template <typename Traits, bool IsCausal, bool HasMask>
__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
const int lane = threadIdx.x;
@@ -22,6 +26,10 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
const int batch = blockIdx.y;
const int split = blockIdx.z;
// Per-request seq_len from device-side kv_indptr — no padding.
const int seq_len = p.kv_indptr[batch + 1] - p.kv_indptr[batch];
const int64_t req_idx = p.req_pool_indices[batch];
constexpr int MAX_G = 16;
const int G_total = p.q_head / p.kv_head;
const int g_begin = pass * MAX_G;
@@ -38,12 +46,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
}
__syncwarp();
const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
const int q_base = batch * p.q_stride_l + q_head0 * p.q_stride_h;
const int qra = gid;
const int qrb = gid + 8;
const bool va = qra < G, vb = qrb < G;
unsigned Qa[Traits::KD][4];
load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
load_q_mma_frags<Traits::KD>(p.q + q_base,
p.q_stride_h, p.q_stride_d,
qra, qrb, va, vb, tid4, Qa);
float Oacc[Traits::DN8][4];
@@ -52,19 +61,19 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
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 tiles_total = (p.kv_len + Traits::BC - 1) / Traits::BC;
const int tiles_total = (seq_len + Traits::BC - 1) / Traits::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 int64_t page_stride = (int64_t)p.page_size * p.kv_head * Traits::HEAD_DIM;
const int64_t pos_stride = (int64_t)p.kv_head * Traits::HEAD_DIM;
// Flat pool stride: [size, kv_head, head_dim] — contiguous.
const int64_t pool_stride = (int64_t)p.kv_head * Traits::HEAD_DIM;
const int64_t head_off = (int64_t)kv_head * Traits::HEAD_DIM;
const int64_t rtt_stride = (int64_t)p.max_context_len;
// ---- Load tile lambda: paged addressing ----
// Unified per-element page-table lookup. When page_size >= BC, all
// elements in a tile share the same page, so the lookup is redundant
// but harmless (L1-cached). This avoids a branch on page_size.
// ---- Load tile lambda: SGLang addressing ----
// slot = req_to_token[req_idx * max_context_len + kc]
// gmem = k_cache[slot * pool_stride + head_off + d]
auto load_tile = [&](int ti, int buf) {
int kv0 = ti * Traits::BC;
bf16* dK = sK + buf * Traits::BC * Traits::LD;
@@ -74,16 +83,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
i += Traits::NUM_THREADS * Traits::VEC) {
int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
int kc = kv0 + r;
bool valid = (kc < p.kv_len);
bool valid = (kc < seq_len);
if constexpr (HasMask) {
valid = valid && p.mask[batch * p.mask_b_stride + kc];
}
int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
valid = valid && (phys_page >= 0);
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;
int64_t slot = valid ? p.req_to_token[req_idx * rtt_stride + kc] : 0;
valid = valid && (slot >= 0);
int64_t gmem_base = slot * pool_stride + head_off;
int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
@@ -91,10 +97,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
cp_async_commit();
};
// ---- Multi-stage cp.async pipeline ----
// Prologue loads STAGES tiles; each loop iteration waits only for the
// oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
// tile loads stay in flight and overlap with the current tile's compute.
constexpr int STAGES = Traits::STAGES;
const int ntiles = ti_end - ti_begin;
@@ -111,8 +113,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
// For decode, maxc = seq_len regardless of IsCausal — the valid
// range [0, seq_len) IS the causal range (query is the last token).
mma_softmax_tile<Traits, HasMask>(kv0, seq_len, seq_len,
0, 0,
p.mask_b_stride, 0, 0,
batch, 0,
@@ -136,7 +139,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
}
} else {
// Fewer tiles than stages: load all, wait for all, process.
for (int i = 0; i < ntiles; i++)
load_tile(ti_begin + i, i);
cp_async_wait_group<0>();
@@ -145,6 +147,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
process_tile(it, it);
}
// ---- write partials ----
auto split_slot = [&](int h) -> size_t {
size_t bh = (size_t)batch * p.q_head + h;
return bh * MAX_SPLITS + split;
csrc/kernels/attn_paged_prefill.cu
+45
View File
@@ -0,0 +1,45 @@
#include "attn_dispatchers.cuh"
#include "attn_entry_utils.cuh"
torch::Tensor attn_paged_prefill(
torch::Tensor q,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor req_to_token,
torch::Tensor req_pool_indices,
torch::Tensor kv_indptr,
torch::Tensor qo_indptr,
c10::optional<torch::Tensor> mask,
int64_t max_q_len,
int64_t causal_offset,
double scale
) {
PagedAttentionParams<bf16> p;
attn_pack_paged_prefill_params(q, k_cache, v_cache,
req_to_token, req_pool_indices,
kv_indptr, qo_indptr, mask,
max_q_len, causal_offset, scale, p);
auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
p.o = (bf16*)O.data_ptr();
DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p);
C10_CUDA_CHECK(cudaGetLastError());
return O;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("attn_paged_prefill", &attn_paged_prefill,
py::arg("q"),
py::arg("k_cache"),
py::arg("v_cache"),
py::arg("req_to_token"),
py::arg("req_pool_indices"),
py::arg("kv_indptr"),
py::arg("qo_indptr"),
py::arg("mask") = py::none(),
py::arg("max_q_len"),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
"SGLang-style paged prefill: flat KV pool + ragged batch.");
}
csrc/kernels/attn_paged_prefill_split_q_mma.cuh
+164
View File
@@ -0,0 +1,164 @@
#pragma once
#include <cfloat>
#include <cuda_bf16.h>
#include "attn_common.h"
#include "attn_mma_utils.cuh"
// SGLang-style split-Q tensor-core prefill.
//
// Reads K/V directly from a flat pool [size, kv_head, head_dim] via
// req_to_token — no gather, no temporary tensor. Supports ragged batches:
// each request has its own q_len and kv_len, addressed via qo_indptr and
// kv_indptr.
//
// Grid: (max_q_tiles, q_head, batch) — one batch element per blockIdx.z.
// Blocks beyond a request's q_len exit early after writing sentinel-free
// no-ops. This avoids the binary-search approach and guarantees every Q
// token is covered, even when q_len < BR*WARPS (e.g. decode-like prefill).
//
// Q layout: [total_q, q_head, head_dim] (3D, flattened across requests).
// O layout: same as Q.
//
// IsCausal is a compile-time bool. When true, each Q row qi (within its
// request) attends to [0, causal_offset_b + qi + 1) where
// causal_offset_b = kv_len_b - q_len_b (position of first Q token).
template <typename Traits, bool IsCausal, bool HasMask>
__global__ void paged_attn_prefill_split_q_mma_kernel(PagedAttentionParams<bf16> p) {
const int warp = threadIdx.x / 32;
const int lane = threadIdx.x % 32;
const int gid = lane >> 2;
const int tid4 = lane & 3;
const int q_head = blockIdx.y;
const int req_b = blockIdx.z;
const int qrow0 = (blockIdx.x * Traits::WARPS + warp) * Traits::BR;
const int seq_len = p.kv_indptr[req_b + 1] - p.kv_indptr[req_b];
const int q_len = p.qo_indptr[req_b + 1] - p.qo_indptr[req_b];
const int causal_off = seq_len - q_len;
const int64_t req_idx = p.req_pool_indices[req_b];
// No per-warp early exit — all warps must participate in __syncthreads.
// Warps beyond q_len get zero-filled Q frags (va=vb=false) and skip output.
const int kv_head = q_head / (p.q_head / p.kv_head);
__shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
__shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
// Q base: offset by qo_indptr[req_b] to get absolute token address.
const int q_base = p.qo_indptr[req_b] * p.q_stride_l + q_head * p.q_stride_h;
const int qra = qrow0 + gid;
const int qrb = qrow0 + gid + 8;
const bool va = qra < q_len, vb = qrb < q_len;
unsigned Qa[Traits::KD][4];
load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
qra, qrb, va, vb, tid4, Qa);
float Oacc[Traits::DN8][4];
#pragma unroll
for (int j = 0; j < Traits::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 int64_t pool_stride = (int64_t)p.kv_head * Traits::HEAD_DIM;
const int64_t head_off = (int64_t)kv_head * Traits::HEAD_DIM;
const int64_t rtt_stride = (int64_t)p.max_context_len;
const int tiles = (seq_len + Traits::BC - 1) / Traits::BC;
const int qr0 = qrow0 + gid;
const int qr1 = qrow0 + gid + 8;
// Causal tile-skip (dead code when IsCausal == false)
const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
const int block_max_kv =
blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
+ causal_off;
int t_end = tiles - 1;
if constexpr (IsCausal) {
int bt = block_max_kv / Traits::BC;
if (bt < t_end) t_end = bt;
}
// ---- Load tile lambda: SGLang addressing ----
auto load_tile = [&](int ti, int buf) {
int kv0 = ti * Traits::BC;
bf16* dK = sK + buf * Traits::BC * Traits::LD;
bf16* dV = sV + buf * Traits::BC * Traits::LD;
#pragma unroll
for (int i = threadIdx.x * Traits::VEC; i < Traits::TOTAL;
i += Traits::NUM_THREADS * Traits::VEC) {
int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
int kc = kv0 + r;
bool valid = kc < seq_len;
int64_t slot = valid ? p.req_to_token[req_idx * rtt_stride + kc] : 0;
valid = valid && (slot >= 0);
int64_t gmem_base = slot * pool_stride + head_off;
int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
}
cp_async_commit();
};
// ---- Prologue + main loop (FA2-style double-buffer) ----
load_tile(0, 0);
for (int ti = 0; ti <= t_end; ti++) {
int buf = ti & 1;
cp_async_wait_group<0>();
__syncthreads();
if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
const bf16* bK = sK + buf * Traits::BC * Traits::LD;
const bf16* bV = sV + buf * Traits::BC * Traits::LD;
int kv0 = ti * Traits::BC;
if (!IsCausal || kv0 <= max_kv) {
float Sacc[Traits::NC8][4];
mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
#pragma unroll
for (int n8 = 0; n8 < Traits::NC8; n8++)
Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
int maxc0 = IsCausal ? min(seq_len, causal_off + qr0 + 1)
: seq_len;
int maxc1 = IsCausal ? min(seq_len, causal_off + qr1 + 1)
: seq_len;
// HasMask: mask[batch, q_head, qi, kc] — kc is request-local.
mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
qr0, qr1,
p.mask_b_stride, p.mask_h_stride,
p.mask_q_stride,
req_b, q_head,
p.mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
}
}
// ---- write output: packed bf16x2 stores ----
float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
const int o_base = p.qo_indptr[req_b] * p.q_stride_l + q_head * p.q_stride_h;
#pragma unroll
for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
int d = dn8 * 8 + 2 * tid4;
if (qr0 < q_len) {
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
Oacc[dn8][1] * rl0);
*reinterpret_cast<__nv_bfloat162*>(
&p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
}
if (qr1 < q_len) {
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
Oacc[dn8][3] * rl1);
*reinterpret_cast<__nv_bfloat162*>(
&p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
}
}
}
csrc/tests/attn_paged_decode_test.cu
-308
View File
@@ -1,308 +0,0 @@
// Compile:
// nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
// --extra-device-vectorization csrc/tests/attn_paged_decode_test.cu \
// -o /tmp/test_paged && /tmp/test_paged
#include <cstring>
#include "test_utils.cuh"
#include "../kernels/attn_dispatchers.cuh"
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 * Hkv + h) * kv_len * head_dim
+ (size_t)pos * 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 causal, int seed) {
printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d causal=%d ... ",
B, Hq, Hkv, kv_len, page_size, HEAD_DIM, causal);
fflush(stdout);
int max_pages = (kv_len + page_size - 1) / page_size;
int n_phys_pages = B * max_pages;
int max_splits = 32;
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);
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);
bf16 *d_q, *d_o_paged;
bf16 *d_k_pool, *d_v_pool;
int64_t* d_pt;
float *d_op, *d_ml;
cudaMalloc(&d_q, sz_q);
cudaMalloc(&d_o_paged, 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);
srand(seed);
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
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);
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);
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);
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_q_f = (float*)malloc((size_t)B * Hq * HEAD_DIM * sizeof(float));
float* h_k_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
float* h_v_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
for (int i = 0; i < B * kv_len * Hkv * HEAD_DIM; i++) {
h_k_f[i] = bf2f(h_k_cont[i]);
h_v_f[i] = bf2f(h_v_cont[i]);
}
float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv,
1, kv_len, HEAD_DIM, causal ? 0 : -1);
PagedAttentionParams<bf16> 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.causal_offset = causal ? 0 : -1;
set_default_paged_strides(p);
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
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_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
cudaDeviceSynchronize();
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]);
float max_abs_err = 0.0f, max_rel_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_abs_err) { max_abs_err = e; bad_idx = i; }
float rel = e / fmaxf(fabsf(h_o_ref[i]), 1e-8f);
if (rel > max_rel_err) max_rel_err = rel;
}
const float atol = 0.01f, rtol = 0.01f;
bool pass = true;
for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
float e = fabsf(h_o_paged[i] - h_o_ref[i]);
if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
}
if (pass) {
printf("PASS (max_abs_err=%.4e max_rel_err=%.4e)\n", max_abs_err, max_rel_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 max_rel_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
max_abs_err, max_rel_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
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_q_f); free(h_k_f); free(h_v_f);
free(h_o_ref); free(h_o_bf16); free(h_o_paged);
cudaFree(d_q); cudaFree(d_o_paged);
cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
cudaFree(d_op); cudaFree(d_ml);
return pass ? 0 : 1;
}
struct TestCase {
int head_dim;
int B, Hq, Hkv, kv_len, page_size, causal, seed;
};
static const TestCase TESTS[] = {
{128, 1, 1, 1, 8, 128, 0, 1},
{128, 1, 4, 4, 128, 128, 0, 2},
{128, 2, 4, 4, 256, 128, 0, 3},
{128, 1, 4, 1, 64, 64, 0, 4},
{128, 1, 8, 2, 64, 128, 0, 5},
{128, 2, 16, 4, 128, 128, 0, 6},
{64, 1, 4, 2, 32, 128, 0, 7},
{256, 1, 2, 1, 16, 128, 0, 8},
{32, 1, 4, 2, 32, 64, 0, 9},
{128, 3, 8, 2, 256, 128, 0, 10},
{128, 2, 32, 8, 512, 128, 0, 11},
{128, 1, 16, 2, 256, 128, 0, 12},
{128, 2, 32, 4, 512, 128, 0, 13},
{128, 2, 8, 2, 128, 128, 1, 14}, // causal
};
static int dispatch_test(const TestCase& tc) {
int r = 0;
dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.causal, tc.seed);
});
return r;
}
template <int HEAD_DIM>
static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
int max_pages = (kv_len + page_size - 1) / page_size;
int n_phys_pages = B * max_pages;
int max_splits = 32;
size_t sz_q = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
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);
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);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t* d_pt;
float *d_op, *d_ml;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
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);
bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
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);
free(h_pt);
PagedAttentionParams<bf16> pa;
pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
pa.use_mask = 0; pa.causal_offset = -1;
set_default_paged_strides(pa);
pa.scale = 1.0f / sqrtf((float)HEAD_DIM);
pa.page_size = page_size; pa.max_pages = max_pages;
pa.page_table = d_pt;
pa.k_cache = d_k_pool; pa.v_cache = d_v_pool;
pa.q = d_q; pa.mask = nullptr; pa.o = d_o;
pa.o_part = d_op; pa.ml_part = d_ml;
const int WARMUP = 10, ITERS = 100;
auto launch = [&]() {
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(pa); });
};
double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
char cfg[64];
snprintf(cfg, sizeof(cfg),
"B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
print_bench_row(cfg, r);
free(tmp);
cudaFree(d_q); cudaFree(d_o);
cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
cudaFree(d_op); cudaFree(d_ml);
}
static void bench() {
printf("\n===== PAGED DECODE BENCH =====\n");
print_bench_header();
bench_config<128>(1, 32, 4, 512, 128);
bench_config<128>(1, 32, 4, 1024, 128);
bench_config<128>(1, 32, 4, 2048, 128);
bench_config<128>(1, 32, 4, 4096, 128);
bench_config<128>(16, 32, 4, 2048, 128);
bench_config<128>(32, 32, 4, 1024, 128);
}
int main() {
int n = sizeof(TESTS) / sizeof(TESTS[0]);
int fail = 0;
printf("=== Paged Decode vs CPU reference (%d cases) ===\n\n", n);
for (int i = 0; i < n; i++) {
fail += dispatch_test(TESTS[i]);
if (fail) break;
}
if (fail) {
printf("\nFAILED (%d/%d tests failed)\n", fail, n);
return fail;
}
printf("\nAll %d tests passed!\n", n);
bench();
return 0;
}
csrc/tests/attn_paged_test.cu
+939
View File
@@ -0,0 +1,939 @@
// Compile:
// nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
// --extra-device-vectorization csrc/tests/attn_paged_test.cu \
// -o /tmp/test_paged && /tmp/test_paged
#include <cstring>
#include <vector>
#include "test_utils.cuh"
#include "../kernels/attn_dispatchers.cuh"
// ---- CPU reference: paged decode with variable seq_lens ----
// Q: [B, Hq, D], K/V pool: [pool_size, Hkv, D]
// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
// kv_indptr: [B+1]. mask: [B, max_seq_len] bool (True=keep) or NULL.
static void cpu_paged_decode_ref(
const float* Q, const float* K_pool, const float* V_pool,
const int64_t* req_to_token, const int64_t* req_pool_indices,
const int* kv_indptr, const bool* mask, int mask_b_stride,
int B, int Hq, int Hkv, int D, int max_ctx_len,
float* O)
{
float scale = 1.0f / sqrtf((float)D);
int n_rep = Hq / Hkv;
for (int b = 0; b < B; b++) {
int seq_len = kv_indptr[b + 1] - kv_indptr[b];
int64_t req_idx = req_pool_indices[b];
for (int h = 0; h < Hq; h++) {
int kv_h = h / n_rep;
float mv = -INFINITY, sv = 0.0f;
float accum[256] = {0.0f};
for (int kj = 0; kj < seq_len; kj++) {
if (mask && !mask[b * mask_b_stride + kj]) continue;
int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
float dot = 0.0f;
for (int d = 0; d < D; d++)
dot += Q[(b * Hq + h) * D + d] *
K_pool[slot * Hkv * D + kv_h * D + d];
dot *= scale;
float nm = fmaxf(mv, dot);
float a = expf(mv - nm);
float be = expf(dot - nm);
sv = sv * a + be;
for (int d = 0; d < D; d++)
accum[d] = accum[d] * a +
V_pool[slot * Hkv * D + kv_h * D + d] * be;
mv = nm;
}
float inv = 1.0f / sv;
for (int d = 0; d < D; d++)
O[(b * Hq + h) * D + d] = accum[d] * inv;
}
}
}
// ---- CPU reference: paged prefill with ragged batch ----
// Q: [total_q, Hq, D], K/V pool: [pool_size, Hkv, D]
// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
// kv_indptr: [B+1], qo_indptr: [B+1].
// mask: [B, max_q_len, max_seq_len] bool (True=keep, q-local + kv-local
// positions) or NULL. Used only when causal==0 to apply an arbitrary
// attention mask on top of the (unused) causal logic.
static void cpu_paged_prefill_ref(
const float* Q, const float* K_pool, const float* V_pool,
const int64_t* req_to_token, const int64_t* req_pool_indices,
const int* kv_indptr, const int* qo_indptr,
const bool* mask, int mask_q_stride, int mask_kv_stride,
int B, int Hq, int Hkv, int D, int max_ctx_len, int causal,
float* O)
{
float scale = 1.0f / sqrtf((float)D);
int n_rep = Hq / Hkv;
for (int b = 0; b < B; b++) {
int seq_len = kv_indptr[b + 1] - kv_indptr[b];
int q_len = qo_indptr[b + 1] - qo_indptr[b];
int causal_off = seq_len - q_len;
int64_t req_idx = req_pool_indices[b];
for (int h = 0; h < Hq; h++) {
int kv_h = h / n_rep;
for (int qi = 0; qi < q_len; qi++) {
float mv = -INFINITY, sv = 0.0f;
float accum[256] = {0.0f};
int lim = causal ? min(seq_len, causal_off + qi + 1) : seq_len;
for (int kj = 0; kj < lim; kj++) {
if (mask && !mask[b * mask_q_stride * mask_kv_stride
+ qi * mask_kv_stride + kj]) continue;
int64_t slot = req_to_token[req_idx * max_ctx_len + kj];
float dot = 0.0f;
for (int d = 0; d < D; d++)
dot += Q[(qo_indptr[b] + qi) * Hq * D + h * D + d] *
K_pool[slot * Hkv * D + kv_h * D + d];
dot *= scale;
float nm = fmaxf(mv, dot);
float a = expf(mv - nm);
float be = expf(dot - nm);
sv = sv * a + be;
for (int d = 0; d < D; d++)
accum[d] = accum[d] * a +
V_pool[slot * Hkv * D + kv_h * D + d] * be;
mv = nm;
}
float inv = 1.0f / sv;
for (int d = 0; d < D; d++)
O[(qo_indptr[b] + qi) * Hq * D + h * D + d] = accum[d] * inv;
}
}
}
}
// ======================================================================
// DECODE TEST
// ======================================================================
template <int HEAD_DIM>
static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
int causal, int seed) {
// Variable seq_lens per request
srand(seed);
std::vector<int> seq_lens(B);
for (int b = 0; b < B; b++)
seq_lens[b] = 8 + rand() % (max_seq - 8);
int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
int max_ctx = max_sl + 16;
int pool_size = B * max_ctx;
int num_reqs = B + 4;
printf("DECODE B=%d Hq=%d Hkv=%d D=%d seqs=[", B, Hq, Hkv, HEAD_DIM);
for (int b = 0; b < B; b++) printf("%d%s", seq_lens[b], b < B-1 ? "," : "");
printf("] causal=%d ... ", causal);
fflush(stdout);
size_t sz_q = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
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);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi;
float *d_op, *d_ml;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi);
cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
bf16* h_q = (bf16*)malloc(sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
bf16* h_k_pool = (bf16*)malloc(sz_kv);
bf16* h_v_pool = (bf16*)malloc(sz_kv);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
h_k_pool[i] = f2bf(rnd());
h_v_pool[i] = f2bf(rnd());
}
cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
// req_to_token: assign unique slots per request (scattered, not contiguous)
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
int next_slot = 0;
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++) {
h_rtt[r * max_ctx + p] = next_slot % pool_size;
next_slot++;
}
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
// req_pool_indices: pick B random request rows
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
for (int b = 0; b < B; b++) h_rpi[b] = b;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
// kv_indptr: prefix sum of seq_lens
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0;
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
// CPU reference
float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
h_k_f[i] = bf2f(h_k_pool[i]);
h_v_f[i] = bf2f(h_v_pool[i]);
}
float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
nullptr, 0,
B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
// Kernel launch
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = max_sl;
p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
p.mask = nullptr; p.mask_b_stride = 0;
p.mask_h_stride = 0; p.mask_q_stride = 0;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
cudaDeviceSynchronize();
bf16* h_o_bf = (bf16*)malloc(sz_q);
cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
const float atol = 0.02f, rtol = 0.02f;
bool pass = true;
float max_err = 0.0f;
for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
float e = fabsf(h_o_got[i] - h_o_ref[i]);
if (e > max_err) max_err = e;
if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
}
if (pass) printf("PASS (max_err=%.4e)\n", max_err);
else printf("FAIL (max_err=%.4e)\n", max_err);
free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
free(h_kvi); free(h_q_f); free(h_k_f); free(h_v_f);
free(h_o_ref); free(h_o_bf); free(h_o_got);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
return pass ? 0 : 1;
}
// ======================================================================
// DECODE WITH MASK TEST (regression: 2D mask on mixed seq_lens)
// ======================================================================
template <int HEAD_DIM>
static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
int seed) {
srand(seed);
std::vector<int> seq_lens(B);
for (int b = 0; b < B; b++)
seq_lens[b] = 8 + rand() % (max_seq - 8);
int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
int max_ctx = max_sl + 16;
int pool_size = B * max_ctx;
int num_reqs = B + 4;
printf("DECODE-MASK B=%d Hq=%d Hkv=%d D=%d max_sl=%d ... ", B, Hq, Hkv, HEAD_DIM, max_sl);
fflush(stdout);
size_t sz_q = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
size_t sz_mask = (size_t)B * max_sl * sizeof(bool);
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);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi;
bool *d_mask;
float *d_op, *d_ml;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi);
cudaMalloc(&d_mask, sz_mask);
cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
bf16* h_q = (bf16*)malloc(sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
bf16* h_k_pool = (bf16*)malloc(sz_kv);
bf16* h_v_pool = (bf16*)malloc(sz_kv);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
h_k_pool[i] = f2bf(rnd());
h_v_pool[i] = f2bf(rnd());
}
cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
int next_slot = 0;
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++) {
h_rtt[r * max_ctx + p] = next_slot % pool_size;
next_slot++;
}
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
for (int b = 0; b < B; b++) h_rpi[b] = b;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0;
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
// Mask: keep first half of each request's kv range, drop the rest —
// exercises the HasMask path with per-request seq_len.
bool* h_mask = (bool*)malloc(sz_mask);
for (int b = 0; b < B; b++)
for (int k = 0; k < max_sl; k++)
h_mask[b * max_sl + k] = (k < seq_lens[b]) && (k % 2 == 0);
cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
h_k_f[i] = bf2f(h_k_pool[i]);
h_v_f[i] = bf2f(h_v_pool[i]);
}
float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
h_mask, max_sl,
B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = max_sl;
p.causal_offset = -1; p.use_mask = 1;
p.mask = d_mask; p.mask_b_stride = max_sl;
p.mask_h_stride = 0; p.mask_q_stride = 0;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
cudaDeviceSynchronize();
bf16* h_o_bf = (bf16*)malloc(sz_q);
cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
const float atol = 0.02f, rtol = 0.02f;
bool pass = true;
float max_err = 0.0f;
for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
float e = fabsf(h_o_got[i] - h_o_ref[i]);
if (e > max_err) max_err = e;
if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
}
if (pass) printf("PASS (max_err=%.4e)\n", max_err);
else printf("FAIL (max_err=%.4e)\n", max_err);
free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
free(h_kvi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
free(h_o_ref); free(h_o_bf); free(h_o_got);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_mask);
cudaFree(d_op); cudaFree(d_ml);
return pass ? 0 : 1;
}
// ======================================================================
// PREFILL TEST
// ======================================================================
template <int HEAD_DIM>
static int run_prefill_test(int B, int Hq, int Hkv,
std::vector<int>& q_lens,
std::vector<int>& kv_lens,
int causal, int seed) {
int total_q = 0;
int max_sl = 0;
for (int b = 0; b < B; b++) {
total_q += q_lens[b];
max_sl = max(max_sl, kv_lens[b]);
}
int max_ctx = max_sl + 16;
int pool_size = B * max_ctx;
int num_reqs = B + 4;
printf("PREFILL B=%d Hq=%d Hkv=%d D=%d q_lens=[", B, Hq, Hkv, HEAD_DIM);
for (int b = 0; b < B; b++) printf("%d%s", q_lens[b], b < B-1 ? "," : "");
printf("] kv_lens=[");
for (int b = 0; b < B; b++) printf("%d%s", kv_lens[b], b < B-1 ? "," : "");
printf("] causal=%d ... ", causal);
fflush(stdout);
size_t sz_q = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi, *d_qoi;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
srand(seed);
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
bf16* h_q = (bf16*)malloc(sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
bf16* h_k_pool = (bf16*)malloc(sz_kv);
bf16* h_v_pool = (bf16*)malloc(sz_kv);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
h_k_pool[i] = f2bf(rnd());
h_v_pool[i] = f2bf(rnd());
}
cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
int next_slot = 0;
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++) {
h_rtt[r * max_ctx + p] = next_slot % pool_size;
next_slot++;
}
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
for (int b = 0; b < B; b++) h_rpi[b] = b;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0;
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_lens[b];
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
int* h_qoi = (int*)malloc(sz_qoi);
h_qoi[0] = 0;
for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_lens[b];
cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
// CPU reference
float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
h_k_f[i] = bf2f(h_k_pool[i]);
h_v_f[i] = bf2f(h_v_pool[i]);
}
float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
nullptr, 0, 0,
B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref);
// Kernel launch
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = max_sl;
int max_ql = 0;
for (int b = 0; b < B; b++) max_ql = max(max_ql, q_lens[b]);
p.max_q_len = max_ql;
p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
p.mask = nullptr; p.mask_b_stride = 0;
p.mask_h_stride = 0; p.mask_q_stride = 0;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
cudaDeviceSynchronize();
bf16* h_o_bf = (bf16*)malloc(sz_q);
cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
const float atol = 0.02f, rtol = 0.02f;
bool pass = true;
float max_err = 0.0f;
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
float e = fabsf(h_o_got[i] - h_o_ref[i]);
if (e > max_err) max_err = e;
if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
}
if (pass) printf("PASS (max_err=%.4e)\n", max_err);
else printf("FAIL (max_err=%.4e)\n", max_err);
free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
free(h_kvi); free(h_qoi); free(h_q_f); free(h_k_f); free(h_v_f);
free(h_o_ref); free(h_o_bf); free(h_o_got);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
return pass ? 0 : 1;
}
// ======================================================================
// PREFILL WITH MASK TEST (regression: 4D causal mask on single request)
// ======================================================================
template <int HEAD_DIM>
static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
srand(seed);
int B = 1;
int total_q = q_len;
int seq_len = q_len; // pure prefill: kv_len == q_len
int max_ctx = seq_len + 16;
int pool_size = B * max_ctx;
int num_reqs = B + 4;
printf("PREFILL-MASK Hq=%d Hkv=%d D=%d q_len=%d ... ", Hq, Hkv, HEAD_DIM, q_len);
fflush(stdout);
size_t sz_q = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
size_t sz_mask = (size_t)B * q_len * q_len * sizeof(bool);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi, *d_qoi;
bool *d_mask;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
cudaMalloc(&d_mask, sz_mask);
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
bf16* h_q = (bf16*)malloc(sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
bf16* h_k_pool = (bf16*)malloc(sz_kv);
bf16* h_v_pool = (bf16*)malloc(sz_kv);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
h_k_pool[i] = f2bf(rnd());
h_v_pool[i] = f2bf(rnd());
}
cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
int next_slot = 0;
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++) {
h_rtt[r * max_ctx + p] = next_slot % pool_size;
next_slot++;
}
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
h_rpi[0] = 0;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0; h_kvi[1] = seq_len;
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
int* h_qoi = (int*)malloc(sz_qoi);
h_qoi[0] = 0; h_qoi[1] = q_len;
cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
// 4D causal mask [B, 1, q_len, q_len], True=keep.
bool* h_mask = (bool*)malloc(sz_mask);
for (int qi = 0; qi < q_len; qi++)
for (int kj = 0; kj < q_len; kj++)
h_mask[qi * q_len + kj] = (kj <= qi);
cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
h_k_f[i] = bf2f(h_k_pool[i]);
h_v_f[i] = bf2f(h_v_pool[i]);
}
float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
// CPU ref with causal=0 so it consults the mask (not the causal flag).
cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
h_mask, q_len, q_len,
B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref);
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = q_len;
p.max_q_len = q_len;
p.causal_offset = -1; p.use_mask = 1;
p.mask = d_mask; p.mask_b_stride = q_len * q_len;
p.mask_h_stride = 0; p.mask_q_stride = q_len;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
cudaDeviceSynchronize();
bf16* h_o_bf = (bf16*)malloc(sz_q);
cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
const float atol = 0.02f, rtol = 0.02f;
bool pass = true;
float max_err = 0.0f;
for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
float e = fabsf(h_o_got[i] - h_o_ref[i]);
if (e > max_err) max_err = e;
if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
}
if (pass) printf("PASS (max_err=%.4e)\n", max_err);
else printf("FAIL (max_err=%.4e)\n", max_err);
free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
free(h_kvi); free(h_qoi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
free(h_o_ref); free(h_o_bf); free(h_o_got);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
cudaFree(d_mask);
return pass ? 0 : 1;
}
// ======================================================================
// BENCH
// ======================================================================
template <int HEAD_DIM>
static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
int max_ctx = seq_len + 16;
int pool_size = B * max_ctx;
int num_reqs = B;
size_t sz_q = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
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);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi;
float *d_op, *d_ml;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi);
cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++)
h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
for (int b = 0; b < B; b++) h_rpi[b] = b;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0;
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_len;
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = seq_len;
p.causal_offset = 0; p.use_mask = 0;
p.mask = nullptr; p.mask_b_stride = 0;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
p.o = d_o; p.o_part = d_op; p.ml_part = d_ml;
auto launch = [&]() {
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
};
// Decode: q_len=1, query is the last token → attends to all [0, seq_len).
// FLOPs = 2 * (QK^T + PV) = 4 * B * Hq * seq_len * D.
double flops = 4.0 * B * Hq * (double)seq_len * HEAD_DIM;
// HBM: K+V read (Q/O negligible for decode).
size_t nKV = (size_t)B * Hkv * seq_len * HEAD_DIM;
double bytes = 2.0 * nKV * sizeof(bf16);
BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
char cfg[64];
snprintf(cfg, sizeof(cfg), "DEC B=%2d Hq=%2d Hk=%d kv=%4d D=%3d",
B, Hq, Hkv, seq_len, HEAD_DIM);
print_bench_row(cfg, r);
free(tmp); free(h_rtt); free(h_rpi); free(h_kvi);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
}
template <int HEAD_DIM>
static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int causal) {
int total_q = B * q_len;
int max_ctx = kv_len + 16;
int pool_size = B * max_ctx;
int num_reqs = B;
size_t sz_q = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int64_t);
size_t sz_rpi = (size_t)B * sizeof(int64_t);
size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
int64_t *d_rtt, *d_rpi;
int *d_kvi, *d_qoi;
cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
int64_t* h_rtt = (int64_t*)malloc(sz_rtt);
for (int r = 0; r < num_reqs; r++)
for (int p = 0; p < max_ctx; p++)
h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
int64_t* h_rpi = (int64_t*)malloc(sz_rpi);
for (int b = 0; b < B; b++) h_rpi[b] = b;
cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
int* h_kvi = (int*)malloc(sz_kvi);
h_kvi[0] = 0;
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_len;
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
int* h_qoi = (int*)malloc(sz_qoi);
h_qoi[0] = 0;
for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len;
cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
PagedAttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
p.max_context_len = max_ctx; p.max_seq_len = kv_len;
p.total_q = total_q; p.max_q_len = q_len;
p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
p.mask = nullptr; p.mask_b_stride = 0;
p.scale = 1.0f / sqrtf((float)HEAD_DIM);
p.q = d_q; p.k_cache = d_k_pool; p.v_cache = d_v_pool;
p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
p.o = d_o; p.o_part = nullptr; p.ml_part = nullptr;
auto launch = [&]() {
dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_prefill<H>(p); });
};
// FLOPs = 2 * (QK^T + PV) = 4 * effective_qk_pairs * Hq * D.
// Non-causal: effective = q_len * kv_len.
// Causal: Q row qi attends to [0, causal_off + qi + 1) where
// causal_off = kv_len - q_len. Total KV accesses per request:
// sum_{qi=0}^{q_len-1} (kv_len - q_len + qi + 1)
// = q_len * (kv_len - q_len) + q_len * (q_len + 1) / 2.
double eff_kv;
if (causal) {
eff_kv = (double)q_len * (kv_len - q_len)
+ (double)q_len * (q_len + 1) / 2.0;
} else {
eff_kv = (double)q_len * kv_len;
}
double flops = 4.0 * B * Hq * eff_kv * HEAD_DIM;
// HBM: Q read + K read + V read + O write.
size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
size_t nQ = (size_t)total_q * Hq * HEAD_DIM;
double bytes = (2.0 * nQ + 2.0 * nKV) * sizeof(bf16);
BenchResult r = bench_kernel(launch, 10, 100, flops, bytes);
char cfg[80];
snprintf(cfg, sizeof(cfg), "PRE B=%d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d c=%d",
B, Hq, Hkv, q_len, kv_len, HEAD_DIM, causal);
print_bench_row(cfg, r);
free(tmp); free(h_rtt); free(h_rpi); free(h_kvi); free(h_qoi);
cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
}
int main() {
int fail = 0;
// ===== DECODE TESTS =====
printf("=== Paged Decode Tests ===\n\n");
fail += run_decode_test<128>(1, 32, 4, 512, 0, 1);
fail += run_decode_test<128>(1, 32, 4, 1024, 0, 2);
fail += run_decode_test<128>(4, 32, 4, 512, 0, 3);
fail += run_decode_test<128>(8, 32, 4, 1024, 0, 4);
fail += run_decode_test<128>(4, 32, 8, 2048, 0, 5);
fail += run_decode_test<128>(1, 16, 1, 256, 0, 6);
fail += run_decode_test<128>(2, 8, 2, 512, 1, 7);
fail += run_decode_test<64>(1, 4, 2, 256, 0, 8);
fail += run_decode_test<256>(1, 2, 1, 256, 0, 9);
fail += run_decode_test<128>(16, 32, 4, 2048, 0, 10);
fail += run_decode_test<128>(32, 32, 4, 1024, 0, 11);
// Decode with 2D mask (regression: mixed seq_lens + HasMask)
fail += run_decode_mask_test<128>(2, 8, 2, 256, 30);
fail += run_decode_mask_test<128>(4, 32, 4, 512, 31);
fail += run_decode_mask_test<64>(2, 4, 2, 128, 32);
if (fail) { printf("\nFAILED decode tests\n"); return fail; }
// ===== PREFILL TESTS =====
printf("\n=== Paged Prefill Tests ===\n\n");
// Single request, pure prefill (q_len == kv_len)
{
std::vector<int> ql = {512};
std::vector<int> kl = {512};
fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 20);
}
{
std::vector<int> ql = {1024};
std::vector<int> kl = {1024};
fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 21);
}
{
std::vector<int> ql = {2048};
std::vector<int> kl = {2048};
fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 22);
}
// Ragged batch: different q_lens and kv_lens
{
std::vector<int> ql = {128, 256, 64};
std::vector<int> kl = {128, 256, 64};
fail += run_prefill_test<128>(3, 32, 4, ql, kl, 1, 23);
}
{
std::vector<int> ql = {64, 128, 256, 32};
std::vector<int> kl = {64, 128, 256, 32};
fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 24);
}
// Extend: kv_len > q_len (append to existing cache)
{
std::vector<int> ql = {64, 128};
std::vector<int> kl = {256, 512};
fail += run_prefill_test<128>(2, 32, 4, ql, kl, 1, 25);
}
// Non-causal
{
std::vector<int> ql = {256, 128};
std::vector<int> kl = {256, 128};
fail += run_prefill_test<128>(2, 32, 4, ql, kl, 0, 26);
}
// Single token (q_len=1 per request, like decode but via prefill path)
{
std::vector<int> ql = {1, 1, 1, 1};
std::vector<int> kl = {128, 256, 64, 512};
fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 27);
}
// D=64
{
std::vector<int> ql = {128, 64};
std::vector<int> kl = {128, 64};
fail += run_prefill_test<64>(2, 4, 2, ql, kl, 1, 28);
}
// D=256
{
std::vector<int> ql = {128, 64};
std::vector<int> kl = {128, 64};
fail += run_prefill_test<256>(2, 2, 1, ql, kl, 1, 29);
}
// Prefill with 4D causal mask (regression: single-request mask path)
fail += run_prefill_mask_test<128>(32, 4, 512, 40);
fail += run_prefill_mask_test<128>(32, 4, 1024, 41);
fail += run_prefill_mask_test<64>(4, 2, 256, 42);
if (fail) { printf("\nFAILED prefill tests\n"); return fail; }
printf("\nAll tests passed!\n");
// ===== BENCH =====
printf("\n===== PAGED DECODE BENCH =====\n");
print_bench_header();
bench_decode<128>(1, 32, 4, 512);
bench_decode<128>(1, 32, 4, 1024);
bench_decode<128>(1, 32, 4, 2048);
bench_decode<128>(1, 32, 4, 4096);
bench_decode<128>(4, 32, 4, 2048);
bench_decode<128>(16, 32, 4, 2048);
bench_decode<128>(32, 32, 4, 1024);
printf("\n===== PAGED PREFILL BENCH =====\n");
print_bench_header();
bench_prefill<128>(1, 32, 4, 512, 512, 0);
bench_prefill<128>(1, 32, 4, 1024, 1024, 0);
bench_prefill<128>(1, 32, 4, 2048, 2048, 0);
bench_prefill<128>(1, 32, 4, 2048, 2048, 1);
bench_prefill<128>(4, 32, 4, 2048, 2048, 1);
bench_prefill<128>(1, 32, 4, 4096, 4096, 1);
return 0;
}
scripts/tools/benchmark.py
+57 -9
View File
@@ -12,9 +12,14 @@ from astrai.model import AutoModel
_DTYPES = ["bfloat16", "float16", "float32"]
_CACHES = ["contiguous", "paged"]
DEFAULT_CKPT = str(Path(__file__).resolve().parents[2] / "ckpt_bucket" / "kami-15bt")
_BACKENDS = ["cuda", "torch_native"]
CACHE_MAX_SEQ = 2048
_BACKEND_MAP = {
"cuda": ATTN_BACKEND.CUDA,
"torch_native": ATTN_BACKEND.TORCH_NATIVE,
}
class BenchmarkResult:
def __init__(
@@ -42,12 +47,14 @@ class GenerationBenchmark:
device: str = "cuda",
dtype: torch.dtype = torch.bfloat16,
cache_type: str = "contiguous",
backend: ATTN_BACKEND = ATTN_BACKEND.CUDA,
):
self.device = device
self.dtype = dtype
self.cache_type = cache_type
self.model = model
self.config = config
self.backend = backend
def _make_pool(self, batch_size: int) -> PagePool:
return PagePool(
@@ -82,7 +89,7 @@ class GenerationBenchmark:
kv_cache = pool.bind_tasks(
task_ids, [prompt_len] * batch_size, self.device, start_pos=0
)
with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode(), attn_backend(self.backend):
self.model(
input_ids,
input_mask=input_mask,
@@ -105,7 +112,7 @@ class GenerationBenchmark:
total_len, device=self.device
)
kv_cache = pool.bind_tasks(task_ids, [seq_len + 1] * batch_size, self.device)
with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode(), attn_backend(self.backend):
self.model(
input_ids,
input_mask=input_mask,
@@ -121,6 +128,11 @@ class GenerationBenchmark:
) -> BenchmarkResult:
import time
pool = self._make_pool(batch_size)
task_ids = [f"bench_prefill_{i}" for i in range(batch_size)]
for tid in task_ids:
pool.task_alloc(tid, list(range(prompt_length)))
input_ids = torch.randint(
0, self.config.vocab_size, (batch_size, prompt_length), device=self.device
)
@@ -129,16 +141,32 @@ class GenerationBenchmark:
.unsqueeze(0)
.expand(batch_size, -1)
)
input_mask = position_ids.unsqueeze(-1) >= torch.arange(
prompt_length, device=self.device
)
kv_cache = pool.bind_tasks(
task_ids, [prompt_length] * batch_size, self.device, start_pos=0
)
for _ in range(3):
with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
self.model(input_ids, position_ids=position_ids)
with torch.inference_mode(), attn_backend(self.backend):
self.model(
input_ids,
input_mask=input_mask,
kv_cache=kv_cache,
position_ids=position_ids,
)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(num_trials):
with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
self.model(input_ids, position_ids=position_ids)
with torch.inference_mode(), attn_backend(self.backend):
self.model(
input_ids,
input_mask=input_mask,
kv_cache=kv_cache,
position_ids=position_ids,
)
torch.cuda.synchronize()
elapsed = time.perf_counter() - t0
tokens = batch_size * prompt_length * num_trials
@@ -208,6 +236,17 @@ def print_benchmark_result(result: BenchmarkResult) -> None:
@click.option(
"--cache", type=click.Choice(_CACHES), default="contiguous", help="KV cache type."
)
@click.option(
"--backend",
type=click.Choice(_BACKENDS),
default="cuda",
help="Attention backend.",
)
@click.option(
"--compare",
is_flag=True,
help="Run both backends and print side-by-side speed comparison.",
)
@click.option("--batch_size", type=int, default=4, help="Batch size.")
@click.option("--prompt_length", type=int, default=512, help="Prompt length.")
@click.option("--gen_length", type=int, default=128, help="Generation length.")
@@ -216,13 +255,16 @@ def print_benchmark_result(result: BenchmarkResult) -> None:
@click.option("--decode_only", is_flag=True, help="Decode benchmark only.")
@click.option(
"--ckpt",
default=DEFAULT_CKPT,
required=True,
type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
help="Checkpoint directory.",
)
def benchmark_command(
device: str,
dtype: str,
cache: str,
backend: str,
compare: bool,
batch_size: int,
prompt_length: int,
gen_length: int,
@@ -244,15 +286,21 @@ def benchmark_command(
model.to(device=device, dtype=dtype_map[dtype])
model.eval()
backends = _BACKENDS if compare else [backend]
for name in backends:
bench = GenerationBenchmark(
model=model,
config=config,
device=device,
dtype=dtype_map[dtype],
cache_type=cache,
backend=_BACKEND_MAP[name],
)
click.secho(f"Benchmark: device={device} dtype={dtype}", bold=True)
click.secho(
f"Benchmark: device={device} dtype={dtype} backend={name}", bold=True
)
if not decode_only:
result = bench.run_prefill_benchmark(
setup.py
+31 -1
View File
@@ -1,4 +1,5 @@
import os
import shutil
import subprocess
import sys
import warnings
@@ -89,14 +90,30 @@ if _should_build():
super().build_extensions()
return
# Each subprocess gets its own build-temp / build-lib so the
# ninja files (build.ninja, .ninja_log) never race. The built
# .so files are then collected into the parent's build_lib so the
# normal setuptools copy steps (inplace / editable wheel) work.
names = [e.name for e in self.extensions]
env = {**os.environ, "BUILD_PARALLEL": "1"}
base = os.path.join("build", "parallel")
os.makedirs(base, exist_ok=True)
procs = {}
for i in range(0, len(names), max_workers):
batch = names[i : i + max_workers]
for name in batch:
e = {**env, "ASTRAI_BUILD_SINGLE_EXT": name}
cmd = [sys.executable, __file__, "build_ext", "--inplace"]
tag = name.replace(".", "_")
subdir = os.path.join(base, tag)
cmd = [
sys.executable,
__file__,
"build_ext",
"--build-temp",
os.path.join(subdir, "temp"),
"--build-lib",
os.path.join(subdir, "lib"),
]
procs[name] = subprocess.Popen(
cmd, env=e, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
)
@@ -108,6 +125,19 @@ if _should_build():
f"parallel build failed for {name} "
f"(exit {procs[name].returncode})"
)
self._collect_extensions(
os.path.join(base, name.replace(".", "_"), "lib")
)
def _collect_extensions(self, sub_lib):
src = os.path.join(sub_lib, "astrai", "extension", "lib")
if not os.path.isdir(src):
return
dst = os.path.join(self.build_lib, "astrai", "extension", "lib")
os.makedirs(dst, exist_ok=True)
for f in os.listdir(src):
if f.endswith(".so"):
shutil.copy2(os.path.join(src, f), os.path.join(dst, f))
cmdclass["build_ext"] = ParallelBuildExtension
tests/extension/test_backend_equivalence.py
+12 -4
View File
@@ -13,18 +13,26 @@ from tests.extension.conftest import D, skip_no_kernel
@skip_no_kernel
def test_training_forward_matches_torch(cuda_model):
"""Training forward (kv_cache=None) should produce identical logits."""
"""Training forward (kv_cache=None) should produce identical logits.
CudaBackend is inference-only: it raises when kv_cache is None. Training
must use TorchNativeBackend (the default). Verify the torch path is
stable and that CudaBackend rejects the training path explicitly.
"""
import pytest
model, _ = cuda_model
input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
with torch.no_grad():
out_torch = model(input_ids)
with pytest.raises(RuntimeError, match="does not support training"):
with attn_backend(ATTN_BACKEND.CUDA):
with torch.no_grad():
out_cuda = model(input_ids)
model(input_ids)
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff == 0.0, f"Training forward diff {diff} should be 0"
assert out_torch["logits"].shape[0] == 2
@skip_no_kernel