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

fix: use filter_zeros for SF allocation + no-branch forward mapping

Browse Source 
- Allocation: cute::size(cute::filter_zeros(layout)) matches CUTLASS examples
- Kernel: layout_sf(make_coord(mn, k_sf*16, 0)) — no branching on LayoutRank
- Avoids silent fallthrough that wrote dst[0] for all threads
This commit is contained in:
biondizzle
2026-05-15 22:58:51 +00:00
parent 6fc8fa61e0
commit 6626b75a2f
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -168,8 +168,8 @@ int cutlass_nvfp4_gemm_run(
using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
int sfa_size = cute::cosize(layout_SFA);
int sfb_size = cute::cosize(layout_SFB);
int sfa_size = cute::size(cute::filter_zeros(layout_SFA));
int sfb_size = cute::size(cute::filter_zeros(layout_SFB));
int K_sf = K / InputSFVectorSize;
cutlass::device_memory::allocation<ElementSF> sfa_cutlass(sfa_size);
@@ -225,7 +225,7 @@ extern "C" int cutlass_nvfp4_sfb_size(
) {
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
*out_size = cute::cosize(layout_SFB);
*out_size = cute::size(cute::filter_zeros(layout_SFB));
return 0;
}
@@ -239,7 +239,7 @@ extern "C" int cutlass_nvfp4_prepack_sfb_run(
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
int sfb_size = cute::cosize(layout_SFB);
int sfb_size = cute::size(cute::filter_zeros(layout_SFB));
int K_sf = K / InputSFVectorSize;
cudaMemsetAsync(static_cast<ElementSF*>(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream);
@@ -281,7 +281,7 @@ extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb(
using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
int sfa_size = cute::cosize(layout_SFA);
int sfa_size = cute::size(cute::filter_zeros(layout_SFA));
int K_sf = K / InputSFVectorSize;
// Only remap SFA (activation scales) — SFB is prepacked