From 10ae8f3346cccfa0f0d0c3ea46206977cff68981 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 30 May 2026 03:46:38 +0000
Subject: [PATCH] auto: pre-test commit

---
 tests/unit/test_fmha_6warp_tma.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unit/test_fmha_6warp_tma.cu b/tests/unit/test_fmha_6warp_tma.cu
index 7a62fc7c..3631c765 100644
--- a/tests/unit/test_fmha_6warp_tma.cu
+++ b/tests/unit/test_fmha_6warp_tma.cu
@@ -28,7 +28,7 @@ constexpr int SK = 128;
 constexpr int MY_MMA_K = 16;
 constexpr int TILE_SZ = 128 * MY_MMA_K;
 constexpr int V_SUB_SZ = 16 * MY_MMA_K;
-constexpr int TMEM_N = (HD <= 128) ? 128 : 256;
+constexpr int TMEM_N = (HD <= 128) ? 128 : (HD <= 256) ? 256 : 512;
 
 #include "dsv4/kernels/attention/fmha_6warp_tma.cuh"