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 15, 01:00 UTC
[feat]: add SM100 support for cutlass FP8 groupGEMM (#20447)
Signed-off-by: Duncan Moss <djm.moss@gmail.com> Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Duncan Moss committed
Jul 22, 2025 at 07:27 UTC
2c8db17cfd6f2026549d7ac708be5109dbf7a2e8
8 files changed
+255
-32
CMakeLists.txt
+21
-1
@@ -577,7 +577,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
577
# if it's possible to compile MoE kernels that use its output.
578
cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
579
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
580
- set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu")
580
+ set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm90.cu")
581
set_gencode_flags_for_srcs(
582
SRCS "${SRCS}"
583
CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -595,6 +595,26 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
595
endif()
596
endif()
597
598
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a" "${CUDA_ARCHS}")
599
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
600
+ set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm100.cu")
601
+ set_gencode_flags_for_srcs(
602
+ SRCS "${SRCS}"
603
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
604
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
605
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
606
+ message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
607
+ else()
608
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
609
+ message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
610
+ "not >= 12.8, we recommend upgrading to CUDA 12.8 or later "
611
+ "if you intend on running FP8 quantized MoE models on Blackwell.")
612
+ else()
613
+ message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
614
+ "in CUDA target architectures.")
615
+ endif()
616
+ endif()
617
+
618
# moe_data.cu is used by all CUTLASS MoE kernels.
619
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a" "${CUDA_ARCHS}")
620
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cuh
+9
-4
@@ -18,7 +18,6 @@ using ProblemShape =
18
cutlass::gemm::GroupProblemShape<cute::Shape<int, int, int>>;
19
20
using ElementAccumulator = float;
21
-using ArchTag = cutlass::arch::Sm90;
21
using OperatorClass = cutlass::arch::OpClassTensorOp;
22
23
using LayoutA = cutlass::layout::RowMajor;
@@ -33,7 +32,7 @@ using LayoutD_Transpose =
32
using LayoutC = LayoutD;
33
using LayoutC_Transpose = LayoutD_Transpose;
34
36
-template <typename ElementAB_, typename ElementC_,
35
+template <typename ElementAB_, typename ElementC_, typename ArchTag_,
36
template <typename, typename, typename> typename Epilogue_,
37
typename TileShape, typename ClusterShape, typename KernelSchedule,
38
typename EpilogueSchedule, bool swap_ab_ = false>
@@ -43,6 +42,7 @@ struct cutlass_3x_group_gemm {
42
using ElementC = void;
43
using ElementD = ElementC_;
44
using ElementAccumulator = float;
45
+ using ArchTag = ArchTag_;
46
47
using Epilogue = Epilogue_<ElementAccumulator, ElementD, TileShape>;
48
@@ -77,7 +77,7 @@ struct cutlass_3x_group_gemm {
77
LayoutB*, AlignmentAB, ElementAccumulator, TileShape, ClusterShape,
78
Stages, KernelSchedule>::CollectiveOp>;
79
80
- using KernelType = enable_sm90_only<cutlass::gemm::kernel::GemmUniversal<
80
+ using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
81
ProblemShape, CollectiveMainloop, CollectiveEpilogue>>;
82
83
struct GemmKernel : public KernelType {};
@@ -156,9 +156,14 @@ void cutlass_group_gemm_caller(
156
static_cast<ElementD**>(out_ptrs.data_ptr()),
157
static_cast<StrideC*>(c_strides.data_ptr())};
158
159
+ int device_id = a_tensors.device().index();
160
+ static const cutlass::KernelHardwareInfo hw_info{
161
+ device_id, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
162
+ device_id)};
163
+
164
typename GemmKernel::Arguments args{
165
cutlass::gemm::GemmUniversalMode::kGrouped, prob_shape, mainloop_args,
161
- epilogue_args};
166
+ epilogue_args, hw_info};
167
168
using GemmOp = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
169
GemmOp gemm_op;
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm100.cu
new
+140
@@ -0,0 +1,140 @@
1
+#include <cudaTypedefs.h>
2
+
3
+#include <c10/cuda/CUDAGuard.h>
4
+#include <torch/all.h>
5
+
6
+#include "cutlass/cutlass.h"
7
+#include "grouped_mm_c3x.cuh"
8
+
9
+using namespace cute;
10
+
11
+namespace {
12
+
13
+template <typename InType, typename OutType,
14
+ template <typename, typename, typename> typename Epilogue>
15
+struct sm100_fp8_config_default {
16
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
17
+ using KernelSchedule =
18
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100;
19
+ using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
20
+ using TileShape = cute::Shape<cute::_128, cute::_256, cute::_128>;
21
+ using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
22
+ using ArchTag = cutlass::arch::Sm100;
23
+
24
+ using Cutlass3xGemm =
25
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
26
+ ClusterShape, KernelSchedule, EpilogueSchedule>;
27
+};
28
+
29
+template <typename InType, typename OutType,
30
+ template <typename, typename, typename> typename Epilogue>
31
+struct sm100_fp8_config_M64 {
32
+ // M in [1,64]
33
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
34
+ using KernelSchedule =
35
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmSm100;
36
+ using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
37
+ using TileShape = cute::Shape<cute::_128, cute::_16, cute::_128>;
38
+ using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
39
+ using ArchTag = cutlass::arch::Sm100;
40
+
41
+ using Cutlass3xGemm =
42
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
43
+ ClusterShape, KernelSchedule, EpilogueSchedule,
44
+ true>;
45
+};
46
+
47
+template <typename InType, typename OutType,
48
+ template <typename, typename, typename> typename Epilogue>
49
+struct sm100_fp8_config_N8192 {
50
+ // N in [8192, inf)
51
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
52
+ using KernelSchedule =
53
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmSm100;
54
+ using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm;
55
+ using TileShape = cute::Shape<cute::_128, cute::_256, cute::_128>;
56
+ using ClusterShape = cute::Shape<cute::_2, cute::_1, cute::_1>;
57
+ using ArchTag = cutlass::arch::Sm100;
58
+
59
+ using Cutlass3xGemm =
60
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
61
+ ClusterShape, KernelSchedule, EpilogueSchedule>;
62
+};
63
+
64
+template <typename InType, typename OutType>
65
+void run_cutlass_moe_mm_sm100(
66
+ torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
67
+ torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
68
+ torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
69
+ torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
70
+ torch::Tensor const& b_strides, torch::Tensor const& c_strides,
71
+ bool per_act_token, bool per_out_ch) {
72
+ TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
73
+ TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
74
+ TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
75
+
76
+ TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn,
77
+ "A tensors must be of type float8_e4m3fn.");
78
+ TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn,
79
+ "B tensors must be of type float8_e4m3fn.");
80
+
81
+ using Cutlass3xGemmDefault = typename sm100_fp8_config_default<
82
+ InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
83
+ using Cutlass3xGemmN8192 = typename sm100_fp8_config_N8192<
84
+ InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
85
+ using Cutlass3xGemmM64 = typename sm100_fp8_config_M64<
86
+ InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
87
+
88
+ uint32_t const m = a_tensors.size(0);
89
+ uint32_t const n = out_tensors.size(1);
90
+
91
+ if (m <= 64) {
92
+ cutlass_group_gemm_caller<Cutlass3xGemmM64>(
93
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
94
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
95
+ per_out_ch);
96
+ } else if (n >= 8192) {
97
+ cutlass_group_gemm_caller<Cutlass3xGemmN8192>(
98
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
99
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
100
+ per_out_ch);
101
+ } else {
102
+ cutlass_group_gemm_caller<Cutlass3xGemmDefault>(
103
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
104
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
105
+ per_out_ch);
106
+ }
107
+}
108
+} // namespace
109
+
110
+void dispatch_moe_mm_sm100(
111
+ torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
112
+ torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
113
+ torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
114
+ torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
115
+ torch::Tensor const& b_strides, torch::Tensor const& c_strides,
116
+ bool per_act_token, bool per_out_ch) {
117
+ if (out_tensors.dtype() == torch::kBFloat16) {
118
+ run_cutlass_moe_mm_sm100<cutlass::float_e4m3_t, cutlass::bfloat16_t>(
119
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
120
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
121
+ per_out_ch);
122
+ } else {
123
+ run_cutlass_moe_mm_sm100<cutlass::float_e4m3_t, cutlass::half_t>(
124
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
125
+ problem_sizes, a_strides, b_strides, c_strides, per_act_token,
126
+ per_out_ch);
127
+ }
128
+}
129
+
130
+void cutlass_moe_mm_sm100(
131
+ torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
132
+ torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
133
+ torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
134
+ torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
135
+ torch::Tensor const& b_strides, torch::Tensor const& c_strides,
136
+ bool per_act_token, bool per_out_ch) {
137
+ dispatch_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
138
+ expert_offsets, problem_sizes, a_strides, b_strides,
139
+ c_strides, per_act_token, per_out_ch);
140
+}
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm90.cu
renamed
+17
-13
@@ -21,10 +21,11 @@ struct sm90_fp8_config_default {
21
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
22
using TileShape = cute::Shape<cute::_64, cute::_256, cute::_128>;
23
using ClusterShape = cute::Shape<cute::_1, cute::_2, cute::_1>;
24
+ using ArchTag = cutlass::arch::Sm90;
25
26
using Cutlass3xGemm =
26
- cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
27
- KernelSchedule, EpilogueSchedule>;
27
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
28
+ ClusterShape, KernelSchedule, EpilogueSchedule>;
29
};
30
31
template <typename InType, typename OutType,
@@ -38,10 +39,12 @@ struct sm90_fp8_config_M4 {
39
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
40
using TileShape = cute::Shape<cute::_128, cute::_16, cute::_128>;
41
using ClusterShape = cute::Shape<cute::_1, cute::_1, cute::_1>;
42
+ using ArchTag = cutlass::arch::Sm90;
43
44
using Cutlass3xGemm =
43
- cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
44
- KernelSchedule, EpilogueSchedule, true>;
45
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
46
+ ClusterShape, KernelSchedule, EpilogueSchedule,
47
+ true>;
48
};
49
50
template <typename InType, typename OutType,
@@ -55,10 +58,12 @@ struct sm90_fp8_config_M64 {
58
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
59
using TileShape = cute::Shape<cute::_128, cute::_16, cute::_256>;
60
using ClusterShape = cute::Shape<cute::_2, cute::_1, cute::_1>;
61
+ using ArchTag = cutlass::arch::Sm90;
62
63
using Cutlass3xGemm =
60
- cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
61
- KernelSchedule, EpilogueSchedule, true>;
64
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
65
+ ClusterShape, KernelSchedule, EpilogueSchedule,
66
+ true>;
67
};
68
69
template <typename InType, typename OutType,
@@ -72,10 +77,11 @@ struct sm90_fp8_config_K8192 {
77
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
78
using TileShape = cute::Shape<cute::_128, cute::_128, cute::_128>;
79
using ClusterShape = cute::Shape<cute::_1, cute::_8, cute::_1>;
80
+ using ArchTag = cutlass::arch::Sm90;
81
82
using Cutlass3xGemm =
77
- cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
78
- KernelSchedule, EpilogueSchedule>;
83
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
84
+ ClusterShape, KernelSchedule, EpilogueSchedule>;
85
};
86
87
template <typename InType, typename OutType,
@@ -89,10 +95,11 @@ struct sm90_fp8_config_N8192 {
95
cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
96
using TileShape = cute::Shape<cute::_64, cute::_128, cute::_256>;
97
using ClusterShape = cute::Shape<cute::_1, cute::_8, cute::_1>;
98
+ using ArchTag = cutlass::arch::Sm90;
99
100
using Cutlass3xGemm =
94
- cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
95
- KernelSchedule, EpilogueSchedule>;
101
+ cutlass_3x_group_gemm<InType, OutType, ArchTag, Epilogue, TileShape,
102
+ ClusterShape, KernelSchedule, EpilogueSchedule>;
103
};
104
105
template <typename InType, typename OutType>
@@ -112,9 +119,6 @@ void run_cutlass_moe_mm_sm90(
119
TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn,
120
"B tensors must be of type float8_e4m3fn.");
121
115
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
116
- TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);
117
-
122
using Cutlass3xGemmN8192 = typename sm90_fp8_config_N8192<
123
InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
124
using Cutlass3xGemmK8192 = typename sm90_fp8_config_K8192<
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+1
-1
@@ -190,4 +190,4 @@ void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
190
static_cast<int32_t*>(problem_sizes2.data_ptr()),
191
static_cast<const int32_t*>(expert_num_tokens.data_ptr()), padded_m, n,
192
k);
193
-}
193
+}
\ No newline at end of file
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+34
-11
@@ -41,6 +41,16 @@ void cutlass_moe_mm_sm90(
41
42
#endif
43
44
+#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
45
+void cutlass_moe_mm_sm100(
46
+ torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
47
+ torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
48
+ torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
49
+ torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
50
+ torch::Tensor const& b_strides, torch::Tensor const& c_strides,
51
+ bool per_act_token, bool per_out_ch);
52
+#endif
53
+
54
#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
55
void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
56
torch::Tensor const& b,
@@ -130,10 +140,10 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
140
// and at least SM90 (Hopper)
141
142
#if defined CUDA_VERSION
133
- if (cuda_device_capability >= 90 && cuda_device_capability < 100) {
134
- return CUDA_VERSION >= 12000;
135
- } else if (cuda_device_capability >= 100) {
143
+ if (cuda_device_capability >= 100) {
144
return CUDA_VERSION >= 12080;
145
+ } else if (cuda_device_capability >= 90) {
146
+ return CUDA_VERSION >= 12000;
147
}
148
#endif
149
@@ -141,11 +151,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
151
}
152
153
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
144
- // CUTLASS grouped FP8 kernels need at least CUDA 12.3
145
- // and SM90 (Hopper)
154
+ // CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
155
+ // or CUDA 12.8 and SM100 (Blackwell)
156
157
#if defined CUDA_VERSION
148
- if (cuda_device_capability == 90) {
158
+ if (cuda_device_capability >= 100) {
159
+ return CUDA_VERSION >= 12080;
160
+ }
161
+ if (cuda_device_capability >= 90) {
162
return CUDA_VERSION >= 12030;
163
}
164
#endif
@@ -234,16 +247,26 @@ void cutlass_moe_mm(
247
torch::Tensor const& b_strides, torch::Tensor const& c_strides,
248
bool per_act_token, bool per_out_ch) {
249
int32_t version_num = get_sm_version_num();
250
+#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
251
+ if (version_num >= 100) {
252
+ cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
253
+ expert_offsets, problem_sizes, a_strides, b_strides,
254
+ c_strides, per_act_token, per_out_ch);
255
+ return;
256
+ }
257
+#endif
258
#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
238
- cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
239
- expert_offsets, problem_sizes, a_strides, b_strides,
240
- c_strides, per_act_token, per_out_ch);
241
- return;
259
+ if (version_num >= 90) {
260
+ cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
261
+ expert_offsets, problem_sizes, a_strides, b_strides,
262
+ c_strides, per_act_token, per_out_ch);
263
+ return;
264
+ }
265
#endif
266
TORCH_CHECK_NOT_IMPLEMENTED(
267
false,
268
"No compiled cutlass_scaled_mm for CUDA device capability: ", version_num,
246
- ". Required capability: 90");
269
+ ". Required capability: 90 or 100");
270
}
271
272
void get_cutlass_moe_mm_data(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+6
@@ -332,6 +332,12 @@ class CompressedTensorsConfig(QuantizationConfig):
332
return (self._check_scheme_supported(90, error=False, match_exact=True)
333
and self._is_fp8_w8a8(weight_quant, input_quant))
334
335
+ def _is_fp8_w8a8_sm100(self, weight_quant: BaseModel,
336
+ input_quant: BaseModel) -> bool:
337
+ return (self._check_scheme_supported(
338
+ 100, error=False, match_exact=True)
339
+ and self._is_fp8_w8a8(weight_quant, input_quant))
340
+
341
def _is_fp8_w8a16(self, weight_quant: BaseModel,
342
input_quant: BaseModel) -> bool:
343
# Confirm weights quantized.
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+27
-2
@@ -83,7 +83,8 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
83
return CompressedTensorsWNA16MarlinMoEMethod(quant_config)
84
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
85
return CompressedTensorsW4A4MoeMethod()
86
- elif quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant):
86
+ elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
87
+ or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)):
88
return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
89
elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
90
return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
@@ -740,6 +741,8 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
741
self.topk_indices_dtype = None
742
self.fused_experts = None # type: ignore
743
self.disable_expert_map = False
744
+ self.is_fp8_w8a8_sm100 = self.quant_config._is_fp8_w8a8_sm100(
745
+ self.weight_quant, self.input_quant)
746
747
def create_weights(self, layer: torch.nn.Module, num_experts: int,
748
hidden_size: int, intermediate_size_per_partition: int,
@@ -931,7 +934,29 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
934
935
per_act_token = (
936
self.input_quant.strategy == QuantizationStrategy.TOKEN)
934
-
937
+ per_channel_quant = (
938
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
939
+ # Triton fused_experts is faster in small batch sizes on SM100.
940
+ # Fall back to fused_experts in small batch sizes.
941
+ if self.is_fp8_w8a8_sm100 and topk_ids.shape[0] <= 8:
942
+ from vllm.model_executor.layers.fused_moe import fused_experts
943
+ return fused_experts(
944
+ x,
945
+ layer.w13_weight,
946
+ layer.w2_weight,
947
+ topk_weights,
948
+ topk_ids,
949
+ inplace=True,
950
+ activation=activation,
951
+ apply_router_weight_on_input=apply_router_weight_on_input,
952
+ use_fp8_w8a8=True,
953
+ per_channel_quant=per_channel_quant,
954
+ global_num_experts=global_num_experts,
955
+ expert_map=None if self.disable_expert_map else expert_map,
956
+ w1_scale=layer.w13_weight_scale,
957
+ w2_scale=layer.w2_weight_scale,
958
+ a1_scale=layer.w13_input_scale,
959
+ a2_scale=layer.w2_input_scale)
960
if self.fused_experts is None:
961
# If no modular kernel is provided, use cutlass_moe_fp8
962
from vllm.model_executor.layers.fused_moe.cutlass_moe import (