Move s_p_vals to dynamic SMEM

This commit is contained in:
2026-05-28 14:38:03 +00:00
parent d2387dd858
commit 97ebb964a2

View File

@@ -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++) {