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

refactor: extract bench_kernel and dispatch_by_head_dim into test_utils

This commit is contained in:
ViperEkura 2026-07-12 00:02:05 +08:00
parent 2c3cef1c87
commit fd6d25ad86
3 changed files with 107 additions and 86 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
csrc/tests/attn_decode_test.cu
View File

@ -61,56 +61,44 @@ static void dispatch_decode_t(AttentionParams<bf16>& p, DecodeScratch& sc) {
}
static void dispatch_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
switch (p.head_dim) {
case 32: dispatch_decode_t<32>(p, sc); break;
case 64: dispatch_decode_t<64>(p, sc); break;
case 128: dispatch_decode_t<128>(p, sc); break;
case 256: dispatch_decode_t<256>(p, sc); break;
default: printf("bench: unsupported D=%d\n", p.head_dim);
}
dispatch_by_head_dim(p.head_dim, [&]<int D>() { dispatch_decode_t<D>(p, sc); });
}
// 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.
// FLOP/s = 2 matmuls (q@K^T, P@V), each 2*B*Hq*kv*D flops.
// Bytes = K + V read = 2 * B*Hk*kv*D * sizeof(bf16).
static void bench() {
const int cfgs[][5] = {
{1, 32, 4, 512, 128}, // B,Hq,Hk,seq,D
{1, 32, 4, 512, 128}, // B, Hq, Hk, kv_len, D
{1, 32, 4, 1024, 128},
{1, 32, 4, 2048, 128},
{1, 32, 4, 4096, 128},
{16, 32, 4, 2048, 128},
{32, 32, 4, 1024, 128},
};
int n = sizeof(cfgs)/sizeof(cfgs[0]);
const int WARMUP = 10, ITERS = 100;
printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
printf("%-46s | %10s | %10s | %10s\n",
"config", "latency", "bandwidth", "throughput");
printf("---------------------------------------------------------------"
"----------------------------\n");
print_bench_header();
for (int ci = 0; ci < n; ci++) {
for (int ci = 0; ci < 6; ci++) {
int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
int sl = cfgs[ci][3], D = cfgs[ci][4];
size_t nQ=(size_t)B*Hq*D, nKV=(size_t)B*Hk*sl*D;
size_t nQ = (size_t)B * Hq * D;
size_t nKV = (size_t)B * Hk * sl * D;
bf16 *dQ,*dK,*dV,*dO,*tmp;
bf16 *dQ, *dK, *dV, *dO;
cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
for (size_t i = 0; i < nQ; i++) tmp[i] = f2bf(randf());
cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
delete[] tmp;
AttentionParams<bf16> p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
p.use_mask=0; p.is_causal=0; p.causal_offset=0;
p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
p.head_dim = D; p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
p.scale = 1.0f / sqrtf((float)D);
p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
@ -118,33 +106,19 @@ static void bench() {
cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
for (int i=0;i<WARMUP;i++) dispatch_decode(p, sc);
cudaDeviceSynchronize();
cudaError_t err=cudaGetLastError();
if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
cudaEventRecord(s);
for (int i=0;i<ITERS;i++) dispatch_decode(p, sc);
cudaEventRecord(e); cudaEventSynchronize(e);
float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
auto launch = [&]() { dispatch_decode(p, sc); };
double flops = 4.0 * B * Hq * (double)sl * D;
double tflops = flops/(ms*1e-3)/1e12;
// HBM traffic: K + V read (B*Hk*sl*D each), bf16; Q/O negligible.
double bytes = 2.0 * (2.0*nKV);
double gbps = bytes/(ms*1e-3)/1e9;
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 causal=%d",
B, Hq, Hk, 1, sl, D, 0);
printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
cfg, ms, gbps, tflops);
print_bench_row(cfg, r);
cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
cudaFree(sc.o_part); cudaFree(sc.ml_part);
delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
}
}

41
csrc/tests/attn_paged_decode_test.cu
View File

@ -221,17 +221,16 @@ static const TestCase TESTS[] = {
};
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;
}
bool matched = false;
int r = 0;
dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
matched = true;
r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
});
return matched ? r : 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) {
@ -281,43 +280,27 @@ static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
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;
auto launch = [&]() { launch_paged_decode<HEAD_DIM>(pa); };
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;
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);
printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
cfg, ms, gbps, tflops);
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);
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");
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);

64
csrc/tests/test_utils.cuh
View File

@ -37,6 +37,70 @@ inline int compute_num_splits(int base_blocks, int tiles_total) {
} \
} while (0)
struct BenchResult {
float ms;
double gbps;
double tflops;
};
template <typename Fn>
BenchResult bench_kernel(Fn launch, int warmup, int iters,
double flops, double bytes) {
for (int i = 0; i < warmup; i++) launch();
cudaDeviceSynchronize();
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
printf("CUDA error before bench: %s\n", cudaGetErrorString(err));
return {0, 0, 0};
}
cudaEvent_t s, e;
cudaEventCreate(&s); cudaEventCreate(&e);
cudaEventRecord(s);
for (int i = 0; i < iters; i++) launch();
cudaEventRecord(e); cudaEventSynchronize(e);
float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= iters;
cudaEventDestroy(s); cudaEventDestroy(e);
return {ms, bytes / (ms * 1e-3) / 1e9, flops / (ms * 1e-3) / 1e12};
}
inline void print_bench_header() {
printf("%-46s | %10s | %10s | %10s\n",
"config", "latency", "bandwidth", "throughput");
printf("---------------------------------------------------------------"
"----------------------------\n");
}
inline void print_bench_row(const char* cfg, const BenchResult& r) {
printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
cfg, r.ms, r.gbps, r.tflops);
}
template <int... Ds>
struct _HeadSwitch;
template <int D>
struct _HeadSwitch<D> {
template <typename Fn>
static void call(int hd, Fn&& fn) { if (hd == D) fn.template operator()<D>(); }
};
template <int D, int... Rest>
struct _HeadSwitch<D, Rest...> {
template <typename Fn>
static void call(int hd, Fn&& fn) {
if (hd == D) fn.template operator()<D>();
else _HeadSwitch<Rest...>::call(hd, fn);
}
};
// Default set: 32, 64, 128, 256
template <typename Fn>
void dispatch_by_head_dim(int head_dim, Fn&& fn) {
_HeadSwitch<32, 64, 128, 256>::call(head_dim, fn);
}
// Generic CPU reference for multi-query / grouped-query attention.
// Tensor shapes (all float*):
// Q : [B, Hq, q_len, D]