diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index fcd7ce93..897a0966 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -218,54 +218,10 @@ class FmhaKernel: cute.gemm(pv_mma, tOtO0, tOrP0[(None,None,kb)], tCrV[(None,None,kb,kvh.index)], tOtO0) pv_mma.set(tcgen05.Field.ACCUMULATE, True) else: - # SMEM-P: DEBUGGING RANK MISMATCH - print(f"[SMEM-P DEEP DEBUG] === START ===") - print(f"[SMEM-P DEEP DEBUG] rP_bf16 shape: {cute.shape(rP_bf16)}") - print(f"[SMEM-P DEEP DEBUG] rP_bf16 layout: {rP_bf16.layout}") - print(f"[SMEM-P DEEP DEBUG] sP shape: {cute.shape(sP)}") - print(f"[SMEM-P DEEP DEBUG] sP layout: {sP.layout}") - - # Try different approaches - print(f"[SMEM-P DEEP DEBUG] --- Approach 1: Direct copy ---") - tSMEM_CPYrP1 = thr_smem_copy.partition_S(rP_bf16) - tSMEM_CPYsP1 = thr_smem_copy.partition_D(sP) - print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP1 shape: {cute.shape(tSMEM_CPYrP1)}") - print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYsP1 shape: {cute.shape(tSMEM_CPYsP1)}") - - print(f"[SMEM-P DEEP DEBUG] --- Approach 2: With sP_2d ---") - sP_2d = cute.group_modes(sP, 0, 3) - print(f"[SMEM-P DEEP DEBUG] sP_2d shape: {cute.shape(sP_2d)}") - tSMEM_CPYrP2 = thr_smem_copy.partition_S(rP_bf16) - tSMEM_CPYsP2 = thr_smem_copy.partition_D(sP_2d) - print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP2 shape: {cute.shape(tSMEM_CPYrP2)}") - print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYsP2 shape: {cute.shape(tSMEM_CPYsP2)}") - - print(f"[SMEM-P DEEP DEBUG] --- Approach 3: Reshape rP_bf16 ---") - # Try to match destination rank - rP_bf16_reshaped = cute.group_modes(rP_bf16, 0, 2) # Group first 3 modes? - print(f"[SMEM-P DEEP DEBUG] rP_bf16_reshaped shape: {cute.shape(rP_bf16_reshaped)}") - tSMEM_CPYrP3 = thr_smem_copy.partition_S(rP_bf16_reshaped) - print(f"[SMEM-P DEEP DEBUG] tSMEM_CPYrP3 shape: {cute.shape(tSMEM_CPYrP3)}") - - # Try copy with approach that has matching rank - if len(cute.shape(tSMEM_CPYrP3)) == len(cute.shape(tSMEM_CPYsP2)): - print(f"[SMEM-P DEEP DEBUG] Rank match! Attempting copy...") - cute.copy(tiled_smem_copy, tSMEM_CPYrP3, tSMEM_CPYsP2) - cute.arch.fence_proxy("async.shared", space="cta") - print(f"[SMEM-P DEEP DEBUG] Copy succeeded with reshaped source") - elif len(cute.shape(tSMEM_CPYrP1)) == len(cute.shape(tSMEM_CPYsP1)): - print(f"[SMEM-P DEEP DEBUG] Direct partition ranks match! Attempting copy...") - cute.copy(tiled_smem_copy, tSMEM_CPYrP1, tSMEM_CPYsP1) - cute.arch.fence_proxy("async.shared", space="cta") - print(f"[SMEM-P DEEP DEBUG] Copy succeeded with direct partitions") - else: - print(f"[SMEM-P DEEP DEBUG] No rank match found. Falling back to stub.") - for j in cutlass.range(cute.size(sP), vectorize=True): - sP[j] = BFloat16(0.0) - cute.arch.fence_proxy("async.shared", space="cta") - - softmax_done_bar.arrive() - print(f"[SMEM-P DEEP DEBUG] === END ===") + # SMEM-P: PV reads P from SMEM + for kb in cutlass.range(cute.size(tCrP, mode=[2]), unroll_full=True): + cute.gemm(pv_mma, tOtO0, tCrP[(None,None,kb,0)], tCrV[(None,None,kb,kvh.index)], tOtO0) + pv_mma.set(tcgen05.Field.ACCUMULATE, True) cute.arch.fence_view_async_tmem_store() kvh.release() acc_pipe.producer_commit(acc_st); acc_st.advance() @@ -388,15 +344,39 @@ class FmhaKernel: cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) cute.arch.fence_view_async_tmem_store() else: - # SMEM-P: write P to SMEM via tiled_smem_copy - # 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))}") + else: + # SMEM-P: DEBUGGING RANK MISMATCH + print(f"[SMEM-P] rP_bf16 shape: {cute.shape(rP_bf16)} rank: {len(cute.shape(rP_bf16))}") + print(f"[SMEM-P] sP shape: {cute.shape(sP)} rank: {len(cute.shape(sP))}") + + # Approach 1: Try sP directly (no group_modes) + tSMEM_CPYrP1 = thr_smem_copy.partition_S(rP_bf16) + tSMEM_CPYsP1 = thr_smem_copy.partition_D(sP) + print(f"[SMEM-P] Approach1 - src shape: {cute.shape(tSMEM_CPYrP1)} rank: {len(cute.shape(tSMEM_CPYrP1))}") + print(f"[SMEM-P] Approach1 - dst shape: {cute.shape(tSMEM_CPYsP1)} rank: {len(cute.shape(tSMEM_CPYsP1))}") + + # Approach 2: Try with sP_2d + sP_2d = cute.group_modes(sP, 0, 3) + print(f"[SMEM-P] sP_2d shape: {cute.shape(sP_2d)} rank: {len(cute.shape(sP_2d))}") + tSMEM_CPYrP2 = thr_smem_copy.partition_S(rP_bf16) + tSMEM_CPYsP2 = thr_smem_copy.partition_D(sP_2d) + print(f"[SMEM-P] Approach2 - src shape: {cute.shape(tSMEM_CPYrP2)} rank: {len(cute.shape(tSMEM_CPYrP2))}") + print(f"[SMEM-P] Approach2 - dst shape: {cute.shape(tSMEM_CPYsP2)} rank: {len(cute.shape(tSMEM_CPYsP2))}") + + # Try copy with matching ranks + if len(cute.shape(tSMEM_CPYrP1)) == len(cute.shape(tSMEM_CPYsP1)): + print(f"[SMEM-P] Attempting copy with direct partitions") + cute.copy(tiled_smem_copy, tSMEM_CPYrP1, tSMEM_CPYsP1) + elif len(cute.shape(tSMEM_CPYrP2)) == len(cute.shape(tSMEM_CPYsP2)): + print(f"[SMEM-P] Attempting copy with sP_2d partitions") + cute.copy(tiled_smem_copy, tSMEM_CPYrP2, tSMEM_CPYsP2) + else: + print(f"[SMEM-P] No rank match. Using stub.") + for j in cutlass.range(cute.size(sP), vectorize=True): + sP[j] = BFloat16(0.0) + + cute.arch.fence_proxy("async.shared", space="cta") + softmax_done_bar.arrive() 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)