From fd6d25ad867a4345a6c0400985659efd1a83eb54 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 12 Jul 2026 00:02:05 +0800 Subject: [PATCH] refactor: extract bench_kernel and dispatch_by_head_dim into test_utils --- csrc/tests/attn_decode_test.cu | 88 ++++++++++------------------ csrc/tests/attn_paged_decode_test.cu | 41 ++++--------- csrc/tests/test_utils.cuh | 64 ++++++++++++++++++++ 3 files changed, 107 insertions(+), 86 deletions(-) diff --git a/csrc/tests/attn_decode_test.cu b/csrc/tests/attn_decode_test.cu index a716856..4e88042 100644 --- a/csrc/tests/attn_decode_test.cu +++ b/csrc/tests/attn_decode_test.cu @@ -61,90 +61,64 @@ static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) { } static void dispatch_decode(AttentionParams& 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, [&]() { dispatch_decode_t(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++) { - 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; + 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; + size_t nKV = (size_t)B * Hk * sl * D; - bf16 *dQ,*dK,*dV,*dO,*tmp; - 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]; - for (size_t i=0;i 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 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.scale=1.0f/sqrtf((float)D); - p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO; + 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; DecodeScratch sc; 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(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, [&]() { + matched = true; + r = run_test(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 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(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; - + auto launch = [&]() { launch_paged_decode(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); diff --git a/csrc/tests/test_utils.cuh b/csrc/tests/test_utils.cuh index d8a8937..f8c05be 100644 --- a/csrc/tests/test_utils.cuh +++ b/csrc/tests/test_utils.cuh @@ -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 +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 +struct _HeadSwitch; + +template +struct _HeadSwitch { + template + static void call(int hd, Fn&& fn) { if (hd == D) fn.template operator()(); } +}; + +template +struct _HeadSwitch { + template + static void call(int hd, Fn&& fn) { + if (hd == D) fn.template operator()(); + else _HeadSwitch::call(hd, fn); + } +}; + +// Default set: 32, 64, 128, 256 +template +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]