refactor: extract bench_kernel and dispatch_by_head_dim into test_utils
This commit is contained in:
parent
2c3cef1c87
commit
fd6d25ad86
|
|
@ -61,90 +61,64 @@ 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++) {
|
||||
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<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);
|
||||
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; 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.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<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;
|
||||
|
||||
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;
|
||||
auto launch = [&]() { dispatch_decode(p, sc); };
|
||||
double flops = 4.0 * B * Hq * (double)sl * D;
|
||||
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);
|
||||
B, Hq, Hk, 1, sl, D, 0);
|
||||
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);
|
||||
cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
|
||||
cudaFree(sc.o_part); cudaFree(sc.ml_part);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in New Issue