From 55c6903980403fa092b1902c97734c816126e42f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 22:02:55 +0000 Subject: [PATCH] D1: fix O rescale identity tensor - use PV MMA shape not QK shape --- dsv4/kernels/attention/fmha.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 5a2f4f94..f42693d4 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -368,7 +368,8 @@ class FmhaKernel: corr_tile_size = 16 n_corr_tiles = self.pv_n_tile // corr_tile_size if const_expr(self.n_kv_tiles > 1): - tOcO = pv_thr.partition_C(cS) + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = pv_thr.partition_C(cO) tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout)