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 776f153e..2f2546cb 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 @@ -102,26 +102,31 @@ __global__ void remap_sf_to_cutlass_kernel( const cutlass::float_ue4m3_t* __restrict__ src, // (MN, K_sf) row-major cutlass::float_ue4m3_t* __restrict__ dst, // CUTLASS interleaved layout LayoutSF layout_sf, // CuTe layout for dst - int MN, int K_sf // Source dimensions + int MN, int K_sf, // Source dimensions (in SF groups) + int SFVecSize // Elements per SF group (16 for NVFP4) ) { // Iterate over SOURCE indices (row-major) and write to CUTLASS destination. - // The layout maps logical (m, k) -> CUTLASS linear index. + // The layout maps logical (m, k_elements) -> CUTLASS linear index. // This is the forward direction, which CuTe handles correctly. // + // IMPORTANT: The CuTe layout uses ELEMENT-SPACE for K, not group-space. + // So k_group=3 with SFVecSize=16 maps to element k=3*16=48 in the layout. + // // Previous approach (iterate over CUTLASS idx, reverse-map with idx2crd+flatten) // was broken: flatten() on the nested CuTe coordinate gives atom sub-indices, - // not logical (m, k). This caused K-group > 0 to always map to m*K_sf+0, - // losing all K-group information in the SFA. + // not logical (m, k). This caused all K-groups > 0 in SFA to map to m*K_sf+0, + // losing K-group information entirely. int src_idx = blockIdx.x * blockDim.x + threadIdx.x; int total = MN * K_sf; if (src_idx >= total) return; int m = src_idx / K_sf; - int k = src_idx % K_sf; + int k_sf = src_idx % K_sf; + int k_elem = k_sf * SFVecSize; // Convert to element-space for CuTe layout - // Use the CuTe layout to find the destination index for this (m, k) - // layout_sf(m, k) returns the linear index in CUTLASS's expected layout - auto dst_idx = layout_sf(cute::make_coord(m, k)); + // Use the CuTe layout to find the destination index for this (m, k_elem) + // layout_sf(m, k_elem) returns the linear index in CUTLASS's expected layout + auto dst_idx = layout_sf(cute::make_coord(m, k_elem)); if (dst_idx < cute::size(layout_sf)) { dst[dst_idx] = src[src_idx]; @@ -159,15 +164,23 @@ int cutlass_nvfp4_gemm_run( int sfb_size = cute::size(layout_SFB); int K_sf = K / InputSFVectorSize; - // Allocate CUTLASS-layout buffers and remap scales + // Allocate CUTLASS-layout buffers, zero-init, and remap scales + // Zero-init is critical: CUTLASS pads to tile boundaries (128x64), + // so dest is larger than M*K_sf. Unmapped slots must be zero to avoid + // garbage values in the GEMM. cutlass::device_memory::allocation sfa_cutlass(sfa_size); cutlass::device_memory::allocation sfb_cutlass(sfb_size); + cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); + cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); 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); - 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); + // Grid size based on SOURCE elements (M*K_sf), not CUTLASS buffer size + 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, InputSFVectorSize); + 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, InputSFVectorSize); typename Gemm::Arguments arguments { cutlass::gemm::GemmUniversalMode::kGemm,