A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 13:30 UTC
[Kernel] Add NVFP4 MoE CUTLASS support for SM120 (#29242)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Nov 25, 2025 at 09:59 UTC
e502098643b863f09dd9d8e74249523641f01298
8 files changed
+264
-30
CMakeLists.txt
+4
-1
@@ -604,12 +604,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
604
set(SRCS
605
"csrc/quantization/fp4/nvfp4_quant_kernels.cu"
606
"csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
607
- "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu")
607
+ "csrc/quantization/fp4/nvfp4_experts_quant.cu"
608
+ "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
609
+ "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
610
set_gencode_flags_for_srcs(
611
SRCS "${SRCS}"
612
CUDA_ARCHS "${FP4_ARCHS}")
613
list(APPEND VLLM_EXT_SRC "${SRCS}")
614
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
615
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
616
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
617
else()
618
message(STATUS "Not building NVFP4 as no compatible archs were found.")
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+224
-6
@@ -22,6 +22,7 @@
22
#include <ATen/cuda/CUDAContext.h>
23
#include <c10/cuda/CUDAGuard.h>
24
#include <c10/cuda/CUDAStream.h>
25
+#include "cutlass_extensions/common.hpp"
26
27
#include "cute/tensor.hpp"
28
#include "cutlass/tensor_ref.h"
@@ -173,7 +174,7 @@ void run_get_group_gemm_starts(
174
}
175
176
template <typename OutType>
176
-void run_fp4_blockwise_scaled_group_mm(
177
+void run_fp4_blockwise_scaled_group_mm_sm100(
178
torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
179
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
180
const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
@@ -343,17 +344,225 @@ void run_fp4_blockwise_scaled_group_mm(
344
345
auto can_implement_status = gemm_op.can_implement(args);
346
TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
346
- "Failed to implement GEMM");
347
+ "Failed to implement GEMM: status=", (int)can_implement_status);
348
+
349
+ // Run the GEMM
350
+ auto status = gemm_op.initialize(args, workspace.data_ptr());
351
+ TORCH_CHECK(status == cutlass::Status::kSuccess,
352
+ "Failed to initialize GEMM: status=", (int)status,
353
+ " workspace_size=", workspace_size, " num_experts=", num_experts,
354
+ " M=", M, " N=", N, " K=", K);
355
+
356
+ status = gemm_op.run(args, workspace.data_ptr(), stream);
357
+ TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
358
+}
359
+
360
+void run_fp4_blockwise_scaled_group_mm_sm120(
361
+ torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
362
+ const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
363
+ const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
364
+ const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
365
+ int N, int K) {
366
+ using ProblemShape =
367
+ cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
368
+ using ElementType = cutlass::float_e2m1_t;
369
+ using ElementSFType = cutlass::float_ue4m3_t;
370
+ using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
371
+ using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
372
+
373
+ // NOTE: For SM120 it seems templating the output type is not supported and
374
+ // we need to hardcode the output type to bfloat16
375
+ using ElementC = cutlass::bfloat16_t;
376
+ using ElementD = ElementC;
377
+ using ElementAccumulator = float;
378
+ // Layout definitions
379
+ using LayoutA = cutlass::layout::RowMajor;
380
+ using LayoutB = cutlass::layout::ColumnMajor;
381
+ using LayoutC = cutlass::layout::RowMajor;
382
+ using LayoutD = LayoutC;
383
+
384
+ // Alignment constraints
385
+ static constexpr int AlignmentA = 32;
386
+ static constexpr int AlignmentB = 32;
387
+ static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
388
+ static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
389
+
390
+ // Architecture definitions
391
+ using ArchTag = cutlass::arch::Sm120;
392
+ using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
393
+
394
+ using ClusterShape = Shape<_1, _1, _1>;
395
+ using MmaTileShape = Shape<_128, _128, _128>;
396
+
397
+ using FusionOperation = cutlass::epilogue::fusion::LinearCombination<
398
+ ElementD, ElementAccumulator, ElementC, ElementAccumulator>;
399
+
400
+ using CollectiveEpilogue =
401
+ typename cutlass::epilogue::collective::CollectiveBuilder<
402
+ ArchTag, OperatorClass, MmaTileShape, ClusterShape,
403
+ cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
404
+ ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
405
+ LayoutD*, AlignmentD,
406
+ cutlass::epilogue::collective::EpilogueScheduleAuto,
407
+ FusionOperation>::CollectiveOp;
408
+
409
+ using CollectiveMainloop =
410
+ typename cutlass::gemm::collective::CollectiveBuilder<
411
+ ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
412
+ LayoutB*, AlignmentB, ElementAccumulator, MmaTileShape, ClusterShape,
413
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
414
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
415
+ cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
416
+
417
+ using GemmKernel =
418
+ cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
419
+ CollectiveEpilogue>;
420
+
421
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
422
+ using StrideA = typename Gemm::GemmKernel::InternalStrideA;
423
+ using StrideB = typename Gemm::GemmKernel::InternalStrideB;
424
+ using StrideC = typename Gemm::GemmKernel::InternalStrideC;
425
+ using StrideD = typename Gemm::GemmKernel::InternalStrideD;
426
+
427
+ using LayoutSFA =
428
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
429
+ using LayoutSFB =
430
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
431
+ using ScaleConfig =
432
+ typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
433
+
434
+ using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
435
+ int num_experts = static_cast<int>(expert_offsets.size(0));
436
+ auto options_int =
437
+ torch::TensorOptions().dtype(torch::kInt64).device(a.device());
438
+
439
+ torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
440
+ torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
441
+ torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
442
+ torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
443
+ torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
444
+ torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
445
+ torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
446
+ torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
447
+ torch::Tensor c_strides1 =
448
+ torch::full({num_experts}, output.stride(0), options_int);
449
+ torch::Tensor a_strides1 =
450
+ torch::full({num_experts}, a.stride(0) * 2, options_int);
451
+ torch::Tensor b_strides1 =
452
+ torch::full({num_experts}, b.stride(1) * 2, options_int);
453
+
454
+ run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
455
+ a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
456
+ layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
457
+ expert_offsets, sf_offsets, problem_sizes, M, N, K);
458
+
459
+ // Create an instance of the GEMM
460
+ Gemm gemm_op;
461
+
462
+ // Initialize problem_sizes_as_shapes correctly
463
+ UnderlyingProblemShape* problem_sizes_as_shapes =
464
+ static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
465
+
466
+ // Set the Scheduler info
467
+ cutlass::KernelHardwareInfo hw_info;
468
+ using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
469
+ typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
470
+ scheduler.raster_order = RasterOrderOptions::AlongM;
471
+ hw_info.device_id = a.get_device();
472
+ static std::unordered_map<int, int> cached_sm_counts;
473
+ if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
474
+ cached_sm_counts[hw_info.device_id] =
475
+ cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
476
+ hw_info.device_id);
477
+ }
478
+ hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
479
+
480
+ // Mainloop Arguments
481
+ typename GemmKernel::MainloopArguments mainloop_args{
482
+ static_cast<const ElementType**>(a_ptrs.data_ptr()),
483
+ static_cast<StrideA*>(a_strides1.data_ptr()),
484
+ static_cast<const ElementType**>(b_ptrs.data_ptr()),
485
+ static_cast<StrideB*>(b_strides1.data_ptr()),
486
+ static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
487
+ reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
488
+ static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
489
+ reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
490
+
491
+ // Epilogue Arguments
492
+ typename GemmKernel::EpilogueArguments epilogue_args{
493
+ {}, // epilogue.thread
494
+ nullptr,
495
+ static_cast<StrideC*>(c_strides1.data_ptr()),
496
+ static_cast<ElementD**>(out_ptrs.data_ptr()),
497
+ static_cast<StrideC*>(c_strides1.data_ptr())};
498
+ auto& fusion_args = epilogue_args.thread;
499
+ fusion_args.alpha_ptr_array =
500
+ reinterpret_cast<float**>(alpha_ptrs.data_ptr());
501
+ fusion_args.dAlpha = {_0{}, _0{}, 1};
502
+ fusion_args.beta = 0.0f;
503
+
504
+ // Gemm Arguments
505
+ typename GemmKernel::Arguments args{
506
+ cutlass::gemm::GemmUniversalMode::kGrouped,
507
+ {num_experts, problem_sizes_as_shapes, nullptr},
508
+ mainloop_args,
509
+ epilogue_args,
510
+ hw_info,
511
+ scheduler};
512
+
513
+ size_t workspace_size = Gemm::get_workspace_size(args);
514
+ auto const workspace_options =
515
+ torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
516
+ auto workspace = torch::empty(workspace_size, workspace_options);
517
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
518
+
519
+ auto can_implement_status = gemm_op.can_implement(args);
520
+ TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
521
+ "Failed to implement GEMM: status=", (int)can_implement_status);
522
523
// Run the GEMM
524
auto status = gemm_op.initialize(args, workspace.data_ptr());
350
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
525
+ TORCH_CHECK(status == cutlass::Status::kSuccess,
526
+ "Failed to initialize GEMM: status=", (int)status,
527
+ " workspace_size=", workspace_size, " num_experts=", num_experts,
528
+ " M=", M, " N=", N, " K=", K);
529
530
status = gemm_op.run(args, workspace.data_ptr(), stream);
531
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
532
}
533
534
+template <typename OutType>
535
+void run_fp4_blockwise_scaled_group_mm(
536
+ torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
537
+ const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
538
+ const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
539
+ const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
540
+ int N, int K) {
541
+ int32_t version_num = get_sm_version_num();
542
+#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
543
+ if (version_num >= 120 && version_num < 130) {
544
+ run_fp4_blockwise_scaled_group_mm_sm120(
545
+ output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
546
+ expert_offsets, sf_offsets, M, N, K);
547
+ return;
548
+ }
549
+#endif
550
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
551
+ if (version_num >= 100 && version_num < 120) {
552
+ run_fp4_blockwise_scaled_group_mm_sm100<OutType>(
553
+ output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
554
+ expert_offsets, sf_offsets, M, N, K);
555
+ return;
556
+ }
557
+#endif
558
+ TORCH_CHECK_NOT_IMPLEMENTED(
559
+ false,
560
+ "No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ",
561
+ version_num, ". Required capability: 100 or 120");
562
+}
563
+
564
+#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
565
+ (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
566
constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
567
constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
568
#endif
@@ -374,7 +583,8 @@ void cutlass_fp4_group_mm(
583
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
584
const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
585
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
377
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
586
+#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
587
+ (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
588
// Input validation
589
CHECK_INPUT(a, FLOAT4_E2M1X2, "a");
590
CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
@@ -408,6 +618,14 @@ void cutlass_fp4_group_mm(
618
output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
619
expert_offsets, sf_offsets, M, N, K);
620
} else {
621
+ #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
622
+ int32_t version_num = get_sm_version_num();
623
+ if (version_num >= 120 && version_num < 130) {
624
+ TORCH_CHECK_NOT_IMPLEMENTED(
625
+ false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ",
626
+ output.scalar_type());
627
+ }
628
+ #endif
629
run_fp4_blockwise_scaled_group_mm<cutlass::half_t>(
630
output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
631
expert_offsets, sf_offsets, M, N, K);
@@ -416,8 +634,8 @@ void cutlass_fp4_group_mm(
634
TORCH_CHECK_NOT_IMPLEMENTED(
635
false,
636
"No compiled cutlass_fp4_group_mm kernel, vLLM must "
419
- "be compiled with ENABLE_NVFP4_SM100 for SM100+ and CUDA "
420
- "12.8 or above.");
637
+ "be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 "
638
+ "and CUDA 12.8 or above.");
639
#endif
640
}
641
csrc/quantization/fp4/nvfp4_experts_quant.cu
+1
-1
@@ -307,7 +307,7 @@ constexpr auto FLOAT = at::ScalarType::Float;
307
constexpr auto INT = at::ScalarType::Int;
308
constexpr auto UINT8 = at::ScalarType::Byte;
309
310
-void scaled_fp4_experts_quant_sm100a(
310
+void scaled_fp4_experts_quant_sm1xxa(
311
torch::Tensor& output, torch::Tensor& output_scale,
312
torch::Tensor const& input, torch::Tensor const& input_global_scale,
313
torch::Tensor const& input_offset_by_experts,
csrc/quantization/fp4/nvfp4_quant_entry.cu
+6
-4
@@ -24,8 +24,9 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
24
torch::Tensor const& input_sf);
25
#endif
26
27
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
28
-void scaled_fp4_experts_quant_sm100a(
27
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
28
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
29
+void scaled_fp4_experts_quant_sm1xxa(
30
torch::Tensor& output, torch::Tensor& output_scale,
31
torch::Tensor const& input, torch::Tensor const& input_global_scale,
32
torch::Tensor const& input_offset_by_experts,
@@ -54,8 +55,9 @@ void scaled_fp4_experts_quant(
55
torch::Tensor const& input, torch::Tensor const& input_global_scale,
56
torch::Tensor const& input_offset_by_experts,
57
torch::Tensor const& output_scale_offset_by_experts) {
57
-#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
58
- return scaled_fp4_experts_quant_sm100a(
58
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
59
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
60
+ return scaled_fp4_experts_quant_sm1xxa(
61
output, output_scale, input, input_global_scale, input_offset_by_experts,
62
output_scale_offset_by_experts);
63
#endif
csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu
+15
-12
@@ -67,9 +67,9 @@ void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
67
std::optional<torch::Tensor> const& bias);
68
#endif
69
70
-#if defined(ENABLE_SCALED_MM_SM90) && ENABLE_SCALED_MM_SM90 || \
71
- defined(ENABLE_SCALED_MM_SM100) && ENABLE_SCALED_MM_SM100 || \
72
- defined(ENABLE_SCALED_MM_SM120) && ENABLE_SCALED_MM_SM120
70
+#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
71
+ (defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
72
+ (defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
73
void get_cutlass_moe_mm_data_caller(
74
const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
75
torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
@@ -284,8 +284,9 @@ void get_cutlass_moe_mm_data(
284
// This function currently gets compiled only if we have a valid cutlass moe
285
// mm to run it for.
286
int32_t version_num = get_sm_version_num();
287
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
288
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
287
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
288
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
289
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
290
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
291
problem_sizes2, input_permutation,
292
output_permutation, num_experts, n, k,
@@ -296,7 +297,7 @@ void get_cutlass_moe_mm_data(
297
false,
298
"No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
299
"CUDA device capability: ",
299
- version_num, ". Required capability: 90 or 100");
300
+ version_num, ". Required capability: 90, 100, or 120");
301
}
302
303
void get_cutlass_moe_mm_problem_sizes(
@@ -304,8 +305,9 @@ void get_cutlass_moe_mm_problem_sizes(
305
torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
306
const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {
307
int32_t version_num = get_sm_version_num();
307
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
308
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
308
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
309
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
310
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
311
get_cutlass_moe_mm_problem_sizes_caller(topk_ids, problem_sizes1,
312
problem_sizes2, num_experts, n, k,
313
blockscale_offsets);
@@ -315,7 +317,7 @@ void get_cutlass_moe_mm_problem_sizes(
317
false,
318
"No compiled get_cutlass_moe_mm_problem_sizes: no cutlass_scaled_mm "
319
"kernel for CUDA device capability: ",
318
- version_num, ". Required capability: 90 or 100");
320
+ version_num, ". Required capability: 90, 100, or 120");
321
}
322
323
void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
@@ -328,8 +330,9 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
330
// This function currently gets compiled only if we have a valid cutlass moe
331
// mm to run it for.
332
int32_t version_num = get_sm_version_num();
331
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
332
- (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
333
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
334
+ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
335
+ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
336
get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1,
337
problem_sizes2, expert_num_tokens,
338
num_local_experts, padded_m, n, k);
@@ -339,7 +342,7 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
342
false,
343
"No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel "
344
"for CUDA device capability: ",
342
- version_num, ". Required capability: 90 or 100");
345
+ version_num, ". Required capability: 90, 100, or 120");
346
}
347
348
void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
docs/design/moe_kernel_features.md
+1
-1
@@ -60,7 +60,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
60
61
- [`ModelOptFp8MoEMethod`][vllm.model_executor.layers.quantization.modelopt.ModelOptFp8MoEMethod]
62
- [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]
63
-- [`CompressedTensorsW4A4MoeMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4MoeMethod]
63
+- [`CompressedTensorsW4A4Nvfp4MoeMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoeMethod]
64
- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW8A8Fp8MoEMethod]
65
- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
66
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+9
-5
@@ -103,7 +103,7 @@ __all__ = [
103
"CompressedTensorsW8A8Int8MoEMethod",
104
"CompressedTensorsWNA16MarlinMoEMethod",
105
"CompressedTensorsWNA16MoEMethod",
106
- "CompressedTensorsW4A4MoeMethod",
106
+ "CompressedTensorsW4A4Nvfp4MoeMethod",
107
"CompressedTensorsW4A8Int8MoEMethod",
108
]
109
@@ -171,7 +171,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
171
quant_config, layer.moe_config
172
)
173
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
174
- return CompressedTensorsW4A4MoeMethod(layer.moe_config)
174
+ return CompressedTensorsW4A4Nvfp4MoeMethod(layer.moe_config)
175
elif (
176
quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
177
or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
@@ -188,7 +188,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
188
)
189
190
191
-class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
191
+class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
192
def __init__(self, moe: FusedMoEConfig):
193
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
194
detect_nvfp4_moe_support,
@@ -205,8 +205,12 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
205
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
206
logger.info_once(
207
f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
208
- " for CompressedTensorsW4A4MoeMethod."
208
+ " for CompressedTensorsW4A4Nvfp4MoeMethod."
209
)
210
+ elif self.use_marlin:
211
+ logger.info_once("Using Marlin for CompressedTensorsW4A4Nvfp4MoeMethod.")
212
+ else:
213
+ logger.info_once("Using Cutlass for CompressedTensorsW4A4Nvfp4MoeMethod.")
214
215
def create_weights(
216
self,
@@ -612,7 +616,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
616
assert expert_map is None, (
617
"Expert Parallelism / expert_map "
618
"is currently not supported for "
615
- "CompressedTensorsW4A4MoeMethod."
619
+ "CompressedTensorsW4A4Nvfp4MoeMethod."
620
)
621
assert self.moe_quant_config is not None
622
vllm/model_executor/layers/quantization/modelopt.py
+4
@@ -1132,6 +1132,10 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1132
f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
1133
" for ModelOptNvFp4FusedMoE."
1134
)
1135
+ elif self.use_marlin:
1136
+ logger.info_once("Using Marlin for ModelOptNvFp4FusedMoE.")
1137
+ else:
1138
+ logger.info_once("Using Cutlass for ModelOptNvFp4FusedMoE.")
1139
1140
def maybe_make_prepare_finalize(
1141
self,