Move s_p_vals to dynamic SMEM
This commit is contained in:
@@ -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++) {
|
||||
|
||||
Reference in New Issue
Block a user