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

fix: replace col_major_src with explicit source strides

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-15 21:23:21 +00:00
parent f6fd549800
commit 7285331395
1 changed files with 24 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -116,49 +116,37 @@ using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
// positions that aren't written are correct zeros.
/////////////////////////////////////////////////////////////////////////////////////////////////
template<typename LayoutSF>
template<typename LayoutSF>
__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<decltype(layout_sf.shape())>;
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<const ElementSF*>(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false);
static_cast<const ElementSF*>(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<const ElementSF*>(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, true);
static_cast<const ElementSF*>(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<const ElementSF*>(SFB_ptr),
static_cast<ElementSF*>(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<const ElementSF*>(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false);
static_cast<const ElementSF*>(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,