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

revert: idx2crd remap approach — source-first needs hierarchical coords

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-15 11:44:38 +00:00
parent d5949a23b4
commit 56e62e916d
1 changed files with 45 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<typename LayoutSF>
__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<decltype(flat)>;
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<<<grid_sfa, block, 0, 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, false);
remap_sf_to_cutlass_kernel<<<grid_sfb, block, 0, stream>>>(
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, true);
typename Gemm::Arguments arguments {
@@ -230,9 +248,8 @@ extern "C" int cutlass_nvfp4_prepack_sfb_run(
cudaMemsetAsync(static_cast<ElementSF*>(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<<<grid, block, 0, stream>>>(
int block = 256;
remap_sf_to_cutlass_kernel<<<(sfb_size + block - 1) / block, block, 0, stream>>>(
static_cast<const ElementSF*>(SFB_ptr),
static_cast<ElementSF*>(SFB_cutlass_ptr),
layout_SFB,
@@ -274,9 +291,8 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb(
cutlass::device_memory::allocation<ElementSF> 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<<<grid_sfa, block, 0, 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, false);
typename Gemm::Arguments arguments {