From 7285331395b89f3268548d30c9fe8a6853a45d86 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 21:23:21 +0000 Subject: [PATCH] fix: replace col_major_src with explicit source strides SFA: src_stride_mn=K_sf, src_stride_ksf=1 (row-major M, K_sf) SFB: src_stride_mn=1, src_stride_ksf=N (row-major K_sf, N after transpose) Removes ambiguity about physical memory layout. The source indexing now uses mn*src_stride_mn + k_sf*src_stride_ksf which works for any contiguous or transposed layout. --- .../cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu | 57 ++++++++----------- 1 file changed, 24 insertions(+), 33 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu index dc7c25de..af37ada1 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu @@ -116,49 +116,37 @@ using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // positions that aren't written are correct zeros. ///////////////////////////////////////////////////////////////////////////////////////////////// +template template __global__ void remap_sf_to_cutlass_kernel( - const cutlass::float_ue4m3_t* __restrict__ src, // (MN, K_sf) or (K_sf, MN) row-major - cutlass::float_ue4m3_t* __restrict__ dst, // CUTLASS interleaved layout (zero-initialized) - LayoutSF layout_sf, // CuTe layout for dst - int MN, int K_sf, // Source dimensions (in SF groups) - bool col_major_src = false // true if source is (K_sf, MN) row-major + const cutlass::float_ue4m3_t* __restrict__ src, + cutlass::float_ue4m3_t* __restrict__ dst, + LayoutSF layout_sf, + int MN, + int K_sf, + int src_stride_mn, + int src_stride_ksf ) { - // Forward-mapping approach: iterate over source indices (mn, k_sf), - // compute the CUTLASS dst index via layout_sf forward mapping. - // k_sf * InputSFVectorSize converts from SF-group index to logical K element - // coordinate, which is what the layout expects. int tid = blockIdx.x * blockDim.x + threadIdx.x; int total = MN * K_sf; if (tid >= total) return; - int mn, k_sf_val, src_idx; + int mn = tid / K_sf; + int k_sf = tid % K_sf; - if (col_major_src) { - // Source is (K_sf, MN) row-major in memory — e.g. SFB after transpose+contiguous - k_sf_val = tid / MN; - mn = tid % MN; - src_idx = tid; // tid = k_sf_val * MN + mn - } else { - // Source is (MN, K_sf) row-major — e.g. SFA - mn = tid / K_sf; - k_sf_val = tid % K_sf; - src_idx = tid; // tid = mn * K_sf + k_sf_val - } - - // Use layout forward mapping: (mn, k_sf*16) -> dst_idx constexpr int LayoutRank = cute::rank_v; - int dst_idx = 0; + int dst_idx; if constexpr (LayoutRank == 2) { - dst_idx = layout_sf(cute::make_coord(mn, k_sf_val * InputSFVectorSize)); + dst_idx = layout_sf(cute::make_coord(mn, k_sf * InputSFVectorSize)); } else if constexpr (LayoutRank == 3) { - dst_idx = layout_sf(cute::make_coord(mn, k_sf_val * InputSFVectorSize, 0)); + dst_idx = layout_sf(cute::make_coord(mn, k_sf * InputSFVectorSize, 0)); + } else { + static_assert(LayoutRank == 2 || LayoutRank == 3, "Unexpected SF layout rank"); } - if (dst_idx >= 0 && dst_idx < cute::cosize(layout_sf)) { - dst[dst_idx] = src[src_idx]; - } + int src_idx = mn * src_stride_mn + k_sf * src_stride_ksf; + dst[dst_idx] = src[src_idx]; } ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -201,9 +189,11 @@ int cutlass_nvfp4_gemm_run( int sfa_src_total = M * K_sf; int sfb_src_total = N * K_sf; remap_sf_to_cutlass_kernel<<<(sfa_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); + static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, + M, K_sf, K_sf, 1); // SFA source: row-major (M, K_sf) remap_sf_to_cutlass_kernel<<<(sfb_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, true); + static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, + N, K_sf, 1, N); // SFB source: row-major (K_sf, N) after transpose typename Gemm::Arguments arguments { cutlass::gemm::GemmUniversalMode::kGemm, @@ -268,7 +258,7 @@ extern "C" int cutlass_nvfp4_prepack_sfb_run( static_cast(SFB_ptr), static_cast(SFB_cutlass_ptr), layout_SFB, - N, K_sf, true // SFB source is (K_sf, N) + N, K_sf, 1, N // SFB source: row-major (K_sf, N) after transpose ); return 0; @@ -309,7 +299,8 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( int block = 256; int sfa_src_total = M * K_sf; remap_sf_to_cutlass_kernel<<<(sfa_src_total + block - 1) / block, block, 0, stream>>>( - static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); + static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, + M, K_sf, K_sf, 1); // SFA source: row-major (M, K_sf) typename Gemm::Arguments arguments { cutlass::gemm::GemmUniversalMode::kGemm,