From 34f5beb767b97a50b06406a002ee13c276bf469f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 23:44:39 +0000 Subject: [PATCH] D2: fix gC coordinate to match 5-mode flat_divide result --- 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 dc71a968..e0cb3555 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -205,7 +205,8 @@ class FmhaKernel: gQ = tC_gQ[None, None, head_cta_idx, None, None, None] c_tiler = cute.slice_(self.pv_mma_tiler, (None, None, 0)) # (128, hd) tC_gC = cute.flat_divide(mC, c_tiler) - gC = tC_gC[None, None, head_cta_idx, None, None, None] + print(f"D2 DEBUG: tC_gC shape={cute.shape(tC_gC)}") + gC = tC_gC[None, None, head_cta_idx, None, None] else: gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))