diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 0585724a..86829eed 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -288,6 +288,8 @@ class FmhaKernel: # sP_addr = 64*m + (k%16) + 16*((k//16)%4) + 8192*(k//64) # This is representable as a CuTe layout: (128, (16, 4, 2)) -> (64, (1, 16, 8192)) _sP_nostage = sP[(None, None, None, 0)] # remove stage dim + print(f"SMEM-P DEBUG: sP_nostage shape={cute.shape(_sP_nostage)}, sP outer={cute.shape(sP)}") + print(f"SMEM-P DEBUG: pv_n_tile={self.pv_n_tile}, p_cols_fp32={p_cols_fp32}") row_max = -Float32.inf row_sum = Float32(0.0)