test: enable both reference + TMEM epilogue tests at hd=64/128

This commit is contained in:
2026-05-28 07:49:48 +00:00
parent 446a0ca9fd
commit 9524b674ab

View File

@@ -129,7 +129,7 @@ int main() {
printf("=== FMHA SM100 Decode Kernel Test Suite ===\n\n");
int all_pass = 1;
int head_dims[] = {64};
int head_dims[] = {64, 128};
int s_ks[] = {128};
for (int t = 0; t < 2; t++) {
@@ -166,7 +166,7 @@ int main() {
cudaMemcpy(dk,hkb,B*sk*HD*2,cudaMemcpyHostToDevice);
cudaMemcpy(dv,hvb,B*HD*sk*2,cudaMemcpyHostToDevice);
// all_pass &= test_kernel("reference", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
all_pass &= test_kernel("reference", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
all_pass &= test_kernel("tmem_epilogue", HD, sk, scale, dq,dk,dv,do_,d_lse,ho_ref,B,H);
cudaFree(dq);cudaFree(dk);cudaFree(dv);cudaFree(do_);cudaFree(d_lse);