diff --git a/tests/unit/test_fmha_sm100_standalone.cu b/tests/unit/test_fmha_sm100_standalone.cu index 9cec4420..8a599730 100644 --- a/tests/unit/test_fmha_sm100_standalone.cu +++ b/tests/unit/test_fmha_sm100_standalone.cu @@ -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);