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

feat: TMA loads for both K and V in 6-warp FMHA kernel

Browse Source
This commit is contained in:
biondizzle
2026-05-29 19:34:48 +00:00
parent 523d3838a2
commit cd2c028b39
2 changed files with 34 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/unit/test_fmha_6warp_tma.cu
View File

@@ -82,11 +82,19 @@ int main() {
// TMA descriptor for K: (SK, HD) with tile (128, 16)
CUtensorMap tma_k; CUtensorMap* d_tma_k;
if (!create_tma_desc_2d_bf16(&tma_k, d_k, SK, HD, 128, 16)) {
printf("TMA desc FAILED\n"); return 1;
printf("TMA K desc FAILED\n"); return 1;
}
cudaMalloc(&d_tma_k, sizeof(CUtensorMap));
cudaMemcpy(d_tma_k, &tma_k, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
// TMA descriptor for V: (HD, SK) with tile (16, 16)
CUtensorMap tma_v; CUtensorMap* d_tma_v;
if (!create_tma_desc_2d_bf16(&tma_v, d_v, HD, SK, 16, 16)) {
printf("TMA V desc FAILED\n"); return 1;
}
cudaMalloc(&d_tma_v, sizeof(CUtensorMap));
cudaMemcpy(d_tma_v, &tma_v, sizeof(CUtensorMap), cudaMemcpyHostToDevice);
// Compute reference
float o_ref[HD];
{
@@ -114,7 +122,7 @@ int main() {
if (smem > 48 * 1024) {
cudaFuncSetAttribute(fmha_6warp_tma_kernel<HD>, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
}
fmha_6warp_tma_kernel<HD><<<1, 192, smem>>>(d_q, d_tma_k, d_v, d_o, d_lse, SK, SCALE);
fmha_6warp_tma_kernel<HD><<<1, 192, smem>>>(d_q, d_tma_k, d_tma_v, d_o, d_lse, SK, SCALE);
cudaError_t launch_err = cudaGetLastError();
if (launch_err != cudaSuccess) { printf("LAUNCH ERROR: %s\n", cudaGetErrorString(launch_err)); return 1; }
@@ -137,7 +145,7 @@ int main() {
printf("Filtered cosine: %.8f\n", cs);
printf("Test %s\n", cs > 0.999f ? "PASSED" : "FAILED");
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); cudaFree(d_tma_k);
cudaFree(d_q); cudaFree(d_k); cudaFree(d_v); cudaFree(d_o); cudaFree(d_lse); cudaFree(d_tma_k); cudaFree(d_tma_v);
free(h_q); free(h_k); free(h_v); free(h_o); free(h_lse);
return cs > 0.999f ? 0 : 1;
}