From 9eb193458e06964fe8692114cd633417b582a112 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 29 May 2026 19:43:41 +0000 Subject: [PATCH] test: refactored multi-row TMA test with multi-head and batch --- tests/unit/test_fmha_6warp_tma_multirow.cu | 161 +++++++++++---------- 1 file changed, 88 insertions(+), 73 deletions(-) diff --git a/tests/unit/test_fmha_6warp_tma_multirow.cu b/tests/unit/test_fmha_6warp_tma_multirow.cu index 1fc2d30c..aa3035b9 100644 --- a/tests/unit/test_fmha_6warp_tma_multirow.cu +++ b/tests/unit/test_fmha_6warp_tma_multirow.cu @@ -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 48*1024) + cudaFuncSetAttribute(fmha_6warp_tma_multirow_kernel, 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<<>>(params); - size_t smem = compute_smem(); - if (smem > 48*1024) - cudaFuncSetAttribute(fmha_6warp_tma_multirow_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); + cudaError_t err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + printf(" CUDA ERROR: %s\n", cudaGetErrorString(err)); + return 1; + } - fmha_6warp_tma_multirow_kernel<<<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 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;