From 7a74fac11fe57203afe50334e67e4cb8823fd80f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 09:34:50 +0000 Subject: [PATCH] Fix sP_2d definition for tSMEM_CPYsP --- dsv4/kernels/attention/fmha.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index b4657739..94007426 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -268,6 +268,7 @@ class FmhaKernel: ) tiled_smem_copy = cute.make_tiled_copy_C(smem_copy_atom, qk_mma) thr_smem_copy = tiled_smem_copy.get_slice(sfw_idx) + sP_2d = cute.group_modes(sP, 0, 3) tSMEM_CPYsP = thr_smem_copy.partition_D(sP_2d) # destination (SMEM) row_max = -Float32.inf