refactor: extract bench_kernel and dispatch_by_head_dim into test_utils

csrc/tests/attn_decode_test.cu

@@ -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) {
+    dispatch_by_head_dim(p.head_dim, [&]<int D>() { dispatch_decode_t<D>(p, sc); });
-        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);
-    }
 }
 
 // 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",
+    print_bench_header();
-           "config", "latency", "bandwidth", "throughput");
-    printf("---------------------------------------------------------------"
-           "----------------------------\n");
 
-    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.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
-        p.use_mask=0; p.is_causal=0; p.causal_offset=0;
+        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);
+        auto launch = [&]() { 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;
-
         double flops = 4.0 * B * Hq * (double)sl * D;
-        double tflops = flops/(ms*1e-3)/1e12;
+        double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-        // HBM traffic: K + V read (B*Hk*sl*D each), bf16; Q/O negligible.
+        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
-        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 causal=%d",
                  B, Hq, Hk, 1, sl, D, 0);
-        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
+        print_bench_row(cfg, r);
-               cfg, ms, gbps, tflops);
 
         cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
         cudaFree(sc.o_part); cudaFree(sc.ml_part);
-        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
     }
 }
 

csrc/tests/attn_paged_decode_test.cu

@@ -221,17 +221,16 @@ static const TestCase TESTS[] = {
 };
 
 static int dispatch_test(const TestCase& tc) {
-    switch (tc.head_dim) {
+    bool matched = false;
-        case 32:  return run_test<32>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+    int r = 0;
-        case 64:  return run_test<64>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+    dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
-        case 128: return run_test<128>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        matched = true;
-        case 256: return run_test<256>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
+        r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.seed);
-        default:  return 1;
+    });
-    }
+    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);
+    auto launch = [&]() { 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 bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-    double gbps = bytes / (ms * 1e-3) / 1e9;
+    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",
+    print_bench_row(cfg, r);
-           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",
+    print_bench_header();
-           "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);

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 <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]