biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: SF remap — element-space K coords + zero-init dest buffer

Browse Source 
Two fixes:
1. CuTe layout uses element-space K, not group-space. k_group=3 with
   SFVecSize=16 maps to k_elem=48 in the layout, not k=3.
   Added SFVecSize param to remap kernel, multiply k_sf * SFVecSize
   before passing to layout_sf().

2. Zero-init CUTLASS dest buffer before remap. The layout pads to
   tile boundaries (128x64), so dest is larger than M*K_sf. Unmapped
   padding slots reading garbage causes sporadic wrong results.
   Also fixed grid size to use source count (M*K_sf), not dest size.
This commit is contained in:
biondizzle
2026-05-14 14:54:18 +00:00
parent 5968ebad9f
commit d90967d6e9
1 changed files with 26 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu
View File

@@ -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<ElementSF> sfa_cutlass(sfa_size);
cutlass::device_memory::allocation<ElementSF> 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<const ElementSF*>(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<const ElementSF*>(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<const ElementSF*>(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<const ElementSF*>(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, InputSFVectorSize);
typename Gemm::Arguments arguments {
cutlass::gemm::GemmUniversalMode::kGemm,