From 9524b674ab0a72e032684b44d2a9b85814ef6a41 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 07:49:48 +0000
Subject: [PATCH] test: enable both reference + TMEM epilogue tests at
 hd=64/128

---
 tests/unit/test_fmha_sm100_standalone.cu | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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);