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

fix: dispatch template HD at compile time

Browse Source
This commit is contained in:
biondizzle
2026-05-28 06:29:10 +00:00
parent bcc5d0b6cb
commit a73fb689f9
1 changed files with 19 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/unit/test_fmha_sm100_standalone.cu
View File

@@ -55,26 +55,34 @@ float cosine_sim(const float* a, const float* b, int n) {
return d > 0 ? dot/d : 0;
}
int test_kernel(const char* name, int HD, int sk, float scale,
int test_kernel(const char* name, int HD_val, int sk, float scale,
uint16_t* dq, uint16_t* dk, uint16_t* dv, uint16_t* do_gpu,
float* d_lse, float* ho_ref, int B, int H) {
dim3 grid(1, H, B);
dim3 block(NTHREADS);
int smem = (HD * sizeof(float)) + 128 + 1024; // Q + row_sums + slack
int smem = (HD_val * sizeof(float)) + 128 + 1024;
cudaMemset(do_gpu, 0, B*H*HD*sizeof(uint16_t));
cudaMemset(do_gpu, 0, B*H*HD_val*sizeof(uint16_t));
// Dispatch based on HD (template param must be compile-time constant)
#define DISPATCH(HD_T, KERNEL) \
KERNEL<HD_T><<<grid, block, smem>>>( \
dq, dk, dv, do_gpu, \
H*HD_T, sk*HD_T, H*HD_T, \
sk, 0, 0, scale, NULL, d_lse)
if (strcmp(name, "reference") == 0) {
fmha_decode_ref<HD><<<grid, block, smem>>>(
dq, dk, dv, do_gpu,
H*HD, sk*HD, H*HD,
sk, 0, 0, scale, NULL, d_lse);
if (HD_val == 64) DISPATCH(64, fmha_decode_ref);
else if (HD_val == 128) DISPATCH(128, fmha_decode_ref);
else if (HD_val == 256) DISPATCH(256, fmha_decode_ref);
else { printf(" ❌ unsupported hd=%d\n", HD_val); return 0; }
} else {
fmha_decode_tmem<HD><<<grid, block, smem>>>(
dq, dk, dv, do_gpu,
H*HD, sk*HD, H*HD,
sk, 0, 0, scale, NULL, d_lse);
if (HD_val == 64) DISPATCH(64, fmha_decode_tmem);
else if (HD_val == 128) DISPATCH(128, fmha_decode_tmem);
else if (HD_val == 256) DISPATCH(256, fmha_decode_tmem);
else { printf(" ❌ unsupported hd=%d\n", HD_val); return 0; }
}
#undef DISPATCH
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) {