diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 5454f85..573e643 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -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", diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py index 5633e9b..63764f4 100644 --- a/astrai/extension/loader.py +++ b/astrai/extension/loader.py @@ -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] = {} diff --git a/astrai/extension/ops.py b/astrai/extension/ops.py index 39f8faa..0dc4d9c 100644 --- a/astrai/extension/ops.py +++ b/astrai/extension/ops.py @@ -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) diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index 7f7c8c4..f87a8f8 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attn_paged_decode.cu @@ -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()); diff --git a/csrc/tests/attn_decode_test.cu b/csrc/tests/attn_decode_test.cu index dfefda6..a716856 100644 --- a/csrc/tests/attn_decode_test.cu +++ b/csrc/tests/attn_decode_test.cu @@ -5,21 +5,12 @@ nvcc -I csrc -arch=sm_89 -O3 \ csrc/tests/attn_decode_test.cu -o test && ./test */ -#include -#include -#include -#include +#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& p) { template static void launch_mma_decode(AttentionParams& 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& p, DecodeScratch& sc) { static void launch_scalar_decode(AttentionParams& 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& 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 +#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 +static void launch_paged_decode(PagedAttentionParams& 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 + <<>>(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); +} + +template +static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) { + printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM); + fflush(stdout); + + int 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 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(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 +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 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(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(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; +} diff --git a/csrc/tests/attn_paged_vs_contiguous.cu b/csrc/tests/attn_paged_vs_contiguous.cu deleted file mode 100644 index 2b26304..0000000 --- a/csrc/tests/attn_paged_vs_contiguous.cu +++ /dev/null @@ -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 -#include -#include -#include -#include -#include "../kernels/attn_paged_decode_split_kv.cuh" -#ifndef ASTRAI_NO_MMA -#include "../kernels/attn_paged_decode_split_kv_mma.cuh" -#endif - -using bf16 = __nv_bfloat16; - -static int num_splits(int base_blocks, int tiles_total) { - int sm_count = 0; - cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); - int n = (2 * sm_count + base_blocks - 1) / base_blocks; - return std::max(1, std::min(n, std::min(tiles_total, 32))); -} - -// Copy contiguous K/V from page pool (reference gather) -static void gather_kv_cpu( - const bf16* h_k_pool, const bf16* h_v_pool, - const int64_t* h_pt, int B, int Hkv, int kv_len, - int page_size, int head_dim, - bf16* h_k, bf16* h_v) -{ - int max_pages = (kv_len + page_size - 1) / page_size; - size_t page_stride = (size_t)page_size * Hkv * head_dim; - for (int b = 0; b < B; b++) { - for (int pos = 0; pos < kv_len; pos++) { - int log_pg = pos / page_size; - int pg_off = pos % page_size; - int phys = (int)h_pt[b * max_pages + log_pg]; - for (int h = 0; h < Hkv; h++) { - size_t src_base = (size_t)phys * page_stride - + (size_t)pg_off * Hkv * head_dim - + h * head_dim; - size_t dst_base = ((size_t)b * kv_len + pos) * Hkv * head_dim + h * head_dim; - memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16)); - memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16)); - } - } - } -} - -template -static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed) { - printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d ... ", B, Hq, Hkv, kv_len, page_size, HEAD_DIM); - fflush(stdout); - - int G = Hq / Hkv; - int max_pages = (kv_len + page_size - 1) / page_size; - int n_phys_pages = B * max_pages; - - // ---- allocate ---- - bf16 *d_q, *d_o_paged, *d_o_ref; - bf16 *d_k_pool, *d_v_pool; - int64_t* d_pt; - float *d_op, *d_ml; - - size_t sz_q = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16); - size_t sz_o = sz_q; - size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16); - size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t); - int max_splits = 32; - size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float); - size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float); - - cudaMalloc(&d_q, sz_q); - cudaMalloc(&d_o_paged, sz_o); - cudaMalloc(&d_o_ref, sz_o); - cudaMalloc(&d_k_pool, sz_kv); - cudaMalloc(&d_v_pool, sz_kv); - cudaMalloc(&d_pt, sz_pt); - cudaMalloc(&d_op, sz_op); - cudaMalloc(&d_ml, sz_ml); - - // ---- init: deterministic random using seed ---- - srand(seed); - auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; }; - - // Q - bf16* h_q = (bf16*)malloc(sz_q); - for (int i = 0; i < B * Hq * HEAD_DIM; i++) - h_q[i] = __float2bfloat16(rnd()); - cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice); - - // Page pool K/V - bf16* h_k_pool = (bf16*)malloc(sz_kv); - bf16* h_v_pool = (bf16*)malloc(sz_kv); - size_t ps = (size_t)page_size * Hkv * HEAD_DIM; - for (int pg = 0; pg < n_phys_pages; pg++) { - for (int off = 0; off < page_size; off++) { - for (int h = 0; h < Hkv; h++) { - for (int d = 0; d < HEAD_DIM; d++) { - float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d)); - size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM + h * HEAD_DIM + d; - h_k_pool[idx] = __float2bfloat16(v); - h_v_pool[idx] = __float2bfloat16(v * 0.3f); - } - } - } - } - cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice); - cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice); - - // Page table - int64_t* h_pt = (int64_t*)malloc(sz_pt); - int next_pg = 0; - for (int b = 0; b < B; b++) - for (int p = 0; p < max_pages; p++) - h_pt[b * max_pages + p] = next_pg++; - cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice); - - // ---- reference: gather contiguous K/V, then run CPU online-softmax ---- - bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16)); - bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16)); - gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont); - - float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float)); - float sscale = 1.0f / sqrtf((float)HEAD_DIM); - for (int b = 0; b < B; b++) { - for (int hq = 0; hq < Hq; hq++) { - int hkv = hq / G; - size_t q_base = (size_t)b * Hq * HEAD_DIM + (size_t)hq * HEAD_DIM; - size_t kv_base = ((size_t)b * kv_len) * Hkv * HEAD_DIM + (size_t)hkv * HEAD_DIM; - - float m = -1e30f, d = 0.0f; - float acc[256] = {0.0f}; - for (int pos = 0; pos < kv_len; pos++) { - float s = 0.0f; - for (int dim = 0; dim < HEAD_DIM; dim++) - s += __bfloat162float(h_q[q_base + dim]) - * __bfloat162float(h_k_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]); - s *= sscale; - - float nm = fmaxf(m, s); - float a = expf(m - nm); - float b = expf(s - nm); - d = d * a + b; - for (int dim = 0; dim < HEAD_DIM; dim++) - acc[dim] = acc[dim] * a + __bfloat162float(h_v_cont[kv_base + (size_t)pos * Hkv * HEAD_DIM + dim]) * b; - m = nm; - } - for (int dim = 0; dim < HEAD_DIM; dim++) - h_o_ref[b * Hq * HEAD_DIM + hq * HEAD_DIM + dim] = acc[dim] / d; - } - } - - // ---- paged decode kernel ---- - float scale_val = 1.0f / sqrtf((float)HEAD_DIM); - PagedAttentionParams p; - p.batch = B; - p.q_head = Hq; - p.kv_head = Hkv; - p.q_len = 1; - p.kv_len = kv_len; - p.head_dim = HEAD_DIM; - p.use_mask = 0; - p.is_causal = 0; - p.causal_offset = 0; - p.num_splits = 1; - p.scale = scale_val; - p.page_size = page_size; - p.max_pages = max_pages; - p.page_table = d_pt; - p.k_cache = d_k_pool; - p.v_cache = d_v_pool; - p.q = d_q; - p.mask = nullptr; - p.o = d_o_paged; - p.o_part = d_op; - p.ml_part = d_ml; - - // Dispatch -#ifndef ASTRAI_NO_MMA - int G_check = p.q_head / p.kv_head; - bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32; - if (use_mma) { - int tiles_total = (p.kv_len + 32 - 1) / 32; - p.num_splits = num_splits(p.batch * p.kv_head, tiles_total); - paged_attn_decode_split_kv_mma_kernel - <<>>(p); - } else -#endif - { - int group_sz = p.q_head / p.kv_head; - int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; - p.num_splits = num_splits(p.batch * p.kv_head, chunks_total); - size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16); - paged_attn_decode_split_kv_kernel<<< - dim3(p.batch * p.kv_head, 1, p.num_splits), - dim3(32, group_sz), smem>>>(p); - } - paged_attn_decode_combine_kernel<<>>(p); - cudaDeviceSynchronize(); - - // Download paged output - bf16* h_o_bf16 = (bf16*)malloc(sz_o); - cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost); - float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float)); - for (int i = 0; i < B * Hq * HEAD_DIM; i++) - h_o_paged[i] = __bfloat162float(h_o_bf16[i]); - - // Compare - float max_err = 0.0f; - int bad_idx = -1; - for (int i = 0; i < B * Hq * HEAD_DIM; i++) { - float e = fabsf(h_o_paged[i] - h_o_ref[i]); - if (e > max_err) { max_err = e; bad_idx = i; } - } - - bool pass = max_err < 0.02f; - - if (pass) { - printf("PASS (max_abs_err=%.4e)\n", max_err); - } else { - int b = bad_idx / (Hq * HEAD_DIM); - int h = (bad_idx / HEAD_DIM) % Hq; - int d = bad_idx % HEAD_DIM; - printf("FAIL (max_abs_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n", - max_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]); - // Print first 8 dims of first head - printf(" ref[0..7]:"); - for (int i = 0; i < 8 && i < HEAD_DIM; i++) - printf(" %.4f", h_o_ref[i]); - printf("\n got[0..7]:"); - for (int i = 0; i < 8 && i < HEAD_DIM; i++) - printf(" %.4f", h_o_paged[i]); - printf("\n"); - } - - free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt); - free(h_k_cont); free(h_v_cont); - free(h_o_ref); free(h_o_bf16); free(h_o_paged); - cudaFree(d_q); cudaFree(d_o_paged); cudaFree(d_o_ref); - cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt); - cudaFree(d_op); cudaFree(d_ml); - - return pass ? 0 : 1; -} - -int main() { - int fail = 0; - printf("=== Paged Decode vs CPU reference ===\n\n"); - - printf("-- scalar (G=1) --\n"); - fail += run_test<128>(1, 1, 1, 8, 128, 1); - fail += run_test<128>(1, 4, 4, 128, 128, 2); - fail += run_test<128>(2, 4, 4, 256, 128, 3); - fail += run_test<128>(1, 4, 1, 64, 64, 4); - - printf("-- scalar (G>1) --\n"); - fail += run_test<128>(1, 8, 2, 64, 128, 5); - fail += run_test<128>(2, 16, 4, 128, 128, 6); - - printf("-- varying head_dim --\n"); - fail += run_test<64>(1, 4, 2, 32, 128, 7); - fail += run_test<256>(1, 2, 1, 16, 128, 8); - fail += run_test<32>(1, 4, 2, 32, 64, 9); - - printf("-- multi-batch --\n"); - fail += run_test<128>(3, 8, 2, 256, 128, 10); - fail += run_test<128>(2, 32, 8, 512, 128, 11); - -#ifndef ASTRAI_NO_MMA - printf("-- MMA (G>1, sm_80+) --\n"); - fail += run_test<128>(1, 16, 2, 256, 128, 12); - fail += run_test<128>(2, 32, 4, 512, 128, 13); -#endif - - printf("\n%s (%d/%d failed)\n", fail ? "FAILED" : "ALL PASSED", fail, fail + (13 - fail + 1)); - return fail; -} diff --git a/csrc/tests/attn_prefill_test.cu b/csrc/tests/attn_prefill_test.cu index b80df36..e67885e 100644 --- a/csrc/tests/attn_prefill_test.cu +++ b/csrc/tests/attn_prefill_test.cu @@ -5,21 +5,12 @@ nvcc -I csrc -arch=sm_89 -O3 \ csrc/tests/attn_prefill_test.cu -o test && ./test */ -#include -#include -#include -#include +#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 @@ -48,45 +39,6 @@ static void dispatch_prefill(AttentionParams& 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 +#include +#include +#include +#include + +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(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; + } + } + } +}