From 8ca07722aede2e55549940a4ce104cd8254a7a3f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 20:02:32 +0000 Subject: [PATCH] SMEM-P: add shape debug prints --- 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 4039dbe5..9a58e436 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -324,7 +324,10 @@ class FmhaKernel: tTMEM_LOADcS_frg = cute.logical_divide(tTMEM_LOADcS, frg_layout) if self.use_smem_p: print(f"[SMEM-P CUTLASS] Created tTMEM_LOADcS_frg shape: {cute.shape(tTMEM_LOADcS_frg)}") + print(f"[SMEM-P CUTLASS] tTMEM_LOADrS shape: {cute.shape(tTMEM_LOADrS)}") + print(f"[SMEM-P CUTLASS] tTMEM_LOADcS shape: {cute.shape(tTMEM_LOADcS)}") print(f"[SMEM-P CUTLASS] frg_tile_size: {frg_tile_size}, frg_layout: {frg_layout}") + print(f"[SMEM-P CUTLASS] tTMEM_LOADrS_frg shape: {cute.shape(tTMEM_LOADrS_frg)}") for j in range(frg_cnt): for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])):