diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu index c8a1620b..4f5a04c8 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu @@ -1,319 +1,156 @@ -/* - * CUTLASS NVFP4 Block-Scaled GEMM Kernel for DeepSeek-V4-Pro on Blackwell (SM100). +/*************************************************************************************************** + * CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro MoE * - * Uses CUTLASS 3.x GemmUniversalAdapter with - * MainloopSm100TmaUmmaWarpSpecializedBlockScaled dispatch policy. - * This invokes the native mxf8f6f4.block_scale tensor core instruction - * (tcgen05.mma) which performs E2M1 × E2M1 with UE4M3 block-16 scaling - * entirely in hardware — no dequantization step. - * - * Layout convention: - * A (activations): (M, K_packed) int8 K-major — packed E2M1, 2 vals/byte - * B (weights): (N, K_packed) int8 K-major — packed E2M1, 2 vals/byte - * SFA: (M, K_sf) float8_e4m3fn K-major — UE4M3 block16 scales - * SFB: (N, K_sf) float8_e4m3fn K-major — UE4M3 block16 scales - * C (output): (M, N) bfloat16 - * - * K_sf = K / 16 (one UE4M3 scale per group of 16 E2M1 elements) - * K_packed = K / 2 (two E2M1 values per int8 byte) - * - * Build with: - * nvcc -std=c++17 -arch=sm_100a \ - * -I/path/to/cutlass/include \ - * -I/path/to/cutlass/tools/util/include \ - * -I/path/to/cutlass/examples/common \ - * -DCUTLASS_ARCH_SM100_ENABLED=1 \ - * cutlass_nvfp4_gemm.cu -o cutlass_nvfp4_gemm \ - * -lcuda - */ + * Based on NVIDIA CUTLASS example 72b_blackwell_nvfp4_nvfp4_gemm.cu + * Uses native tcgen05.mma kind::mxf8f6f4.block_scale instructions on Blackwell SM100. + **************************************************************************************************/ #pragma once #include -#include -#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include -// CUTLASS 3.x includes -#include "cutlass/cutlass.h" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/util/packed_stride.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass/detail/sm100_blockscaled_layout.hpp" - -#include "cutlass/float_subbyte.h" -#include "cute/layout.hpp" +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) using namespace cute; -// ============================================================================ -// NVFP4 type definitions -// ============================================================================ -using ElementA = cutlass::float_e2m1_t; // Packed FP4 (2 per byte, K-major) -using ElementB = cutlass::float_e2m1_t; // Packed FP4 (2 per byte, K-major) -using ElementSF = cutlass::float_ue4m3_t; // UE4M3 block scale factor (group_size=16) -using ElementAccum = float; // Accumulator (float32) -using ElementC = cutlass::bfloat16_t; // Output type -using ElementD = cutlass::bfloat16_t; // Output type +///////////////////////////////////////////////////////////////////////////////////////////////// +// NVFP4 × NVFP4 → BF16 GEMM (for L1 and L2 MoE layers) +///////////////////////////////////////////////////////////////////////////////////////////////// -// Layout: K-major (ColumnMajor for A, RowMajor transposed interpretation for B) -using LayoutA = cutlass::layout::ColumnMajor; -using LayoutB = cutlass::layout::ColumnMajor; -using LayoutC = cutlass::layout::RowMajor; -using LayoutD = cutlass::layout::RowMajor; +// A matrix configuration (activation, FP4 + UE4M3 block scales) +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 32; -// Stride types -using StrideA = cutlass::gemm::TagToStrideA_t; -using StrideB = cutlass::gemm::TagToStrideB_t; -using StrideC = cutlass::gemm::TagToStrideC_t; -using StrideD = cutlass::gemm::TagToStrideD_t; +// B matrix configuration (weight, FP4 + UE4M3 block scales) +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; // K-major +constexpr int AlignmentB = 32; -// Block scale factor strides — K-major -using StrideSFA = Stride; -using StrideSFB = Stride; +// C/D matrix configuration (BF16 output) +using ElementD = cutlass::bfloat16_t; +using ElementC = float; // Epilogue compute type +using LayoutCTag = cutlass::layout::RowMajor; +using LayoutDTag = cutlass::layout::RowMajor; +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; +constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; -// ============================================================================ -// GEMM configuration for DeepSeek-V4-Pro -// ============================================================================ -// Tile shape: chosen to match UMMA atom requirements for f8f6f4.block_scale -// SFVecSize = 16 (one UE4M3 scale per 16 E2M1 elements) -// CTA_N must be one of 64/128/192/256 per CUTLASS requirement -// ============================================================================ -constexpr int SFVecSize = 16; // NVFP4 block group size +// Kernel functional config +using ElementAccumulator = float; +using ElementCompute = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; -// Tile shape: 128x128x64 (M, N, K) where K is in E2M1 elements -using TileShape = Shape<_128, _128, _64>; -using ClusterShape = Shape<_1, _1, _1>; -constexpr int Stages = 2; +// Kernel perf config +using MmaTileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; -// ============================================================================ -// Tiled MMA: Use UMMA atom for mxf8f6f4.block_scale -// ============================================================================ -// The MMA atom for block-scaled f8f6f4 on SM100 is: -// UMMA::rs64x128x16tb0x0_base_op_C, kind=mxf8f6f4.block_scale -// We use TiledMma that's compatible with the dispatch policy. -// The CollectiveMma specialization for BlockScaled will deduce the TiledMma -// from the dispatch policy and tile shape. -// ============================================================================ +constexpr int InputSFVectorSize = 16; // UE4M3 group_size = 16 -// ============================================================================ -// Smem layout atoms — must match UMMA requirements -// ============================================================================ -// For SM100 UMMA with FP4 (4-bit) elements: -// SmemLayoutAtomA: (128, 16) in E2M1 elements (K-major tiled) -// SmemLayoutAtomB: (128, 16) in E2M1 elements (K-major tiled) -// ============================================================================ -using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::ss_smem_selector< - cutlass::gemm::collective::Sm100UmmaInterwarpsplit, - ElementA, cutlass::layout::ColumnMajor, TileShape, - decltype(cute::tile_size<0>(typename cutlass::gemm::collective::detail::ss_smem_layout_helper< - cutlass::gemm::collective::Sm100UmmaInterwarpsplit, - ElementA, cutlass::layout::ColumnMajor>::tiled_mma_op{})), - 2>{}()); +// Epilogue: standard LinComb (alpha * AB + beta * C) → BF16 output +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; -// Simplified: let the CollectiveBuilder deduce all smem layouts. -// We'll use the builder approach below. - -// ============================================================================ -// CUTLASS CollectiveBuilder approach (recommended for CUTLASS 3.5+) -// ============================================================================ -// The builder auto-deduces TiledMma, smem layouts, TMA copies, etc. -// based on arch, op type, element types, and tile shape. -// ============================================================================ - -#ifdef CUTLASS_ENABLE_COLLECTIVE_BUILDER - -using CollectiveOp = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm100, - cutlass::gemm::collective::OpBlockScaled, - ElementA, LayoutA, 2, // A: float_e2m1, K-major, alignment=2 - ElementB, LayoutB, 2, // B: float_e2m1, K-major, alignment=2 - ElementAccum, - TileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout<0>, +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, cutlass::gemm::collective::KernelScheduleAuto >::CollectiveOp; -#else - -// ============================================================================ -// Manual collective specialization (for CUTLASS < 3.5 or when builder -// doesn't support OpBlockScaled) -// ============================================================================ -// This directly uses the CollectiveMma specialization with -// MainloopSm100TmaUmmaWarpSpecializedBlockScaled dispatch policy. -// ============================================================================ - -#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" -#include "cutlass/epilogue/collective/sm100_epilogue.hpp" -#include "cutlass/pipeline/pipeline.hpp" - -// UMMA atom for block-scaled MMA -using TiledMma = decltype(cutlass::gemm::collective::detail::make_tiled_mma_sm100< - ElementA, LayoutA, ElementB, LayoutB, ElementAccum, TileShape>()); - -// Smem layout atoms (let CUTLASS deduce from TiledMma and TileShape) -using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::ss_smem_selector< - cutlass::gemm::collective::Sm100UmmaInterwarpsplit, - ElementA, LayoutA, TileShape, - decltype(cute::tile_size<0>(TiledMma{})), - Stages>()); - -using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::ss_smem_selector< - cutlass::gemm::collective::Sm100UmmaInterwarpsplit, - ElementB, LayoutB, TileShape, - decltype(cute::tile_size<1>(TiledMma{})), - Stages>()); - -// Smem layout atoms for scale factors -using SmemLayoutAtomSFA = decltype( - cutlass::detail::Sm1xxBlockScaledConfig::deduce_smem_layoutSFA( - TiledMma{}, TileShape{})); -using SmemLayoutAtomSFB = decltype( - cutlass::detail::Sm1xxBlockScaledConfig::deduce_smem_layoutSFB( - TiledMma{}, TileShape{})); - -using CollectiveOp = cutlass::gemm::collective::CollectiveMma< - cutlass::gemm::collective::MainloopSm100TmaUmmaWarpSpecializedBlockScaled< - Stages, 1, 1, ClusterShape, cutlass::arch::Sm100>, - TileShape, - cute::tuple, // ElementPairA - cute::tuple, // StridePairA - cute::tuple, // ElementPairB - cute::tuple, // StridePairB - TiledMma, - cute::tuple, // GmemTiledCopyPairA - cute::tuple, // SmemLayoutAtomPairA - void, // SmemCopyAtomA (void for UMMA) - cute::identity, // TransformA - cute::tuple, // GmemTiledCopyPairB - cute::tuple, // SmemLayoutAtomPairB - void, // SmemCopyAtomB (void for UMMA) - cute::identity // TransformB ->; - -#endif // CUTLASS_ENABLE_COLLECTIVE_BUILDER - -// ============================================================================ -// Epilogue -// ============================================================================ -using EpilogueOp = cutlass::epilogue::thread::LinearCombination< - ElementD, 1, ElementAccum, ElementAccum>; - -using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue< - cutlass::gemm::collective::EpilogueSm100TmaUmma<2, ClusterShape, cutlass::arch::Sm100>, - TileShape, - ElementAccum, - StrideC, - ElementC, - StrideD, - ElementD, - EpilogueOp ->; - -// ============================================================================ -// Gemm kernel and device adapter -// ============================================================================ using GemmKernel = cutlass::gemm::kernel::GemmUniversal< Shape, - CollectiveOp, - CollectiveEpilogue ->; + CollectiveMainloop, + CollectiveEpilogue, + void>; -using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter; +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +// Layout and stride types from the kernel +using StrideA = typename Gemm::GemmKernel::StrideA; +using StrideB = typename Gemm::GemmKernel::StrideB; +using StrideC = typename Gemm::GemmKernel::StrideC; +using StrideD = typename Gemm::GemmKernel::StrideD; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; +using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; -// ============================================================================ -// C API for PyTorch extension -// ============================================================================ +///////////////////////////////////////////////////////////////////////////////////////////////// +// C API for PyTorch binding +///////////////////////////////////////////////////////////////////////////////////////////////// extern "C" { -/** - * Run a single NVFP4 block-scaled GEMM: C = A @ B^T - * - * @param A_packed (M, K_packed) int8 — packed E2M1, 2 values per byte - * @param SFA (M, K_sf) uint8 — UE4M3 block16 scales - * @param B_packed (N, K_packed) int8 — packed E2M1, 2 values per byte - * @param SFB (N, K_sf) uint8 — UE4M3 block16 scales - * @param C_out (M, N) bfloat16 output - * @param M, N, K Problem dimensions (K in E2M1 elements, must be even) - * @param stream CUDA stream - * @return 0 on success, non-zero on error - */ +// Run a single NVFP4 × NVFP4 → BF16 GEMM +// A: (M, K/2) int8 packed E2M1, SFA: UE4M3 block scales, B: (N, K/2) int8 packed E2M1, SFB: UE4M3 block scales +// Output: (M, N) bfloat16 int cutlass_nvfp4_gemm_run( - const int8_t* A_packed, - const uint8_t* SFA, - const int8_t* B_packed, - const uint8_t* SFB, - __nv_bfloat16* C_out, - int M, int N, int K, + const void* A_ptr, const void* SFA_ptr, // A matrix + scale factors + const void* B_ptr, const void* SFB_ptr, // B matrix + scale factors + void* D_ptr, // Output D matrix (bfloat16) + int M, int N, int K, // Problem dimensions (K in E2M1 elements, so actual K = K_arg * 2) + float alpha, float beta, cudaStream_t stream ) { - int K_packed = K / 2; - int K_sf = K / SFVecSize; + // 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}); - // Compute scale factor layouts using Sm1xxBlockScaledConfig - auto problem_shape = cute::make_shape(M, N, K, 1); - auto layout_SFA = cutlass::detail::Sm1xxBlockScaledConfig::tile_atom_to_shape_SFA(problem_shape); - auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig::tile_atom_to_shape_SFB(problem_shape); + // 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)); - // Compute strides for A and B (K-major: contiguous in K, then MN) - // A is (M, K_packed) K-major: stride = (K_packed, 1) → ColumnMajor - StrideA dA = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); - StrideB dB = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); - StrideC dC = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); - StrideD dD = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); + 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 + }, + { // Epilogue arguments + { alpha, beta }, + nullptr, stride_C, // C matrix (not used, beta=0) + D_ptr, stride_D + } + }; - typename GemmDevice::Arguments args; - args.mode = cutlass::gemm::GemmUniversalMode::kGemm; - args.problem_shape = {M, N, K, 1}; + Gemm gemm; + CUTLASS_CHECK(gemm.can_implement(arguments)); - // Mainloop: paired (element, scale) inputs - args.mainloop.ptr_A = reinterpret_cast(A_packed); - args.mainloop.dA = dA; - args.mainloop.ptr_B = reinterpret_cast(B_packed); - args.mainloop.dB = dB; - args.mainloop.ptr_SFA = reinterpret_cast(SFA); - args.mainloop.layout_SFA = layout_SFA; - args.mainloop.ptr_SFB = reinterpret_cast(SFB); - args.mainloop.layout_SFB = layout_SFB; + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); - // Epilogue: C = 1*D + 0*C - args.epilogue.ptr_C = nullptr; - args.epilogue.dC = dC; - args.epilogue.ptr_D = reinterpret_cast(C_out); - args.epilogue.dD = dD; - args.epilogue.thread = {ElementAccum(1), ElementAccum(0)}; - - // Hardware info - args.hw_info.device_id = 0; // Will be set by caller - args.hw_info.sm_count = 0; - - GemmDevice gemm; - - auto can_impl = GemmDevice::can_implement(args); - if (can_impl != cutlass::Status::kSuccess) { - fprintf(stderr, "CUTLASS NVFP4 GEMM: can_implement returned false\n"); - return -1; - } - - auto status = gemm.initialize(args); - if (status != cutlass::Status::kSuccess) { - fprintf(stderr, "CUTLASS NVFP4 GEMM: initialize failed\n"); - return -2; - } - - status = gemm.run(stream); - if (status != cutlass::Status::kSuccess) { - fprintf(stderr, "CUTLASS NVFP4 GEMM: run failed\n"); - return -3; - } + CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); + CUTLASS_CHECK(gemm.run(stream)); return 0; } } // extern "C" + +#endif // CUTLASS_ARCH_MMA_SM100_SUPPORTED diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp index 93c31d6c..b4e33d33 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp @@ -1,222 +1,48 @@ -/* - * PyTorch CUDA extension binding for CUTLASS NVFP4 block-scaled GEMM. - * - * Build via setup.py or torch.utils.cpp_extension.load(). - * - * Requires: - * - CUDA 13.0+ (for SM100/Blackwell support) - * - CUTLASS headers (include path passed via extra_include_paths) - * - PyTorch with CUDA support - */ +/** PyTorch binding for CUTLASS NVFP4 block-scaled GEMM */ #include -#include #include -#include -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) -// CUTLASS includes -#include "cutlass/cutlass.h" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/util/packed_stride.hpp" -#include "cutlass/kernel_hardware_info.hpp" -#include "cutlass/detail/sm100_blockscaled_layout.hpp" - -#include "cutlass/float_subbyte.h" -#include "cute/layout.hpp" - -using namespace cute; - -// ============================================================================ -// NVFP4 types -// ============================================================================ -using ElementA = cutlass::float_e2m1_t; -using ElementB = cutlass::float_e2m1_t; -using ElementSF = cutlass::float_ue4m3_t; -using ElementAccum = float; -using ElementC = cutlass::bfloat16_t; -using ElementD = cutlass::bfloat16_t; - -using LayoutA = cutlass::layout::ColumnMajor; // K-major -using LayoutB = cutlass::layout::ColumnMajor; // K-major -using LayoutC = cutlass::layout::RowMajor; -using LayoutD = cutlass::layout::RowMajor; - -using StrideA = cutlass::gemm::TagToStrideA_t; -using StrideB = cutlass::gemm::TagToStrideB_t; -using StrideC = cutlass::gemm::TagToStrideC_t; -using StrideD = cutlass::gemm::TagToStrideD_t; -using StrideSFA = Stride; -using StrideSFB = Stride; - -constexpr int SFVecSize = 16; -constexpr int Stages = 2; - -// ============================================================================ -// Use CollectiveBuilder if available (CUTLASS 3.5+), otherwise manual -// ============================================================================ -#if defined(CUTLASS_ENABLE_COLLECTIVE_BUILDER) && CUTLASS_ENABLE_COLLECTIVE_BUILDER - -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" - -using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - cutlass::arch::Sm100, - cutlass::gemm::collective::OpBlockScaled, - ElementA, LayoutA, 2, - ElementB, LayoutB, 2, - ElementAccum, - Shape<_128, _128, _64>, - Shape<_1, _1, _1>, - cutlass::gemm::collective::StageCountAutoCarveout<0>, - cutlass::gemm::collective::KernelScheduleAuto ->::CollectiveOp; - -using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - cutlass::arch::Sm100, - cutlass::gemm::EpilogueDefault, - ElementAccum, - ElementC, LayoutC, 1, - ElementD, LayoutD, 1, - cutlass::gemm::collective::EpilogueScheduleAuto ->::CollectiveOp; - -#else - -// Manual specialization -#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" -#include "cutlass/epilogue/collective/sm100_epilogue.hpp" -#include "cutlass/pipeline/pipeline.hpp" - -// The TiledMma for block-scaled MMA is configured by the dispatch policy. -// We define minimal types and let the CollectiveMma specialization handle deduction. -using TileShape = Shape<_128, _128, _64>; -using ClusterShape = Shape<_1, _1, _1>; - -// Use the block-scaled dispatch policy -using DispatchPolicy = cutlass::gemm::collective::MainloopSm100TmaUmmaWarpSpecializedBlockScaled< - Stages, 1, 1, ClusterShape, cutlass::arch::Sm100>; - -// Element pairs: (element, scale_factor) -using ElementPairA = cute::tuple; -using ElementPairB = cute::tuple; -using StridePairA = cute::tuple; -using StridePairB = cute::tuple; - -// Smem layout atoms — we need to compute these based on the TiledMma. -// For now, use the CUTLASS internal deduction helpers. -// NOTE: The exact SmemLayoutAtom types depend on the TiledMma which is -// deduced inside CollectiveMma. For the PyTorch extension, we rely on -// the CollectiveBuilder path above. This manual path is provided as -// a fallback for older CUTLASS versions and may need adjustment. - -// For the manual path, we need to specify all template parameters explicitly. -// This is complex, so we provide a simplified version that uses the builder -// when available and falls back to a direct GEMM call otherwise. - -// Simplified: just include the header and let the linker resolve -using CollectiveMainloop = void; // Placeholder — use builder path -using CollectiveEpilogue = void; // Placeholder — use builder path +extern "C" int cutlass_nvfp4_gemm_run( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +); #endif -// Gemm kernel -using GemmKernel = cutlass::gemm::kernel::GemmUniversal< - Shape, - CollectiveMainloop, - CollectiveEpilogue ->; - -using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter; - - -// ============================================================================ -// PyTorch binding -// ============================================================================ - torch::Tensor cutlass_nvfp4_gemm_forward( - torch::Tensor A_packed, // (M, K//2) int8 — E2M1 packed - torch::Tensor SFA, // (M, K//16) uint8 — UE4M3 block16 scales (raw bytes) - torch::Tensor B_packed, // (N, K//2) int8 — E2M1 packed - torch::Tensor SFB, // (N, K//16) uint8 — UE4M3 block16 scales (raw bytes) - int64_t M, int64_t N, int64_t K + torch::Tensor A_packed, // (M, K_half) int8 packed E2M1 + torch::Tensor SFA, // UE4M3 block scales for A + torch::Tensor B_packed, // (N, K_half) int8 packed E2M1 + torch::Tensor SFB, // UE4M3 block scales for B + int64_t M, int64_t N, int64_t K // K in FP4 elements (actual K = K * 2 since 2 FP4 per byte) ) { - TORCH_CHECK(A_packed.is_cuda(), "A must be CUDA tensor"); - TORCH_CHECK(B_packed.is_cuda(), "B must be CUDA tensor"); - TORCH_CHECK(A_packed.dtype() == torch::kInt8, "A must be int8"); - TORCH_CHECK(B_packed.dtype() == torch::kInt8, "B must be int8"); +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + auto D = torch::empty({M, N}, torch::dtype(torch::kBF16).device(A_packed.device())); - auto options = torch::TensorOptions() - .dtype(torch::kBF16) - .device(A_packed.device()); - auto C = torch::zeros({M, N}, options); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - // Compute scale factor layouts - auto problem_shape = cute::make_shape( - static_cast(M), static_cast(N), static_cast(K), 1); - auto layout_SFA = cutlass::detail::Sm1xxBlockScaledConfig::tile_atom_to_shape_SFA(problem_shape); - auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig::tile_atom_to_shape_SFB(problem_shape); + cutlass_nvfp4_gemm_run( + A_packed.data_ptr(), SFA.data_ptr(), + B_packed.data_ptr(), SFB.data_ptr(), + D.data_ptr(), + static_cast(M), static_cast(N), static_cast(K), + 1.0f, 0.0f, + stream + ); - // K-major strides - StrideA dA = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); - StrideB dB = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); - StrideC dC = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1}); - StrideD dD = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); - - typename GemmDevice::Arguments args; - args.mode = cutlass::gemm::GemmUniversalMode::kGemm; - args.problem_shape = {static_cast(M), static_cast(N), static_cast(K), 1}; - - args.mainloop.ptr_A = reinterpret_cast(A_packed.data_ptr()); - args.mainloop.dA = dA; - args.mainloop.ptr_B = reinterpret_cast(B_packed.data_ptr()); - args.mainloop.dB = dB; - args.mainloop.ptr_SFA = reinterpret_cast(SFA.data_ptr()); - args.mainloop.layout_SFA = layout_SFA; - args.mainloop.ptr_SFB = reinterpret_cast(SFB.data_ptr()); - args.mainloop.layout_SFB = layout_SFB; - - args.epilogue.ptr_C = nullptr; - args.epilogue.dC = dC; - args.epilogue.ptr_D = reinterpret_cast(C.data_ptr()); - args.epilogue.dD = dD; - args.epilogue.thread = {ElementAccum(1), ElementAccum(0)}; - - int device_id = A_packed.get_device(); - args.hw_info.device_id = device_id; - args.hw_info.sm_count = 0; - - GemmDevice gemm; - - auto can_impl = GemmDevice::can_implement(args); - TORCH_CHECK(can_impl == cutlass::Status::kSuccess, - "CUTLASS NVFP4 GEMM: can_implement returned false"); - - auto status = gemm.initialize(args); - TORCH_CHECK(status == cutlass::Status::kSuccess, - "CUTLASS NVFP4 GEMM: initialize failed"); - - auto stream = c10::cuda::getCurrentCUDAStream(device_id); - status = gemm.run(stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, - "CUTLASS NVFP4 GEMM: run failed"); - - return C; + return D; +#else + TORCH_CHECK(false, "CUTLASS NVFP4 GEMM requires SM100 (Blackwell) architecture"); +#endif } - -// ============================================================================ -// Module bindings -// ============================================================================ - PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gemm_forward", &cutlass_nvfp4_gemm_forward, - "CUTLASS NVFP4 block-scaled GEMM forward (native SM100)"); + m.def("forward", &cutlass_nvfp4_gemm_forward, "CUTLASS NVFP4 block-scaled GEMM forward"); } - -#pragma GCC diagnostic pop