diff --git a/CMakeLists.txt b/CMakeLists.txt index 54374c2d01c0..fea927d75bfb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -979,7 +979,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu" "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu" + "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu" "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu") + if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + message(STATUS + "Building mxfp4_experts_quant unsupported stubs for SM12x because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).") + endif() set_gencode_flags_for_srcs( SRCS "${FP4_SM120_SRCS}" CUDA_ARCHS "${FP4_SM120_ARCHS}") @@ -991,8 +997,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.") endif() - # SM10x/11x FP4 kernels. MXFP4 experts quantization is currently compiled - # only in this block; SM12x has separate NVFP4 matmul/MoE kernels above. + # SM10x/11x FP4 kernels. The MXFP4 experts-quant and grouped-GEMM sources + # are shared with the SM12x block above: CUTLASS resolves the same + # block-scaled scale-factor layout for arch::Sm100 and arch::Sm120, so a + # single SF swizzle serves both families. if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(FP4_SM100_ARCHS "10.0f;10.7f;11.0f" "${CUDA_ARCHS}") else() diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu index e4d2f2201250..58fb3636908a 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu @@ -2,8 +2,9 @@ * SPDX-License-Identifier: Apache-2.0 * SPDX-FileCopyrightText: Copyright contributors to the vLLM project * - * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100. - * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups. + * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM10x/11x and + * SM12x. Uses Cutlass mx_float4_t operands, E8M0 block scales, and + * 32-element groups. */ #include @@ -166,8 +167,41 @@ void mxfp4_run_get_group_gemm_starts( } } -template -void run_mxfp4_blockwise_scaled_group_mm_sm100( +// Architecture-specific kernel configuration. +// +// SM100 (datacenter Blackwell) has a dedicated 1-SM MXFP4 Ptr-Array schedule +// and a 2-SM cluster capable block-scaled MMA. SM120 (GeForce/RTX PRO +// Blackwell) has no TMEM and no 2-SM block-scaled cluster, so it relies on +// the generic cooperative Ptr-Array block-scaled schedule that CUTLASS +// auto-selects from a pointer-typed StrideA. +template +struct Mxfp4GroupGemmArchConfig; + +#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 +template <> +struct Mxfp4GroupGemmArchConfig { + using ClusterShape = Shape<_1, _1, _1>; + using MmaTileShape = Shape<_128, _128, _128>; + using EpilogueTile = Shape<_128, _64>; + using KernelSchedule = + cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; +}; +#endif + +#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120 +template <> +struct Mxfp4GroupGemmArchConfig { + using ClusterShape = Shape<_1, _1, _1>; + using MmaTileShape = Shape<_128, _128, _128>; + using EpilogueTile = cutlass::epilogue::collective::EpilogueTileAuto; + using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; +}; +#endif + +template +void run_mxfp4_blockwise_scaled_group_mm_impl( torch::stable::Tensor& output, const torch::stable::Tensor& a, const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, const torch::stable::Tensor& b_blockscales, @@ -198,35 +232,31 @@ void run_mxfp4_blockwise_scaled_group_mm_sm100( static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; // Architecture definitions - using ArchTag = cutlass::arch::Sm100; + using ArchTag = Arch; using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; using StageCountType = cutlass::gemm::collective::StageCountAuto; - using ClusterShape = Shape<_1, _1, _1>; - struct MMA1SMConfig { - using MmaTileShape = Shape<_128, _128, _128>; - using KernelSchedule = - cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; - }; + using ArchConfig = Mxfp4GroupGemmArchConfig; + using ClusterShape = typename ArchConfig::ClusterShape; + using MmaTileShape = typename ArchConfig::MmaTileShape; using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape, - ClusterShape, Shape<_128, _64>, ElementAccumulator, + ArchTag, EpilogueOperatorClass, MmaTileShape, ClusterShape, + typename ArchConfig::EpilogueTile, ElementAccumulator, ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, LayoutC*, AlignmentD, - typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; + typename ArchConfig::EpilogueSchedule>::CollectiveOp; using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA, - ElementB, LayoutB*, AlignmentB, ElementAccumulator, - typename MMA1SMConfig::MmaTileShape, ClusterShape, + ElementB, LayoutB*, AlignmentB, ElementAccumulator, MmaTileShape, + ClusterShape, cutlass::gemm::collective::StageCountAutoCarveout( sizeof(typename CollectiveEpilogue::SharedStorage))>, - typename MMA1SMConfig::KernelSchedule>::CollectiveOp; + typename ArchConfig::KernelSchedule>::CollectiveOp; using GemmKernel = cutlass::gemm::kernel::GemmUniversal= 100 && version_num < 120) { - run_mxfp4_blockwise_scaled_group_mm_sm100( + run_mxfp4_blockwise_scaled_group_mm_impl( + output, a, b, a_blockscale, b_blockscales, problem_sizes, + expert_offsets, sf_offsets, M, N, K); + return; + } +#endif +#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120 + if (version_num >= 120 && version_num < 130) { + run_mxfp4_blockwise_scaled_group_mm_impl( output, a, b, a_blockscale, b_blockscales, problem_sizes, expert_offsets, sf_offsets, M, N, K); return; @@ -388,10 +426,13 @@ void run_mxfp4_blockwise_scaled_group_mm( STD_TORCH_CHECK_NOT_IMPLEMENTED( false, "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ", - version_num, ". Required capability: 100"); + version_num, + ". Required capability: 100-119 (SM10x/11x) or 120-129 " + "(SM12x)."); } -#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 +#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \ + (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120) constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte; // E8M0 scale factors stored as uint8 constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte; @@ -417,7 +458,8 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, const torch::stable::Tensor& problem_sizes, const torch::stable::Tensor& expert_offsets, const torch::stable::Tensor& sf_offsets) { -#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 +#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \ + (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120) // Input validation CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a"); CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b"); @@ -463,7 +505,8 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, STD_TORCH_CHECK_NOT_IMPLEMENTED( false, "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with " - "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+."); + "SM10x/11x (ENABLE_NVFP4_SM100) or SM12x (ENABLE_NVFP4_SM120) " + "block-scaled FP4 MoE and CUDA 12.8+."); #endif } diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 20f024bcef51..219cd3c98e7f 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -377,7 +377,11 @@ static void validate_mxfp4_experts_quant_inputs( static bool mxfp4_experts_quant_sm_supported(int64_t cuda_device_capability) { #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED - return cuda_device_capability >= 100 && cuda_device_capability < 120; + // SM10x/11x and SM12x share this kernel: CUTLASS resolves the same + // Sm1xxBlkScaledConfig scale-factor layout for cutlass::arch::Sm100 and + // cutlass::arch::Sm120, so the swizzle emitted by + // cvt_quant_to_fp4_get_sf_out_offset() is valid for both families. + return cuda_device_capability >= 100 && cuda_device_capability < 130; #else return false; #endif @@ -391,9 +395,10 @@ void mxfp4_experts_quant( int64_t n_experts) { #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED int32_t sm = get_sm_version_num(); - STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), - "No compiled MXFP4 experts quant kernel for SM ", sm, - ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9."); + STD_TORCH_CHECK( + mxfp4_experts_quant_sm_supported(sm), + "No compiled MXFP4 experts quant kernel for SM ", sm, + ". Recompile with SM10x/11x or SM12x FP4 support and CUDA >= 12.9."); auto m_topk = input.size(0); auto k = input.size(1); @@ -429,9 +434,10 @@ void silu_and_mul_mxfp4_experts_quant( int64_t n_experts) { #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED int32_t sm = get_sm_version_num(); - STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), - "No compiled SiLU+Mul MXFP4 experts quant kernel for SM ", sm, - ". Recompile with SM10x/11x FP4 support and CUDA >= 12.9."); + STD_TORCH_CHECK( + mxfp4_experts_quant_sm_supported(sm), + "No compiled SiLU+Mul MXFP4 experts quant kernel for SM ", sm, + ". Recompile with SM10x/11x or SM12x FP4 support and CUDA >= 12.9."); auto m_topk = input.size(0); auto k_times_2 = input.size(1); diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py index 16b233b935e0..2d8192cb94f5 100644 --- a/tests/kernels/moe/test_mxfp4_moe.py +++ b/tests/kernels/moe/test_mxfp4_moe.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels.""" +"""Tests for CUTLASS MXFP4 x MXFP4 grouped MoE kernels (SM10x/11x and SM12x).""" import random @@ -30,9 +30,13 @@ def calc_diff(x, y): return 1 - sim -def is_sm100_supported() -> bool: - return current_platform.is_cuda() and current_platform.is_device_capability_family( - 100 +def is_mxfp4_moe_supported() -> bool: + """Ask the build which devices carry compiled MXFP4 MoE kernels.""" + capability = current_platform.get_device_capability() + return ( + current_platform.is_cuda() + and capability is not None + and ops.mxfp4_experts_quant_supported(capability.to_int()) ) @@ -61,8 +65,8 @@ def compute_ref_output( @pytest.mark.skipif( - not is_sm100_supported(), - reason="cutlass_mxfp4_group_mm requires CUDA SM100", + not is_mxfp4_moe_supported(), + reason="cutlass_mxfp4_group_mm requires a build with MXFP4 MoE kernels", ) @pytest.mark.parametrize("num_experts", [8, 16, 32]) @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) @@ -201,8 +205,8 @@ def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype): @pytest.mark.skipif( - not is_sm100_supported(), - reason="mxfp4_experts_quant requires CUDA SM100", + not is_mxfp4_moe_supported(), + reason="mxfp4_experts_quant requires a build with MXFP4 MoE kernels", ) def test_mxfp4_experts_quant_basic(): """ @@ -289,8 +293,8 @@ def compute_reference_e8m0_scale(block_max: float) -> int: @pytest.mark.skipif( - not is_sm100_supported(), - reason="mxfp4_experts_quant requires CUDA SM100", + not is_mxfp4_moe_supported(), + reason="mxfp4_experts_quant requires a build with MXFP4 MoE kernels", ) @pytest.mark.parametrize("k", [256, 7168]) @pytest.mark.parametrize("m", [16, 64]) @@ -412,8 +416,8 @@ def test_mxfp4_experts_quant_e8m0_scale_correctness(m, k): @pytest.mark.skipif( - not is_sm100_supported(), - reason="mxfp4_experts_quant requires CUDA SM100", + not is_mxfp4_moe_supported(), + reason="mxfp4_experts_quant requires a build with MXFP4 MoE kernels", ) def test_mxfp4_experts_quant_no_saturation(): """