From 97ebb964a26f639054f003db0e934cd3094c1c7e Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 14:38:03 +0000
Subject: [PATCH] Move s_p_vals to dynamic SMEM

---
 tests/unit/test_fmha_v4.cu | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_fmha_v4.cu b/tests/unit/test_fmha_v4.cu
index 7f44d6ea..458e4811 100644
--- a/tests/unit/test_fmha_v4.cu
+++ b/tests/unit/test_fmha_v4.cu
@@ -87,7 +87,8 @@ test_fmha_smem_p(const bf16_t* __restrict__ q, const bf16_t* __restrict__ k,
     // ===== STEP 2: Softmax — read S from TMEM, keep P in registers =====
     // Warp 0 reads S, computes softmax, stores P values in registers
     // Then for each PV K-tile, warp 0 fills sPk from registers and signals
-    __shared__ float s_p_vals[SK];  // Share P values across warps
+    // s_p_vals lives in dynamic SMEM after sV
+    float* s_p_vals = (float*)(sV + NKT_PV * 256);
     if (wid == 0) {
         float s_vals[SK], row_max = -INFINITY;
         for (int n = 0; n < SK / 8; n++) {