test: bench production MMA attention path with FLOP/s and bandwidth
This commit is contained in:
parent
cbd140340d
commit
9027fdc546
|
|
@ -10,6 +10,9 @@ nvcc -I csrc -arch=sm_89 -O3 \
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
#include <sys/time.h>
|
#include <sys/time.h>
|
||||||
#include "../kernels/gqa_decode_attn.cuh"
|
#include "../kernels/gqa_decode_attn.cuh"
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
#include "../kernels/gqa_decode_attn_mma.cuh"
|
||||||
|
#endif
|
||||||
|
|
||||||
static double now_ms() {
|
static double now_ms() {
|
||||||
struct timeval tv;
|
struct timeval tv;
|
||||||
|
|
@ -17,6 +20,80 @@ static double now_ms() {
|
||||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
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).
|
||||||
|
struct DecodeScratch {
|
||||||
|
float* o_part = nullptr;
|
||||||
|
float* ml_part = nullptr;
|
||||||
|
};
|
||||||
|
|
||||||
|
// Launch the production decode path (tensor-core head-packing MMA on sm_80+,
|
||||||
|
// scalar fallback otherwise), mirroring dispatch_decode() in gqa_decode_attn.cu.
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
static bool decode_use_mma(const GQAParams& p) {
|
||||||
|
int G = p.q_head / p.kv_head;
|
||||||
|
return !p.use_mask && G > 1 && G <= 16;
|
||||||
|
}
|
||||||
|
|
||||||
|
static int decode_num_splits(const GQAParams& p, int tiles_total) {
|
||||||
|
int sm_count = 0;
|
||||||
|
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||||
|
int base_blocks = p.kv_head * p.batch;
|
||||||
|
int desired = 2 * (sm_count > 0 ? sm_count : 64);
|
||||||
|
int n = (desired + base_blocks - 1) / base_blocks;
|
||||||
|
return max(1, min(n, min(tiles_total, 32)));
|
||||||
|
}
|
||||||
|
|
||||||
|
template <int HEAD_DIM, int BC>
|
||||||
|
static void launch_mma_decode(GQAParams& p, DecodeScratch& sc) {
|
||||||
|
constexpr int BR = 16, LD = HEAD_DIM;
|
||||||
|
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||||
|
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||||
|
int num_splits = decode_num_splits(p, tiles_total);
|
||||||
|
|
||||||
|
if (num_splits <= 1) {
|
||||||
|
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
||||||
|
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||||
|
gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
|
||||||
|
<<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
|
||||||
|
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||||
|
gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
|
||||||
|
<<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
|
||||||
|
p, sc.o_part, sc.ml_part, num_splits);
|
||||||
|
gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
|
||||||
|
sc.o_part, sc.ml_part, p.o, num_splits, p.head_dim);
|
||||||
|
}
|
||||||
|
#endif
|
||||||
|
|
||||||
|
static void launch_scalar_decode(const GQAParams& p) {
|
||||||
|
int gs = p.q_head / p.kv_head;
|
||||||
|
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||||
|
gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, gs), smem>>>(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <int HEAD_DIM>
|
||||||
|
static void dispatch_decode_t(GQAParams& p, DecodeScratch& sc) {
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
|
||||||
|
#endif
|
||||||
|
(void)sc;
|
||||||
|
launch_scalar_decode(p);
|
||||||
|
}
|
||||||
|
|
||||||
|
static void dispatch_decode(GQAParams& 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
static void cpu_decode(const float* Q, const float* K, const float* V,
|
static void cpu_decode(const float* Q, const float* K, const float* V,
|
||||||
const bool* mask, float* O,
|
const bool* mask, float* O,
|
||||||
int B, int Hq, int Hk, int seq_len, int D) {
|
int B, int Hq, int Hk, int seq_len, int D) {
|
||||||
|
|
@ -54,6 +131,84 @@ static bf16 f2bf(float x) { return __float2bfloat16(x); }
|
||||||
static float bf2f(bf16 x) { return __bfloat162float(x); }
|
static float bf2f(bf16 x) { return __bfloat162float(x); }
|
||||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
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.
|
||||||
|
// 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, 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");
|
||||||
|
|
||||||
|
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;
|
||||||
|
|
||||||
|
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);
|
||||||
|
|
||||||
|
GQAParams 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;
|
||||||
|
|
||||||
|
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;
|
||||||
|
|
||||||
|
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);
|
||||||
|
|
||||||
|
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
|
||||||
|
cudaFree(sc.o_part);cudaFree(sc.ml_part);
|
||||||
|
delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
int main() {
|
int main() {
|
||||||
const int configs[][5] = {
|
const int configs[][5] = {
|
||||||
{1, 2, 1, 64, 32}, // B,Hq,Hk,seq_len,D
|
{1, 2, 1, 64, 32}, // B,Hq,Hk,seq_len,D
|
||||||
|
|
@ -92,18 +247,17 @@ int main() {
|
||||||
|
|
||||||
GQAParams p;
|
GQAParams 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.head_dim=D;
|
||||||
p.use_mask=1; p.is_causal=0; p.causal_offset=0;
|
p.use_mask=0; p.is_causal=0; p.causal_offset=0;
|
||||||
p.scale=1.0f/sqrtf((float)D);
|
p.scale=1.0f/sqrtf((float)D);
|
||||||
p.q=dQ; p.k=dK; p.v=dV; p.mask=dMask; p.o=dO;
|
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
|
||||||
|
|
||||||
size_t smem=DC_CHUNK*D*sizeof(bf16);
|
// Split-K scratch (max 32 splits), sized for the production MMA path.
|
||||||
dim3 block(32, gs);
|
DecodeScratch sc;
|
||||||
dim3 grid(B*Hk);
|
cudaMalloc(&sc.o_part, (size_t)B*Hq*32*D*sizeof(float));
|
||||||
printf("grid=(%d,1,1) block=(%d,%d,1) smem=%zu\n",
|
cudaMalloc(&sc.ml_part, (size_t)B*Hq*32*2*sizeof(float));
|
||||||
grid.x, block.x, block.y, smem);
|
|
||||||
|
|
||||||
double t0=now_ms();
|
double t0=now_ms();
|
||||||
gqa_decode_attn_kernel<<<grid,block,smem>>>(p);
|
dispatch_decode(p, sc);
|
||||||
cudaDeviceSynchronize();
|
cudaDeviceSynchronize();
|
||||||
double kms=now_ms()-t0;
|
double kms=now_ms()-t0;
|
||||||
cudaError_t err=cudaGetLastError();
|
cudaError_t err=cudaGetLastError();
|
||||||
|
|
@ -123,8 +277,10 @@ int main() {
|
||||||
printf("kernel: %.3f ms max_err: %.6e\n\n",kms,max_err);
|
printf("kernel: %.3f ms max_err: %.6e\n\n",kms,max_err);
|
||||||
|
|
||||||
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
|
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
|
||||||
|
cudaFree(sc.o_part);cudaFree(sc.ml_part);
|
||||||
delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
|
delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
|
||||||
}
|
}
|
||||||
printf("All tests passed!\n");
|
printf("All tests passed!\n");
|
||||||
|
bench();
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,9 @@ nvcc -I csrc -arch=sm_89 -O3 \
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
#include <sys/time.h>
|
#include <sys/time.h>
|
||||||
#include "../kernels/gqa_prefill_attn.cuh"
|
#include "../kernels/gqa_prefill_attn.cuh"
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
#include "../kernels/gqa_prefill_attn_mma.cuh"
|
||||||
|
#endif
|
||||||
|
|
||||||
static double now_ms() {
|
static double now_ms() {
|
||||||
struct timeval tv;
|
struct timeval tv;
|
||||||
|
|
@ -17,6 +20,32 @@ static double now_ms() {
|
||||||
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
|
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 gqa_prefill_attn.cu.
|
||||||
|
template <int HEAD_DIM>
|
||||||
|
static void launch_prefill(GQAParams& p) {
|
||||||
|
#ifndef ASTRAI_NO_MMA
|
||||||
|
constexpr int WARPS = 4, BC = 16, BR = 16;
|
||||||
|
constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2;
|
||||||
|
dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
|
||||||
|
dim3 block(WARPS * 32, 1, 1);
|
||||||
|
gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
|
||||||
|
#else
|
||||||
|
constexpr int G = 8, ROWS = 32, P_BC = 32;
|
||||||
|
dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
|
||||||
|
dim3 block(G, ROWS, 1);
|
||||||
|
gqa_prefill_attn_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
|
||||||
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
|
static void dispatch_prefill(GQAParams& p) {
|
||||||
|
switch (p.head_dim) {
|
||||||
|
case 64: launch_prefill<64>(p); break;
|
||||||
|
case 128: launch_prefill<128>(p); break;
|
||||||
|
default: printf("bench: unsupported D=%d\n", p.head_dim);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
static void cpu_attention(const float* Q, const float* K, const float* V, float* O,
|
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 B, int Hq, int Hk, int q_len, int kv_len, int D,
|
||||||
int is_causal, int causal_off) {
|
int is_causal, int causal_off) {
|
||||||
|
|
@ -56,6 +85,78 @@ static __nv_bfloat16 f2bf(float x) { return __float2bfloat16(x); }
|
||||||
static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
|
static float bf2f(__nv_bfloat16 x) { return __bfloat162float(x); }
|
||||||
static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; }
|
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).
|
||||||
|
static void bench() {
|
||||||
|
const int cfgs[][7] = {
|
||||||
|
{1,32,4,512,512,128,0},
|
||||||
|
{1,32,4,1024,1024,128,0},
|
||||||
|
{1,32,4,2048,2048,128,0},
|
||||||
|
{1,32,4,2048,2048,128,1},
|
||||||
|
{4,32,4,2048,2048,128,1},
|
||||||
|
{1,32,4,4096,4096,128,1},
|
||||||
|
};
|
||||||
|
int n = sizeof(cfgs)/sizeof(cfgs[0]);
|
||||||
|
const int WARMUP = 10, ITERS = 50;
|
||||||
|
printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
|
||||||
|
printf("%-46s | %10s | %10s | %10s\n",
|
||||||
|
"config", "latency", "bandwidth", "throughput");
|
||||||
|
printf("---------------------------------------------------------------"
|
||||||
|
"----------------------------\n");
|
||||||
|
|
||||||
|
for (int ci = 0; ci < n; ci++) {
|
||||||
|
int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
|
||||||
|
int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
|
||||||
|
size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*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);
|
||||||
|
|
||||||
|
GQAParams p;
|
||||||
|
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
|
||||||
|
p.use_mask=0; p.is_causal=causal; 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;
|
||||||
|
|
||||||
|
for (int i=0;i<WARMUP;i++) dispatch_prefill(p);
|
||||||
|
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_prefill(p);
|
||||||
|
cudaEventRecord(e); cudaEventSynchronize(e);
|
||||||
|
float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
|
||||||
|
|
||||||
|
double flops = 4.0*B*Hq*(double)ql*kl*D;
|
||||||
|
if (causal) flops *= 0.5;
|
||||||
|
double tflops = flops/(ms*1e-3)/1e12;
|
||||||
|
// HBM traffic: Q + O (B*Hq*ql*D each) + K + V (B*Hk*kl*D each), bf16.
|
||||||
|
double bytes = 2.0 * (2.0*nQ + 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,ql,kl,D,causal);
|
||||||
|
printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
|
||||||
|
cfg, ms, gbps, tflops);
|
||||||
|
|
||||||
|
cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
|
||||||
|
delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
int main() {
|
int main() {
|
||||||
const int configs[][7] = {
|
const int configs[][7] = {
|
||||||
{1,2,1,64,128,64,0}, // tiny: B,Hq,Hk,q,kv,D,causal
|
{1,2,1,64,128,64,0}, // tiny: B,Hq,Hk,q,kv,D,causal
|
||||||
|
|
@ -94,18 +195,8 @@ int main() {
|
||||||
p.scale=1.0f/sqrtf((float)D);
|
p.scale=1.0f/sqrtf((float)D);
|
||||||
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
|
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
|
||||||
|
|
||||||
constexpr int G=8, ROWS=32, P_BC=32;
|
|
||||||
dim3 grid((ql+ROWS-1)/ROWS, Hq, B);
|
|
||||||
dim3 block(G, ROWS, 1);
|
|
||||||
printf("grid=(%d,%d,%d) block=(%d,%d,%d)\n",
|
|
||||||
grid.x,grid.y,grid.z, block.x,block.y,block.z);
|
|
||||||
|
|
||||||
double t0=now_ms();
|
double t0=now_ms();
|
||||||
switch (D) {
|
dispatch_prefill(p);
|
||||||
case 64: gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<<grid,block>>>(p); break;
|
|
||||||
case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<<grid,block>>>(p); break;
|
|
||||||
default: printf("unsupported D=%d\n",D); return 1;
|
|
||||||
}
|
|
||||||
cudaDeviceSynchronize();
|
cudaDeviceSynchronize();
|
||||||
double kms=now_ms()-t0;
|
double kms=now_ms()-t0;
|
||||||
cudaError_t err=cudaGetLastError();
|
cudaError_t err=cudaGetLastError();
|
||||||
|
|
@ -128,5 +219,6 @@ int main() {
|
||||||
delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
|
delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
|
||||||
}
|
}
|
||||||
printf("All tests passed!\n");
|
printf("All tests passed!\n");
|
||||||
|
bench();
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue