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

Fix type casts and includes for CUTLASS NVFP4 GEMM

Browse Source
This commit is contained in:
biondizzle
2026-05-13 23:26:18 +00:00
parent 1eb9c43217
commit 3299d22ad6
1 changed files with 16 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,6 +19,8 @@
#include <cutlass/gemm/kernel/gemm_universal.hpp>
#include <cutlass/gemm/kernel/tile_scheduler_params.h>
#include <cutlass/float_subbyte.h>
#include <cutlass/util/packed_stride.hpp>
#include <cutlass/device_memory.h>
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
@@ -113,29 +115,34 @@ int cutlass_nvfp4_gemm_run(
cudaStream_t stream
) {
// Compute strides
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
// Compute scale factor layouts
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
// Cast void* to proper element types
using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF;
typename Gemm::Arguments arguments {
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{ // Mainloop arguments
A_ptr, stride_A,
B_ptr, stride_B,
SFA_ptr, layout_SFA,
SFB_ptr, layout_SFB
static_cast<const ArrayElementA*>(A_ptr), stride_A,
static_cast<const ArrayElementB*>(B_ptr), stride_B,
static_cast<const ElementSF*>(SFA_ptr), layout_SFA,
static_cast<const ElementSF*>(SFB_ptr), layout_SFB
},
{ // Epilogue arguments
{ alpha, beta },
nullptr, stride_C, // C matrix (not used, beta=0)
D_ptr, stride_D
static_cast<typename Gemm::GemmKernel::CollectiveEpilogue::ElementD*>(D_ptr), stride_D
}
};