feat: wire up paged decode CUDA kernel to Python extension
- Add attn_paged_decode wrapper in ops.py with gather fallback - Register kernel in loader.py and export from __init__.py - Extract test_utils.cuh shared by all attention unit tests - Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu - Refactor decode/prefill tests to use common bf16 helpers and cpu ref - Fix k_cache dim check in attn_paged_decode.cu
This commit is contained in:
parent
89ece26c25
commit
2c3cef1c87
|
|
@ -3,16 +3,18 @@
|
|||
Public API:
|
||||
- ``attn_decode`` — single-query decode attention
|
||||
- ``attn_prefill`` — multi-query prefill attention
|
||||
- ``attn_paged_decode`` — paged decode attention (direct page-table access)
|
||||
|
||||
Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
|
||||
when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
|
||||
"""
|
||||
|
||||
from astrai.extension.loader import KERNEL_NAMES, is_available
|
||||
from astrai.extension.ops import attn_decode, attn_prefill
|
||||
from astrai.extension.ops import attn_decode, attn_paged_decode, attn_prefill
|
||||
|
||||
__all__ = [
|
||||
"attn_decode",
|
||||
"attn_paged_decode",
|
||||
"attn_prefill",
|
||||
"is_available",
|
||||
"KERNEL_NAMES",
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import logging
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
KERNEL_NAMES = ["attn_decode", "attn_prefill"]
|
||||
KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode"]
|
||||
|
||||
_available: dict[str, bool] = {}
|
||||
_modules: dict[str, object] = {}
|
||||
|
|
|
|||
|
|
@ -42,6 +42,40 @@ def _torch_fallback(
|
|||
)
|
||||
|
||||
|
||||
def _gather_kv_from_pages(
|
||||
page_table: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
page_size: int,
|
||||
kv_len: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Gather contiguous K/V from paged cache for torch SDPA fallback.
|
||||
|
||||
Shapes:
|
||||
page_table : [batch, max_pages] (int64)
|
||||
k_cache : [n_pages, page_size, n_kv_heads, head_dim]
|
||||
v_cache : same as k_cache
|
||||
Returns:
|
||||
k, v : [batch, n_kv_heads, kv_len, head_dim]
|
||||
"""
|
||||
batch, max_pages = page_table.shape
|
||||
n_pages, ps, n_kv_heads, head_dim = k_cache.shape
|
||||
if ps != page_size:
|
||||
raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}")
|
||||
|
||||
k = k_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
|
||||
v = v_cache.new_empty(batch, n_kv_heads, kv_len, head_dim)
|
||||
|
||||
for b in range(batch):
|
||||
for pos in range(kv_len):
|
||||
log_pg = pos // page_size
|
||||
pg_off = pos % page_size
|
||||
phys = int(page_table[b, log_pg].item())
|
||||
k[b, :, pos, :] = k_cache[phys, pg_off, :, :]
|
||||
v[b, :, pos, :] = v_cache[phys, pg_off, :, :]
|
||||
return k, v
|
||||
|
||||
|
||||
def attn_decode(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
|
|
@ -84,3 +118,32 @@ def attn_prefill(
|
|||
scale=scale,
|
||||
)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
||||
|
||||
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,
|
||||
mask: torch.Tensor | None = None,
|
||||
is_causal: bool = False,
|
||||
causal_offset: int = 0,
|
||||
scale: float | None = None,
|
||||
) -> torch.Tensor:
|
||||
if _available["attn_paged_decode"]:
|
||||
return _modules["attn_paged_decode"].attn_paged_decode(
|
||||
q,
|
||||
page_table,
|
||||
k_cache,
|
||||
v_cache,
|
||||
page_size,
|
||||
kv_len,
|
||||
mask=mask,
|
||||
is_causal=is_causal,
|
||||
causal_offset=causal_offset,
|
||||
scale=scale,
|
||||
)
|
||||
k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
|
||||
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ torch::Tensor attn_paged_decode(
|
|||
int batch = q.size(0);
|
||||
int q_head = q.size(1);
|
||||
int head_dim = q.size(3);
|
||||
int kv_head = k_cache.size(3);
|
||||
int kv_head = k_cache.size(2);
|
||||
int max_pages = page_table.size(1);
|
||||
|
||||
TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
|
||||
|
|
|
|||
|
|
@ -5,21 +5,12 @@ nvcc -I csrc -arch=sm_89 -O3 \
|
|||
csrc/tests/attn_decode_test.cu -o test && ./test
|
||||
*/
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <sys/time.h>
|
||||
#include "test_utils.cuh"
|
||||
#include "../kernels/attn_decode_split_kv.cuh"
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "../kernels/attn_decode_split_kv_mma.cuh"
|
||||
#endif
|
||||
|
||||
static double now_ms() {
|
||||
struct timeval tv;
|
||||
gettimeofday(&tv, NULL);
|
||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
||||
}
|
||||
|
||||
// Split-K scratch (torch-free): the production launcher allocates these from
|
||||
// torch; here we pass pre-allocated device buffers so the bench loop doesn't
|
||||
// pay a cudaMalloc per iteration. Size for the maximum split count (32).
|
||||
|
|
@ -28,13 +19,6 @@ struct DecodeScratch {
|
|||
float* ml_part = nullptr;
|
||||
};
|
||||
|
||||
static int decode_num_splits(int base_blocks, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
|
||||
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||
}
|
||||
|
||||
// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
|
||||
// scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu.
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
|
|
@ -46,7 +30,7 @@ static bool decode_use_mma(const AttentionParams<bf16>& p) {
|
|||
template <int HEAD_DIM, int BC>
|
||||
static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
|
||||
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||
p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
p.o_part = sc.o_part;
|
||||
p.ml_part = sc.ml_part;
|
||||
|
||||
|
|
@ -59,7 +43,7 @@ static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
|
|||
static void launch_scalar_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
|
||||
int gs = p.q_head / p.kv_head;
|
||||
int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
|
||||
p.num_splits = decode_num_splits(p.batch * p.kv_head, chunks_total);
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
|
||||
p.o_part = sc.o_part;
|
||||
p.ml_part = sc.ml_part;
|
||||
|
||||
|
|
@ -86,43 +70,6 @@ static void dispatch_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
|
|||
}
|
||||
}
|
||||
|
||||
static void cpu_decode(const float* Q, const float* K, const float* V,
|
||||
const bool* mask, float* O,
|
||||
int B, int Hq, int Hk, int seq_len, int D) {
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
int n_rep = Hq / Hk;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int h = 0; h < Hq; h++) {
|
||||
int kv_h = h / n_rep;
|
||||
float mv = -INFINITY, sv = 0.0f;
|
||||
float accum[256] = {0};
|
||||
for (int s = 0; s < seq_len; s++) {
|
||||
if (!mask[b * seq_len + s]) continue;
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < D; d++)
|
||||
dot += Q[((b * Hq + h) * 1 + 0) * D + d]
|
||||
* K[((b * Hk + kv_h) * seq_len + s) * D + d];
|
||||
dot *= scale;
|
||||
float nm = fmaxf(mv, dot);
|
||||
float al = expf(mv - nm);
|
||||
float be = expf(dot - nm);
|
||||
sv = sv * al + be;
|
||||
for (int d = 0; d < D; d++)
|
||||
accum[d] = accum[d] * al
|
||||
+ V[((b * Hk + kv_h) * seq_len + s) * D + d] * be;
|
||||
mv = nm;
|
||||
}
|
||||
float inv = 1.0f / sv;
|
||||
for (int d = 0; d < D; d++)
|
||||
O[((b * Hq + h) * 1 + 0) * D + d] = accum[d] * inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static bf16 f2bf(float x) { return __float2bfloat16(x); }
|
||||
static float bf2f(bf16 x) { return __bfloat162float(x); }
|
||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
||||
|
||||
// Warmed-up, CUDA-event timed sweep over the production decode MMA path.
|
||||
// Decode (q_len==1) is memory-bound: the two matmuls are GEMV-shaped, so we
|
||||
// report both effective K/V read bandwidth and the (small) attention FLOP/s.
|
||||
|
|
@ -259,7 +206,7 @@ int main() {
|
|||
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
|
||||
|
||||
float* ref=new float[nQ];
|
||||
cpu_decode(hQ,hK,hV,hMask,ref,B,Hq,Hk,sl,D);
|
||||
cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 0);
|
||||
|
||||
float max_err=0;
|
||||
for (size_t i=0;i<nQ;i++){
|
||||
|
|
|
|||
|
|
@ -0,0 +1,346 @@
|
|||
// 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_paged_decode_split_kv.cuh"
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "../kernels/attn_paged_decode_split_kv_mma.cuh"
|
||||
#endif
|
||||
|
||||
// Copy contiguous K/V from page pool (reference gather)
|
||||
static void gather_kv_cpu(
|
||||
const bf16* h_k_pool, const bf16* h_v_pool,
|
||||
const int64_t* h_pt, int B, int Hkv, int kv_len,
|
||||
int page_size, int head_dim,
|
||||
bf16* h_k, bf16* h_v)
|
||||
{
|
||||
int max_pages = (kv_len + page_size - 1) / page_size;
|
||||
size_t page_stride = (size_t)page_size * Hkv * head_dim;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int pos = 0; pos < kv_len; pos++) {
|
||||
int log_pg = pos / page_size;
|
||||
int pg_off = pos % page_size;
|
||||
int phys = (int)h_pt[b * max_pages + log_pg];
|
||||
for (int h = 0; h < Hkv; h++) {
|
||||
size_t src_base = (size_t)phys * page_stride
|
||||
+ (size_t)pg_off * Hkv * head_dim
|
||||
+ h * head_dim;
|
||||
size_t dst_base = ((size_t)b * 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 void launch_paged_decode(PagedAttentionParams<bf16, float>& p) {
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
int G_check = p.q_head / p.kv_head;
|
||||
bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
|
||||
if (use_mma) {
|
||||
int tiles_total = (p.kv_len + 32 - 1) / 32;
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
int group_sz = p.q_head / p.kv_head;
|
||||
int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
|
||||
size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
paged_attn_decode_split_kv_kernel<<<
|
||||
dim3(p.batch * p.kv_head, 1, p.num_splits),
|
||||
dim3(32, group_sz), smem>>>(p);
|
||||
}
|
||||
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
|
||||
}
|
||||
|
||||
template <int HEAD_DIM>
|
||||
static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) {
|
||||
printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM);
|
||||
fflush(stdout);
|
||||
|
||||
int max_pages = (kv_len + page_size - 1) / page_size;
|
||||
int n_phys_pages = B * max_pages;
|
||||
|
||||
size_t sz_q = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
|
||||
size_t sz_o = sz_q;
|
||||
size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
|
||||
size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
|
||||
int max_splits = 32;
|
||||
size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
|
||||
size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
|
||||
|
||||
bf16 *d_q, *d_o_paged, *d_o_ref;
|
||||
bf16 *d_k_pool, *d_v_pool;
|
||||
int64_t* d_pt;
|
||||
float *d_op, *d_ml;
|
||||
|
||||
cudaMalloc(&d_q, sz_q);
|
||||
cudaMalloc(&d_o_paged, sz_o);
|
||||
cudaMalloc(&d_o_ref, sz_o);
|
||||
cudaMalloc(&d_k_pool, sz_kv);
|
||||
cudaMalloc(&d_v_pool, sz_kv);
|
||||
cudaMalloc(&d_pt, sz_pt);
|
||||
cudaMalloc(&d_op, sz_op);
|
||||
cudaMalloc(&d_ml, sz_ml);
|
||||
|
||||
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, 0, 0);
|
||||
|
||||
float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
|
||||
PagedAttentionParams<bf16, float> p;
|
||||
p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
|
||||
p.kv_len = kv_len; p.head_dim = HEAD_DIM;
|
||||
p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
|
||||
p.num_splits = 1; p.scale = scale_val;
|
||||
p.page_size = page_size; p.max_pages = max_pages;
|
||||
p.page_table = d_pt;
|
||||
p.k_cache = d_k_pool; p.v_cache = d_v_pool;
|
||||
p.q = d_q; p.mask = nullptr; p.o = d_o_paged;
|
||||
p.o_part = d_op; p.ml_part = d_ml;
|
||||
|
||||
launch_paged_decode<HEAD_DIM>(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_err = 0.0f;
|
||||
int bad_idx = -1;
|
||||
for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
|
||||
float e = fabsf(h_o_paged[i] - h_o_ref[i]);
|
||||
if (e > max_err) { max_err = e; bad_idx = i; }
|
||||
}
|
||||
|
||||
bool pass = max_err < 0.02f;
|
||||
|
||||
if (pass) {
|
||||
printf("PASS (max_abs_err=%.4e)\n", max_err);
|
||||
} else {
|
||||
int b = bad_idx / (Hq * HEAD_DIM);
|
||||
int h = (bad_idx / HEAD_DIM) % Hq;
|
||||
int d = bad_idx % HEAD_DIM;
|
||||
printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
|
||||
max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
|
||||
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_o_ref);
|
||||
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, seed;
|
||||
};
|
||||
|
||||
static const TestCase TESTS[] = {
|
||||
{128, 1, 1, 1, 8, 128, 1},
|
||||
{128, 1, 4, 4, 128, 128, 2},
|
||||
{128, 2, 4, 4, 256, 128, 3},
|
||||
{128, 1, 4, 1, 64, 64, 4},
|
||||
{128, 1, 8, 2, 64, 128, 5},
|
||||
{128, 2, 16, 4, 128, 128, 6},
|
||||
{64, 1, 4, 2, 32, 128, 7},
|
||||
{256, 1, 2, 1, 16, 128, 8},
|
||||
{32, 1, 4, 2, 32, 64, 9},
|
||||
{128, 3, 8, 2, 256, 128, 10},
|
||||
{128, 2, 32, 8, 512, 128, 11},
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
{128, 1, 16, 2, 256, 128, 12},
|
||||
{128, 2, 32, 4, 512, 128, 13},
|
||||
#endif
|
||||
};
|
||||
|
||||
static int dispatch_test(const TestCase& tc) {
|
||||
switch (tc.head_dim) {
|
||||
case 32: return run_test<32>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
|
||||
case 64: return run_test<64>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
|
||||
case 128: return run_test<128>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
|
||||
case 256: return run_test<256>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
|
||||
default: return 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Warmed-up, CUDA-event timed sweep over paged decode configs.
|
||||
// Reports per-call latency and effective K/V read bandwidth.
|
||||
// Bytes = K + V read through page table (B*Hk*kv*D each), bf16.
|
||||
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;
|
||||
|
||||
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);
|
||||
int max_splits = 32;
|
||||
size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
|
||||
size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
|
||||
|
||||
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);
|
||||
|
||||
float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
|
||||
PagedAttentionParams<bf16, float> 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.is_causal = 0; pa.causal_offset = 0;
|
||||
pa.num_splits = 1; pa.scale = scale_val;
|
||||
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;
|
||||
for (int i = 0; i < WARMUP; i++) launch_paged_decode<HEAD_DIM>(pa);
|
||||
cudaDeviceSynchronize();
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
cudaEvent_t s, e;
|
||||
cudaEventCreate(&s); cudaEventCreate(&e);
|
||||
cudaEventRecord(s);
|
||||
for (int i = 0; i < ITERS; i++) launch_paged_decode<HEAD_DIM>(pa);
|
||||
cudaEventRecord(e); cudaEventSynchronize(e);
|
||||
float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= ITERS;
|
||||
|
||||
double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
|
||||
double tflops = flops / (ms * 1e-3) / 1e12;
|
||||
size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
|
||||
double bytes = 2.0 * (2.0 * nKV);
|
||||
double gbps = bytes / (ms * 1e-3) / 1e9;
|
||||
|
||||
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);
|
||||
printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
|
||||
cfg, ms, gbps, tflops);
|
||||
|
||||
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);
|
||||
cudaEventDestroy(s); cudaEventDestroy(e);
|
||||
}
|
||||
|
||||
static void bench() {
|
||||
printf("\n===== PAGED DECODE BENCH =====\n");
|
||||
printf("%-46s | %10s | %10s | %10s\n",
|
||||
"config", "latency", "bandwidth", "throughput");
|
||||
printf("---------------------------------------------------------------"
|
||||
"----------------------------\n");
|
||||
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;
|
||||
}
|
||||
|
|
@ -1,279 +0,0 @@
|
|||
// Compile:
|
||||
// nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
|
||||
// --extra-device-vectorization csrc/tests/attn_paged_vs_contiguous.cu \
|
||||
// -o /tmp/test_pv && /tmp/test_pv
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cassert>
|
||||
#include "../kernels/attn_paged_decode_split_kv.cuh"
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "../kernels/attn_paged_decode_split_kv_mma.cuh"
|
||||
#endif
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
static int num_splits(int base_blocks, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
|
||||
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||
}
|
||||
|
||||
// Copy contiguous K/V from page pool (reference gather)
|
||||
static void gather_kv_cpu(
|
||||
const bf16* h_k_pool, const bf16* h_v_pool,
|
||||
const int64_t* h_pt, int B, int Hkv, int kv_len,
|
||||
int page_size, int head_dim,
|
||||
bf16* h_k, bf16* h_v)
|
||||
{
|
||||
int max_pages = (kv_len + page_size - 1) / page_size;
|
||||
size_t page_stride = (size_t)page_size * Hkv * head_dim;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int pos = 0; pos < kv_len; pos++) {
|
||||
int log_pg = pos / page_size;
|
||||
int pg_off = pos % page_size;
|
||||
int phys = (int)h_pt[b * max_pages + log_pg];
|
||||
for (int h = 0; h < Hkv; h++) {
|
||||
size_t src_base = (size_t)phys * page_stride
|
||||
+ (size_t)pg_off * Hkv * head_dim
|
||||
+ h * head_dim;
|
||||
size_t dst_base = ((size_t)b * kv_len + pos) * Hkv * head_dim + h * head_dim;
|
||||
memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16));
|
||||
memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int HEAD_DIM>
|
||||
static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) {
|
||||
printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM);
|
||||
fflush(stdout);
|
||||
|
||||
int G = Hq / Hkv;
|
||||
int max_pages = (kv_len + page_size - 1) / page_size;
|
||||
int n_phys_pages = B * max_pages;
|
||||
|
||||
// ---- allocate ----
|
||||
bf16 *d_q, *d_o_paged, *d_o_ref;
|
||||
bf16 *d_k_pool, *d_v_pool;
|
||||
int64_t* d_pt;
|
||||
float *d_op, *d_ml;
|
||||
|
||||
size_t sz_q = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
|
||||
size_t sz_o = sz_q;
|
||||
size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
|
||||
size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
|
||||
int max_splits = 32;
|
||||
size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
|
||||
size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
|
||||
|
||||
cudaMalloc(&d_q, sz_q);
|
||||
cudaMalloc(&d_o_paged, sz_o);
|
||||
cudaMalloc(&d_o_ref, sz_o);
|
||||
cudaMalloc(&d_k_pool, sz_kv);
|
||||
cudaMalloc(&d_v_pool, sz_kv);
|
||||
cudaMalloc(&d_pt, sz_pt);
|
||||
cudaMalloc(&d_op, sz_op);
|
||||
cudaMalloc(&d_ml, sz_ml);
|
||||
|
||||
// ---- init: deterministic random using seed ----
|
||||
srand(seed);
|
||||
auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
|
||||
|
||||
// Q
|
||||
bf16* h_q = (bf16*)malloc(sz_q);
|
||||
for (int i = 0; i < B * Hq * HEAD_DIM; i++)
|
||||
h_q[i] = __float2bfloat16(rnd());
|
||||
cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
|
||||
|
||||
// Page pool K/V
|
||||
bf16* h_k_pool = (bf16*)malloc(sz_kv);
|
||||
bf16* h_v_pool = (bf16*)malloc(sz_kv);
|
||||
size_t ps = (size_t)page_size * Hkv * HEAD_DIM;
|
||||
for (int pg = 0; pg < n_phys_pages; pg++) {
|
||||
for (int off = 0; off < page_size; off++) {
|
||||
for (int h = 0; h < Hkv; h++) {
|
||||
for (int d = 0; d < HEAD_DIM; d++) {
|
||||
float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d));
|
||||
size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM + h * HEAD_DIM + d;
|
||||
h_k_pool[idx] = __float2bfloat16(v);
|
||||
h_v_pool[idx] = __float2bfloat16(v * 0.3f);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
|
||||
|
||||
// Page table
|
||||
int64_t* h_pt = (int64_t*)malloc(sz_pt);
|
||||
int next_pg = 0;
|
||||
for (int b = 0; b < B; b++)
|
||||
for (int p = 0; p < max_pages; p++)
|
||||
h_pt[b * max_pages + p] = next_pg++;
|
||||
cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
|
||||
|
||||
// ---- reference: gather contiguous K/V, then run CPU online-softmax ----
|
||||
bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
|
||||
bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
|
||||
gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont);
|
||||
|
||||
float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
|
||||
float sscale = 1.0f / sqrtf((float)HEAD_DIM);
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int hq = 0; hq < Hq; hq++) {
|
||||
int hkv = hq / G;
|
||||
size_t q_base = (size_t)b * Hq * HEAD_DIM + (size_t)hq * HEAD_DIM;
|
||||
size_t kv_base = ((size_t)b * kv_len) * Hkv * HEAD_DIM + (size_t)hkv * HEAD_DIM;
|
||||
|
||||
float m = -1e30f, d = 0.0f;
|
||||
float acc[256] = {0.0f};
|
||||
for (int pos = 0; pos < kv_len; pos++) {
|
||||
float s = 0.0f;
|
||||
for (int dim = 0; dim < HEAD_DIM; dim++)
|
||||
s += __bfloat162float(h_q[q_base + dim])
|
||||
* __bfloat162float(h_k_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]);
|
||||
s *= sscale;
|
||||
|
||||
float nm = fmaxf(m, s);
|
||||
float a = expf(m - nm);
|
||||
float b = expf(s - nm);
|
||||
d = d * a + b;
|
||||
for (int dim = 0; dim < HEAD_DIM; dim++)
|
||||
acc[dim] = acc[dim] * a + __bfloat162float(h_v_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]) * b;
|
||||
m = nm;
|
||||
}
|
||||
for (int dim = 0; dim < HEAD_DIM; dim++)
|
||||
h_o_ref[b * Hq * HEAD_DIM + hq * HEAD_DIM + dim] = acc[dim] / d;
|
||||
}
|
||||
}
|
||||
|
||||
// ---- paged decode kernel ----
|
||||
float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
|
||||
PagedAttentionParams<bf16, float> p;
|
||||
p.batch = B;
|
||||
p.q_head = Hq;
|
||||
p.kv_head = Hkv;
|
||||
p.q_len = 1;
|
||||
p.kv_len = kv_len;
|
||||
p.head_dim = HEAD_DIM;
|
||||
p.use_mask = 0;
|
||||
p.is_causal = 0;
|
||||
p.causal_offset = 0;
|
||||
p.num_splits = 1;
|
||||
p.scale = scale_val;
|
||||
p.page_size = page_size;
|
||||
p.max_pages = max_pages;
|
||||
p.page_table = d_pt;
|
||||
p.k_cache = d_k_pool;
|
||||
p.v_cache = d_v_pool;
|
||||
p.q = d_q;
|
||||
p.mask = nullptr;
|
||||
p.o = d_o_paged;
|
||||
p.o_part = d_op;
|
||||
p.ml_part = d_ml;
|
||||
|
||||
// Dispatch
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
int G_check = p.q_head / p.kv_head;
|
||||
bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
|
||||
if (use_mma) {
|
||||
int tiles_total = (p.kv_len + 32 - 1) / 32;
|
||||
p.num_splits = num_splits(p.batch * p.kv_head, tiles_total);
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
int group_sz = p.q_head / p.kv_head;
|
||||
int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
|
||||
p.num_splits = num_splits(p.batch * p.kv_head, chunks_total);
|
||||
size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
paged_attn_decode_split_kv_kernel<<<
|
||||
dim3(p.batch * p.kv_head, 1, p.num_splits),
|
||||
dim3(32, group_sz), smem>>>(p);
|
||||
}
|
||||
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
|
||||
cudaDeviceSynchronize();
|
||||
|
||||
// Download paged output
|
||||
bf16* h_o_bf16 = (bf16*)malloc(sz_o);
|
||||
cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost);
|
||||
float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
|
||||
for (int i = 0; i < B * Hq * HEAD_DIM; i++)
|
||||
h_o_paged[i] = __bfloat162float(h_o_bf16[i]);
|
||||
|
||||
// Compare
|
||||
float max_err = 0.0f;
|
||||
int bad_idx = -1;
|
||||
for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
|
||||
float e = fabsf(h_o_paged[i] - h_o_ref[i]);
|
||||
if (e > max_err) { max_err = e; bad_idx = i; }
|
||||
}
|
||||
|
||||
bool pass = max_err < 0.02f;
|
||||
|
||||
if (pass) {
|
||||
printf("PASS (max_abs_err=%.4e)\n", max_err);
|
||||
} else {
|
||||
int b = bad_idx / (Hq * HEAD_DIM);
|
||||
int h = (bad_idx / HEAD_DIM) % Hq;
|
||||
int d = bad_idx % HEAD_DIM;
|
||||
printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
|
||||
max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
|
||||
// Print first 8 dims of first head
|
||||
printf(" ref[0..7]:");
|
||||
for (int i = 0; i < 8 && i < HEAD_DIM; i++)
|
||||
printf(" %.4f", h_o_ref[i]);
|
||||
printf("\n got[0..7]:");
|
||||
for (int i = 0; i < 8 && i < HEAD_DIM; i++)
|
||||
printf(" %.4f", h_o_paged[i]);
|
||||
printf("\n");
|
||||
}
|
||||
|
||||
free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt);
|
||||
free(h_k_cont); free(h_v_cont);
|
||||
free(h_o_ref); free(h_o_bf16); free(h_o_paged);
|
||||
cudaFree(d_q); cudaFree(d_o_paged); cudaFree(d_o_ref);
|
||||
cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
|
||||
cudaFree(d_op); cudaFree(d_ml);
|
||||
|
||||
return pass ? 0 : 1;
|
||||
}
|
||||
|
||||
int main() {
|
||||
int fail = 0;
|
||||
printf("=== Paged Decode vs CPU reference ===\n\n");
|
||||
|
||||
printf("-- scalar (G=1) --\n");
|
||||
fail += run_test<128>(1, 1, 1, 8, 128, 1);
|
||||
fail += run_test<128>(1, 4, 4, 128, 128, 2);
|
||||
fail += run_test<128>(2, 4, 4, 256, 128, 3);
|
||||
fail += run_test<128>(1, 4, 1, 64, 64, 4);
|
||||
|
||||
printf("-- scalar (G>1) --\n");
|
||||
fail += run_test<128>(1, 8, 2, 64, 128, 5);
|
||||
fail += run_test<128>(2, 16, 4, 128, 128, 6);
|
||||
|
||||
printf("-- varying head_dim --\n");
|
||||
fail += run_test<64>(1, 4, 2, 32, 128, 7);
|
||||
fail += run_test<256>(1, 2, 1, 16, 128, 8);
|
||||
fail += run_test<32>(1, 4, 2, 32, 64, 9);
|
||||
|
||||
printf("-- multi-batch --\n");
|
||||
fail += run_test<128>(3, 8, 2, 256, 128, 10);
|
||||
fail += run_test<128>(2, 32, 8, 512, 128, 11);
|
||||
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
printf("-- MMA (G>1, sm_80+) --\n");
|
||||
fail += run_test<128>(1, 16, 2, 256, 128, 12);
|
||||
fail += run_test<128>(2, 32, 4, 512, 128, 13);
|
||||
#endif
|
||||
|
||||
printf("\n%s (%d/%d failed)\n", fail ? "FAILED" : "ALL PASSED", fail, fail + (13 - fail + 1));
|
||||
return fail;
|
||||
}
|
||||
|
|
@ -5,21 +5,12 @@ nvcc -I csrc -arch=sm_89 -O3 \
|
|||
csrc/tests/attn_prefill_test.cu -o test && ./test
|
||||
*/
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <sys/time.h>
|
||||
#include "test_utils.cuh"
|
||||
#include "../kernels/attn_prefill_split_q.cuh"
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
#include "../kernels/attn_prefill_split_q_mma.cuh"
|
||||
#endif
|
||||
|
||||
static double now_ms() {
|
||||
struct timeval tv;
|
||||
gettimeofday(&tv, NULL);
|
||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
||||
}
|
||||
|
||||
// Launch the production prefill path (tensor-core MMA on sm_80+, else the
|
||||
// scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
|
||||
template <int HEAD_DIM>
|
||||
|
|
@ -48,45 +39,6 @@ static void dispatch_prefill(AttentionParams<bf16>& p) {
|
|||
}
|
||||
}
|
||||
|
||||
static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
|
||||
int B, int Hq, int Hk, int q_len, int kv_len, int D,
|
||||
int is_causal, int causal_off) {
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
int n_rep = Hq / Hk;
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int h = 0; h < Hq; h++) {
|
||||
for (int qi = 0; qi < q_len; qi++) {
|
||||
int kv_h = h / n_rep;
|
||||
float mv = -INFINITY, sv = 0.0f;
|
||||
float accum[256] = {0};
|
||||
int lim = is_causal ? min(kv_len, qi + causal_off + 1) : kv_len;
|
||||
for (int kj = 0; kj < lim; kj++) {
|
||||
float dot = 0.0f;
|
||||
for (int d = 0; d < D; d++)
|
||||
dot += Q[((b*Hq + h)*q_len + qi)*D + d]
|
||||
* K[((b*Hk + kv_h)*kv_len + kj)*D + d];
|
||||
dot *= scale;
|
||||
float nm = fmaxf(mv, dot);
|
||||
float al = expf(mv - nm);
|
||||
float be = expf(dot - nm);
|
||||
sv = sv * al + be;
|
||||
for (int d = 0; d < D; d++)
|
||||
accum[d] = accum[d] * al
|
||||
+ V[((b*Hk + kv_h)*kv_len + kj)*D + d] * be;
|
||||
mv = nm;
|
||||
}
|
||||
float inv = 1.0f / sv;
|
||||
for (int d = 0; d < D; d++)
|
||||
O[((b*Hq + h)*q_len + qi)*D + d] = accum[d] * inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
|
||||
static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
|
||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
||||
|
||||
// Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
|
||||
// Reports per-call latency and effective tensor-core TFLOP/s (2 matmuls:
|
||||
// QK^T and P@V, each 2*B*Hq*ql*kl*D flops; halved for causal).
|
||||
|
|
@ -208,7 +160,7 @@ int main() {
|
|||
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
|
||||
|
||||
float* ref=new float[nQ];
|
||||
cpu_attention(hQ,hK,hV,ref,B,Hq,Hk,ql,kl,D,causal,0);
|
||||
cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0);
|
||||
|
||||
float max_err=0;
|
||||
for (size_t i=0;i<nQ;i++) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
#pragma once
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cmath>
|
||||
#include <chrono>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
inline bf16 f2bf(float x) { return __float2bfloat16(x); }
|
||||
inline float bf2f(bf16 x) { return __bfloat162float(x); }
|
||||
|
||||
inline float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
||||
|
||||
inline double now_ms() {
|
||||
using namespace std::chrono;
|
||||
return duration_cast<milliseconds>(steady_clock::now().time_since_epoch()).count();
|
||||
}
|
||||
|
||||
inline int compute_num_splits(int base_blocks, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
|
||||
if (n > tiles_total) n = tiles_total;
|
||||
if (n > 32) n = 32;
|
||||
if (n < 1) n = 1;
|
||||
return n;
|
||||
}
|
||||
|
||||
#define CUDA_CHECK(call) \
|
||||
do { \
|
||||
cudaError_t _e = (call); \
|
||||
if (_e != cudaSuccess) { \
|
||||
printf("CUDA error %s at %s:%d\n", cudaGetErrorString(_e), __FILE__, __LINE__); \
|
||||
exit(1); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Generic CPU reference for multi-query / grouped-query attention.
|
||||
// Tensor shapes (all float*):
|
||||
// Q : [B, Hq, q_len, D]
|
||||
// K : [B, Hk, kv_len, D]
|
||||
// V : [B, Hk, kv_len, D]
|
||||
// O : [B, Hq, q_len, D]
|
||||
// mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported.
|
||||
static void cpu_attention_ref(
|
||||
const float* Q, const float* K, const float* V, const bool* mask,
|
||||
float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D,
|
||||
int is_causal, int causal_offset
|
||||
) {
|
||||
float scale = 1.0f / sqrtf((float)D);
|
||||
int n_rep = Hq / Hk;
|
||||
for (int b = 0; b < B; 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 = kv_len;
|
||||
if (is_causal) {
|
||||
int c = qi + causal_offset + 1;
|
||||
lim = (c < kv_len) ? c : kv_len;
|
||||
}
|
||||
for (int kj = 0; kj < lim; kj++) {
|
||||
if (mask != nullptr && q_len == 1) {
|
||||
if (!mask[b * kv_len + kj]) continue;
|
||||
}
|
||||
float dot = 0.0f;
|
||||
size_t q_idx = ((size_t)b * Hq + h) * q_len + qi;
|
||||
size_t kv_idx = ((size_t)b * Hk + kv_h) * kv_len + kj;
|
||||
for (int d = 0; d < D; d++)
|
||||
dot += Q[q_idx * D + d] * K[kv_idx * D + d];
|
||||
dot *= scale;
|
||||
float nm = fmaxf(mv, dot);
|
||||
float a = expf(mv - nm);
|
||||
float b_exp = expf(dot - nm);
|
||||
sv = sv * a + b_exp;
|
||||
for (int d = 0; d < D; d++)
|
||||
accum[d] = accum[d] * a + V[kv_idx * D + d] * b_exp;
|
||||
mv = nm;
|
||||
}
|
||||
float inv = 1.0f / sv;
|
||||
size_t o_idx = ((size_t)b * Hq + h) * q_len + qi;
|
||||
for (int d = 0; d < D; d++)
|
||||
O[o_idx * D + d] = accum[d] * inv;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue