biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: refactored multi-row TMA test with multi-head and batch

Browse Source
This commit is contained in:
biondizzle
2026-05-29 19:43:41 +00:00
parent 832a04181d
commit 9eb193458e
1 changed files with 88 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

161
tests/unit/test_fmha_6warp_tma_multirow.cu
View File

@@ -34,13 +34,13 @@ static size_t compute_smem() {
size_t off = 0;
off += 4; off = (off+127)&~(size_t)127;
off += 16; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sTmaBuf
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sQ0
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sK0
off += TILE_SZ * 2; off = (off+127)&~(size_t)127; // sPk
off += 16 * MY_MMA_K * 2; // sV
off += 128 * 4; // sRowMax
off += 128 * 4; // sRowSum
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += TILE_SZ * 2; off = (off+127)&~(size_t)127;
off += 16 * MY_MMA_K * 2;
off += 128 * 4;
off += 128 * 4;
return off;
}
@@ -70,93 +70,108 @@ static void reference_attention(
}
}
int main() {
printf("=== 6-warp TMA FMHA multi-row HD=%d ===\n", HD);
static int test_single(int T, int n_h = 1, int batch = 1) {
printf("=== T=%d, n_h=%d, batch=%d, HD=%d ===\n", T, n_h, batch, HD);
const float SCALE = 1.0f / sqrtf((float)HD);
int total_heads = batch * n_h;
int total_fail = 0;
bf16_t* h_q = (bf16_t*)calloc(total_heads * MAX_T * HD, sizeof(bf16_t));
bf16_t* h_k = (bf16_t*)calloc(total_heads * SK * HD, sizeof(bf16_t));
bf16_t* h_v = (bf16_t*)calloc(total_heads * HD * SK, sizeof(bf16_t));
bf16_t* h_o = (bf16_t*)calloc(total_heads * MAX_T * HD, sizeof(bf16_t));
float* h_lse = (float*)calloc(total_heads * MAX_T, sizeof(float));
// Test T=1,4,32,128
for (int T : {1, 4, 32, 128}) {
printf("\n--- T=%d ---\n", T);
srand(42);
for (int h=0;h<total_heads;h++) {
for (int i=0;i<T*HD;i++) h_q[h*MAX_T*HD+i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
for (int i=0;i<SK*HD;i++) h_k[h*SK*HD+i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
for (int i=0;i<HD*SK;i++) h_v[h*HD*SK+i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
}
bf16_t* h_q = (bf16_t*)calloc(MAX_T * HD, sizeof(bf16_t));
bf16_t* h_k = (bf16_t*)calloc(SK * HD, sizeof(bf16_t));
bf16_t* h_v = (bf16_t*)calloc(HD * SK, sizeof(bf16_t));
bf16_t* h_o = (bf16_t*)calloc(MAX_T * HD, sizeof(bf16_t));
float* h_lse = (float*)calloc(MAX_T, sizeof(float));
bf16_t *d_q,*d_k,*d_v,*d_o; float *d_lse;
cudaMalloc(&d_q, total_heads*MAX_T*HD*sizeof(bf16_t));
cudaMalloc(&d_k, total_heads*SK*HD*sizeof(bf16_t));
cudaMalloc(&d_v, total_heads*HD*SK*sizeof(bf16_t));
cudaMalloc(&d_o, total_heads*MAX_T*HD*sizeof(bf16_t));
cudaMalloc(&d_lse, total_heads*MAX_T*sizeof(float));
cudaMemcpy(d_q, h_q, total_heads*MAX_T*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, total_heads*SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_v, h_v, total_heads*HD*SK*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemset(d_o, 0, total_heads*MAX_T*HD*sizeof(bf16_t));
cudaMemset(d_lse, 0, total_heads*MAX_T*sizeof(float));
srand(42);
for (int i=0;i<T*HD;i++) h_q[i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
for (int i=0;i<SK*HD;i++) h_k[i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
for (int i=0;i<HD*SK;i++) h_v[i] = f32_to_bf16_host((float)(rand()%100)/100.0f-0.5f);
// TMA descriptor for K: (SK, HD) — same for all heads (for now)
CUtensorMap tma_k; CUtensorMap* d_tma_k;
create_tma_desc_2d_bf16(&tma_k, d_k, SK, HD, 128, 16);
cudaMalloc(&d_tma_k, sizeof(CUtensorMap));
cudaMemcpy(d_tma_k, &tma_k, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
bf16_t *d_q,*d_k,*d_v,*d_o; float *d_lse;
cudaMalloc(&d_q, MAX_T*HD*sizeof(bf16_t));
cudaMalloc(&d_k, SK*HD*sizeof(bf16_t));
cudaMalloc(&d_v, HD*SK*sizeof(bf16_t));
cudaMalloc(&d_o, MAX_T*HD*sizeof(bf16_t));
cudaMalloc(&d_lse, MAX_T*sizeof(float));
cudaMemcpy(d_q, h_q, T*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_k, h_k, SK*HD*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemcpy(d_v, h_v, HD*SK*sizeof(bf16_t), cudaMemcpyHostToDevice);
cudaMemset(d_o, 0, MAX_T*HD*sizeof(bf16_t));
cudaMemset(d_lse, 0, MAX_T*sizeof(float));
FmhaTmaMultiRowParams params;
params.q = d_q; params.tma_k = d_tma_k; params.v = d_v;
params.o = d_o; params.lse = d_lse;
params.s_k = SK; params.T = T; params.scale = SCALE;
params.head_dim = HD;
params.q_head_stride = MAX_T*HD; params.q_batch_stride = n_h*MAX_T*HD;
params.k_head_stride = SK*HD; params.k_batch_stride = n_h*SK*HD;
params.v_head_stride = HD*SK; params.v_batch_stride = n_h*HD*SK;
params.o_head_stride = MAX_T*HD; params.o_batch_stride = n_h*MAX_T*HD;
params.lse_head_stride = MAX_T; params.lse_batch_stride = n_h*MAX_T;
// TMA descriptor for K
CUtensorMap tma_k; CUtensorMap* d_tma_k;
create_tma_desc_2d_bf16(&tma_k, d_k, SK, HD, 128, 16);
cudaMalloc(&d_tma_k, sizeof(CUtensorMap));
cudaMemcpy(d_tma_k, &tma_k, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
size_t smem = compute_smem();
if (smem > 48*1024)
cudaFuncSetAttribute(fmha_6warp_tma_multirow_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
FmhaTmaMultiRowParams params;
params.q = d_q; params.tma_k = d_tma_k; params.v = d_v;
params.o = d_o; params.lse = d_lse;
params.s_k = SK; params.T = T; params.scale = SCALE;
params.head_dim = HD;
params.q_head_stride = T*HD; params.q_batch_stride = T*HD;
params.k_head_stride = SK*HD; params.k_batch_stride = SK*HD;
params.v_head_stride = HD*SK; params.v_batch_stride = HD*SK;
params.o_head_stride = MAX_T*HD; params.o_batch_stride = MAX_T*HD;
params.lse_head_stride = MAX_T; params.lse_batch_stride = MAX_T;
dim3 grid(1, n_h, batch);
fmha_6warp_tma_multirow_kernel<HD><<<grid, 192, smem>>>(params);
size_t smem = compute_smem();
if (smem > 48*1024)
cudaFuncSetAttribute(fmha_6warp_tma_multirow_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) {
printf(" CUDA ERROR: %s\n", cudaGetErrorString(err));
return 1;
}
fmha_6warp_tma_multirow_kernel<HD><<<1, 192, smem>>>(params);
cudaMemcpy(h_o, d_o, total_heads*MAX_T*HD*sizeof(bf16_t), cudaMemcpyDeviceToHost);
cudaMemcpy(h_lse, d_lse, total_heads*MAX_T*sizeof(float), cudaMemcpyDeviceToHost);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) {
printf(" CUDA ERROR: %s\n", cudaGetErrorString(err));
total_fail++; continue;
}
cudaMemcpy(h_o, d_o, T*HD*sizeof(bf16_t), cudaMemcpyHostToHost);
cudaMemcpy(h_o, d_o, T*HD*sizeof(bf16_t), cudaMemcpyDeviceToHost);
cudaMemcpy(h_lse, d_lse, T*sizeof(float), cudaMemcpyDeviceToHost);
// Reference
// Check each head
int total_bad = 0;
float min_cos = 1.0f;
for (int h = 0; h < total_heads; h++) {
float* o_ref = (float*)calloc(T*HD, sizeof(float));
reference_attention(h_q, h_k, h_v, o_ref, nullptr, HD, T, SK, SCALE);
reference_attention(
h_q + h*MAX_T*HD, h_k + h*SK*HD, h_v + h*HD*SK,
o_ref, nullptr, HD, T, SK, SCALE);
// Check
float cs=0,na=0,nb=0; int bad=0;
float cs=0,na=0,nb=0;
for (int t=0;t<T;t++) {
for (int d=0;d<HD;d++) {
float a = bf16_to_f32_host(h_o[t*HD+d]), b = o_ref[t*HD+d];
float a = bf16_to_f32_host(h_o[h*MAX_T*HD + t*HD+d]), b = o_ref[t*HD+d];
if (fabsf(b) > 1e-4f) { cs+=a*b; na+=a*a; nb+=b*b; }
float rel = fabsf(b)>1e-4f ? fabsf(a-b)/fabsf(b) : fabsf(a-b);
if (rel > 0.05f) bad++;
}
}
cs /= (sqrtf(na)*sqrtf(nb)+1e-10f);
printf(" T=%d: cosine=%.8f bad=%d %s\n", T, cs, bad, bad==0&&cs>0.999f?"PASS":"FAIL");
if (cs < 0.999f) total_fail++;
if (cs < min_cos) min_cos = cs;
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); cudaFree(d_tma_k);
free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse); free(o_ref);
free(o_ref);
}
printf(" min_cos=%.8f %s\n", min_cos, min_cos>0.999f?"PASS":"FAIL");
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); cudaFree(d_tma_k);
free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse);
return min_cos > 0.999f ? 0 : 1;
}
int main() {
int total_fail = 0;
for (int T : {1, 4, 32, 128}) {
total_fail += test_single(T);
}
// Multi-head
total_fail += test_single(1, 4, 1); // n_h=4, T=1
total_fail += test_single(4, 2, 2); // n_h=2, batch=2, T=4
printf("\nOverall: %s\n", total_fail==0?"ALL PASSED":"SOME FAILED");
return total_fail == 0 ? 0 : 1;