From 24b9310682b2ff6a243b07bd1a0d5cbecce96c39 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 04:43:12 +0000 Subject: [PATCH] D1: Debug TMA partition shapes at hd=512 --- 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 2e91e554..e8afb4e5 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -181,6 +181,7 @@ class FmhaKernel: tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3)) tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3)) tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,0,None,0)]; tVgV = tVgV[(None,0,None,0)] + print(f"TMA: tAgQ shape={cute.shape(tAgQ)}, tBgK shape={cute.shape(tBgK)}, tVgV shape={cute.shape(tVgV)}") tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK) tCrV = pv_mma.make_fragment_B(sV)