From 69f458ac306337a36cad6fddb4119ee1f1fdd882 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 09:25:48 +0000 Subject: [PATCH] Add debug prints to SMEM-P path to understand rank mismatch --- dsv4/kernels/attention/fmha.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 79af5205..4e371c5a 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -346,6 +346,13 @@ class FmhaKernel: # rP_bf16 contains P values in QK C-fragment layout (BF16) # Use rP_bf16 directly (already in correct layout for QK C-fragment) tSMEM_CPYrP = thr_smem_copy.partition_S(rP_bf16) + # DEBUG: Print shapes before copy + print(f"[SMEM-P DEBUG] rP_bf16 shape: {cute.shape(rP_bf16)}") + print(f"[SMEM-P DEBUG] tSMEM_CPYrP shape: {cute.shape(tSMEM_CPYrP)}") + print(f"[SMEM-P DEBUG] tSMEM_CPYsP shape: {cute.shape(tSMEM_CPYsP)}") + print(f"[SMEM-P DEBUG] rP_bf16 rank: {len(cute.shape(rP_bf16))}") + print(f"[SMEM-P DEBUG] tSMEM_CPYrP rank: {len(cute.shape(tSMEM_CPYrP))}") + print(f"[SMEM-P DEBUG] tSMEM_CPYsP rank: {len(cute.shape(tSMEM_CPYsP))}") cute.copy(tiled_smem_copy, tSMEM_CPYrP, tSMEM_CPYsP) cute.arch.fence_proxy("async.shared", space="cta") softmax_done_bar.arrive()