Rewrite CUTLASS kernel based on NVIDIA example 72b (nv_float4_t, CollectiveBuilder, OpClassBlockScaledTensorOp)
This commit is contained in:
@@ -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 <cuda_runtime.h>
|
||||
#include <cstdio>
|
||||
#include <type_traits>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cute/tensor.hpp>
|
||||
#include <cutlass/tensor_ref.h>
|
||||
#include <cutlass/gemm/dispatch_policy.hpp>
|
||||
#include <cutlass/gemm/collective/collective_builder.hpp>
|
||||
#include <cutlass/epilogue/collective/collective_builder.hpp>
|
||||
#include <cutlass/detail/sm100_blockscaled_layout.hpp>
|
||||
#include <cutlass/gemm/device/gemm_universal_adapter.h>
|
||||
#include <cutlass/gemm/kernel/gemm_universal.hpp>
|
||||
#include <cutlass/gemm/kernel/tile_scheduler_params.h>
|
||||
#include <cutlass/float_subbyte.h>
|
||||
|
||||
// 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<cutlass::float_e2m1_t>;
|
||||
using LayoutATag = cutlass::layout::RowMajor;
|
||||
constexpr int AlignmentA = 32;
|
||||
|
||||
// Stride types
|
||||
using StrideA = cutlass::gemm::TagToStrideA_t<LayoutA>;
|
||||
using StrideB = cutlass::gemm::TagToStrideB_t<LayoutB>;
|
||||
using StrideC = cutlass::gemm::TagToStrideC_t<LayoutC>;
|
||||
using StrideD = cutlass::gemm::TagToStrideD_t<LayoutD>;
|
||||
// B matrix configuration (weight, FP4 + UE4M3 block scales)
|
||||
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor; // K-major
|
||||
constexpr int AlignmentB = 32;
|
||||
|
||||
// Block scale factor strides — K-major
|
||||
using StrideSFA = Stride<int64_t, _1, int64_t>;
|
||||
using StrideSFB = Stride<int64_t, _1, int64_t>;
|
||||
// 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<ElementD>::value;
|
||||
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::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<static_cast<int>(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<SFVecSize>::deduce_smem_layoutSFA(
|
||||
TiledMma{}, TileShape{}));
|
||||
using SmemLayoutAtomSFB = decltype(
|
||||
cutlass::detail::Sm1xxBlockScaledConfig<SFVecSize>::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<ElementA, ElementSF>, // ElementPairA
|
||||
cute::tuple<StrideA, StrideSFA>, // StridePairA
|
||||
cute::tuple<ElementB, ElementSF>, // ElementPairB
|
||||
cute::tuple<StrideB, StrideSFB>, // StridePairB
|
||||
TiledMma,
|
||||
cute::tuple<SM90_TMA_LOAD, SM90_TMA_LOAD>, // GmemTiledCopyPairA
|
||||
cute::tuple<SmemLayoutAtomA, SmemLayoutAtomSFA>, // SmemLayoutAtomPairA
|
||||
void, // SmemCopyAtomA (void for UMMA)
|
||||
cute::identity, // TransformA
|
||||
cute::tuple<SM90_TMA_LOAD, SM90_TMA_LOAD>, // GmemTiledCopyPairB
|
||||
cute::tuple<SmemLayoutAtomB, SmemLayoutAtomSFB>, // 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<int, int, int, int>,
|
||||
CollectiveOp,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue,
|
||||
void>;
|
||||
|
||||
using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
// 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<SFVecSize>::tile_atom_to_shape_SFA(problem_shape);
|
||||
auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig<SFVecSize>::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<ElementA const*>(A_packed);
|
||||
args.mainloop.dA = dA;
|
||||
args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(B_packed);
|
||||
args.mainloop.dB = dB;
|
||||
args.mainloop.ptr_SFA = reinterpret_cast<ElementSF const*>(SFA);
|
||||
args.mainloop.layout_SFA = layout_SFA;
|
||||
args.mainloop.ptr_SFB = reinterpret_cast<ElementSF const*>(SFB);
|
||||
args.mainloop.layout_SFB = layout_SFB;
|
||||
size_t workspace_size = Gemm::get_workspace_size(arguments);
|
||||
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
|
||||
|
||||
// Epilogue: C = 1*D + 0*C
|
||||
args.epilogue.ptr_C = nullptr;
|
||||
args.epilogue.dC = dC;
|
||||
args.epilogue.ptr_D = reinterpret_cast<ElementD*>(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
|
||||
|
||||
@@ -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 <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
#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<LayoutA>;
|
||||
using StrideB = cutlass::gemm::TagToStrideB_t<LayoutB>;
|
||||
using StrideC = cutlass::gemm::TagToStrideC_t<LayoutC>;
|
||||
using StrideD = cutlass::gemm::TagToStrideD_t<LayoutD>;
|
||||
using StrideSFA = Stride<int64_t, _1, int64_t>;
|
||||
using StrideSFB = Stride<int64_t, _1, int64_t>;
|
||||
|
||||
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<ElementA, ElementSF>;
|
||||
using ElementPairB = cute::tuple<ElementB, ElementSF>;
|
||||
using StridePairA = cute::tuple<StrideA, StrideSFA>;
|
||||
using StridePairB = cute::tuple<StrideB, StrideSFB>;
|
||||
|
||||
// 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<int,int,int,int>,
|
||||
CollectiveMainloop,
|
||||
CollectiveEpilogue
|
||||
>;
|
||||
|
||||
using GemmDevice = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// 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<int>(M), static_cast<int>(N), static_cast<int>(K), 1);
|
||||
auto layout_SFA = cutlass::detail::Sm1xxBlockScaledConfig<SFVecSize>::tile_atom_to_shape_SFA(problem_shape);
|
||||
auto layout_SFB = cutlass::detail::Sm1xxBlockScaledConfig<SFVecSize>::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<int>(M), static_cast<int>(N), static_cast<int>(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<int>(M), static_cast<int>(N), static_cast<int>(K), 1};
|
||||
|
||||
args.mainloop.ptr_A = reinterpret_cast<ElementA const*>(A_packed.data_ptr<int8_t>());
|
||||
args.mainloop.dA = dA;
|
||||
args.mainloop.ptr_B = reinterpret_cast<ElementB const*>(B_packed.data_ptr<int8_t>());
|
||||
args.mainloop.dB = dB;
|
||||
args.mainloop.ptr_SFA = reinterpret_cast<ElementSF const*>(SFA.data_ptr<uint8_t>());
|
||||
args.mainloop.layout_SFA = layout_SFA;
|
||||
args.mainloop.ptr_SFB = reinterpret_cast<ElementSF const*>(SFB.data_ptr<uint8_t>());
|
||||
args.mainloop.layout_SFB = layout_SFB;
|
||||
|
||||
args.epilogue.ptr_C = nullptr;
|
||||
args.epilogue.dC = dC;
|
||||
args.epilogue.ptr_D = reinterpret_cast<ElementD*>(C.data_ptr<cutlass::bfloat16_t>());
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user