From 56e62e916d9002ade060482afaed1e92754e759a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 11:44:38 +0000 Subject: [PATCH] =?UTF-8?q?revert:=20idx2crd=20remap=20approach=20?= =?UTF-8?q?=E2=80=94=20source-first=20needs=20hierarchical=20coords?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cute::crd2idx requires hierarchical coordinates matching the layout's nested shape, which we don't have from flat (m, k_sf). Reverted to idx2crd dest-first approach. The real bug was cute::size vs cute::cosize for allocation, not the remap direction. --- .../cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu | 74 +++++++++++-------- 1 file changed, 45 insertions(+), 29 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 8b208e1f..da9196b2 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 @@ -96,36 +96,57 @@ using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; ///////////////////////////////////////////////////////////////////////////////////////////////// // Scale factor remap: source (row-major or col-major) -> CUTLASS interleaved layout // -// Source-first remap: iterates over the logical (m, k_sf) source grid (which is -// a simple 2D row-major tensor), and for each valid element, computes the -// physical CUTLASS dest index using the layout's forward mapping. +// Scale factor remap: source (row-major or col-major) -> CUTLASS interleaved layout // -// This is the inverse of the old idx2crd approach (which iterated dest → source -// and relied on fragile coordinate extraction from flattened hierarchical coords). -// Source-first is simpler, correct by construction, and avoids the idx2crd bugs. +// Iterates over CUTLASS dest indices, uses idx2crd to get the hierarchical coordinate, +// then extracts logical (m, k_sf) from the flattened result. // -// 2D kernel launch: blockDim = (32, 8), grid over (K_sf, MN). -// SFA: source is (MN, K_sf) row-major → src[m * K_sf + k_sf] -// SFB: source is (K_sf, MN) row-major → src[k_sf * MN + m] +// The flattened coordinate from idx2crd has nested structure. +// For SFA with Step<_2,_1> tiling, the layout shape is: +// ((32, 4, n_m_tiles), (16, 4, n_k_tiles)) +// Flattening gives: (inner_m, sub_m, tile_m, inner_k, sub_k, tile_k) +// where inner_m in [0,32), sub_m in [0,4), tile_m in [0, n_m_tiles) +// inner_k in [0,16) (within one SF group), sub_k in [0,4), tile_k in [0, n_k_tiles) +// +// Logical m = tile_m * 128 + inner_m * 4 + sub_m +// Logical k_sf = tile_k * 4 + sub_k (inner_k is within one SF group — same byte) +// +// NOTE: Allocation must use cute::cosize() (physical size including tile padding), +// not cute::size() (logical size). The dest buffer is zero-initialized so padding +// positions that aren't written are correct zeros. ///////////////////////////////////////////////////////////////////////////////////////////////// template __global__ void remap_sf_to_cutlass_kernel( - const cutlass::float_ue4m3_t* __restrict__ src, + const cutlass::float_ue4m3_t* __restrict__ src, // (MN, K_sf) or (K_sf, MN) depending on col_major_src 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 ) { - int k_sf = blockIdx.x * blockDim.x + threadIdx.x; - int m = blockIdx.y * blockDim.y + threadIdx.y; - if (m >= MN || k_sf >= K_sf) return; + int dst_idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = cute::size(layout_sf); + if (dst_idx >= total) return; - // Compute the CUTLASS physical index from logical (m, k_sf) via the layout - int dst_idx = cute::crd2idx(cute::make_coord(m, k_sf), layout_sf); + auto coord = cute::idx2crd(dst_idx, layout_sf.shape(), layout_sf.stride()); + auto flat = cute::flatten(coord); - // Read from source (row-major or col-major) and write to CUTLASS position - dst[dst_idx] = col_major_src ? src[k_sf * MN + m] : src[m * K_sf + k_sf]; + constexpr int R = cute::rank_v; + + int m = 0, k_sf = 0; + + if constexpr (R == 8) { + m = cute::get<0>(flat) + cute::get<1>(flat) * 32 + cute::get<2>(flat) * 128; + k_sf = cute::get<4>(flat) + cute::get<5>(flat) * 4; + } else { + m = 0; k_sf = 0; + } + + if (m < MN && k_sf < K_sf) { + // SFA: source is (MN, K_sf) row-major → src[m * K_sf + k_sf] + // SFB: source is (K_sf, MN) row-major → src[k_sf * MN + m] + dst[dst_idx] = col_major_src ? src[k_sf * MN + m] : src[m * K_sf + k_sf]; + } } ///////////////////////////////////////////////////////////////////////////////////////////////// @@ -164,13 +185,10 @@ int cutlass_nvfp4_gemm_run( cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); - dim3 block(32, 8); - dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y); - dim3 grid_sfb((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y); - - remap_sf_to_cutlass_kernel<<>>( + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>( static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); - remap_sf_to_cutlass_kernel<<>>( + remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>( static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, true); typename Gemm::Arguments arguments { @@ -230,9 +248,8 @@ extern "C" int cutlass_nvfp4_prepack_sfb_run( cudaMemsetAsync(static_cast(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream); - dim3 block(32, 8); - dim3 grid((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y); - remap_sf_to_cutlass_kernel<<>>( + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>( static_cast(SFB_ptr), static_cast(SFB_cutlass_ptr), layout_SFB, @@ -274,9 +291,8 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( cutlass::device_memory::allocation sfa_cutlass(sfa_size); cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); - dim3 block(32, 8); - dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y); - remap_sf_to_cutlass_kernel<<>>( + int block = 256; + remap_sf_to_cutlass_kernel<<<(sfa_size + block - 1) / block, block, 0, stream>>>( static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); typename Gemm::Arguments arguments {