From b5cd1b88c966fd6ab74dcd556a65c9d4c6db9ee3 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 24 May 2026 23:37:10 +0000
Subject: [PATCH] D2: add shape debug print for mQ/mK

---
 dsv4/kernels/attention/fmha.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 247e58c5..2812ca3a 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -189,6 +189,9 @@ class FmhaKernel:
         # D2: Multi-CTA grid. Use block_idx_y to select Q and O for this CTA's head.
         _bidx, head_cta_idx, _bidz = cute.arch.block_idx()  # grid=(1, num_ctas, 1)
 
+        # Print mQ shape at trace time to understand mode structure
+        print(f"D2 DEBUG: mQ shape={cute.shape(mQ)}, mK shape={cute.shape(mK)}")
+
         # Q: if num_ctas > 1, mQ has a head dimension. local_tile indexes into it.
         # K/V: shared (MQA), always coordinate 0.
         # For single-CTA (num_ctas=1), head_cta_idx=0 and coordinates are the same as before.