[Kernel] Enable CUTLASS MXFP4 W4A4 MoE on SM12x

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()

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 <torch/csrc/stable/library.h>
@@ -166,8 +167,41 @@ void mxfp4_run_get_group_gemm_starts(
   }
 }
 
-template <typename OutType>
-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 <typename Arch>
+struct Mxfp4GroupGemmArchConfig;
+
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+template <>
+struct Mxfp4GroupGemmArchConfig<cutlass::arch::Sm100> {
+  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<cutlass::arch::Sm120> {
+  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 <typename Arch, typename OutType>
+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<ElementD>::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<Arch>;
+  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<static_cast<int>(
               sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
+          typename ArchConfig::KernelSchedule>::CollectiveOp;
 
   using GemmKernel =
       cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
@@ -379,7 +409,15 @@ void run_mxfp4_blockwise_scaled_group_mm(
   int32_t version_num = get_sm_version_num();
 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
   if (version_num >= 100 && version_num < 120) {
-    run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
+    run_mxfp4_blockwise_scaled_group_mm_impl<cutlass::arch::Sm100, OutType>(
+        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<cutlass::arch::Sm120, OutType>(
         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
 }
 

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);

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():
     """