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 5, 01:00 UTC
[Kernel] Add MXFP4 W4A4 CUTLASS MoE kernel for SM100 (#37463)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Apr 17, 2026 at 19:42 UTC
a8bffaa13337236408e7fe6387e1da764191b05a
11 files changed
+1701
-14
.buildkite/test_areas/kernels.yaml
+1
@@ -141,6 +141,7 @@ steps:
141
- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
142
- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
143
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
144
+ - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
145
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
146
- pytest -v -s tests/kernels/moe/test_flashinfer.py
147
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
CMakeLists.txt
+3
-1
@@ -952,7 +952,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
952
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
953
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
954
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
955
- "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
955
+ "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
956
+ "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
957
+ "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
958
set_gencode_flags_for_srcs(
959
SRCS "${SRCS}"
960
CUDA_ARCHS "${FP4_ARCHS}")
csrc/libtorch_stable/ops.h
+23
@@ -134,4 +134,27 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
134
torch::stable::Tensor& input,
135
torch::stable::Tensor& input_global_scale);
136
137
+void mxfp4_experts_quant(
138
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
139
+ torch::stable::Tensor const& input,
140
+ torch::stable::Tensor const& input_offset_by_experts,
141
+ torch::stable::Tensor const& output_scale_offset_by_experts,
142
+ int64_t n_experts);
143
+
144
+void silu_and_mul_mxfp4_experts_quant(
145
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
146
+ torch::stable::Tensor const& input,
147
+ torch::stable::Tensor const& input_offset_by_experts,
148
+ torch::stable::Tensor const& output_scale_offset_by_experts,
149
+ int64_t n_experts);
150
+
151
+void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
152
+ const torch::stable::Tensor& a,
153
+ const torch::stable::Tensor& b,
154
+ const torch::stable::Tensor& a_blockscale,
155
+ const torch::stable::Tensor& b_blockscales,
156
+ const torch::stable::Tensor& problem_sizes,
157
+ const torch::stable::Tensor& expert_offsets,
158
+ const torch::stable::Tensor& sf_offsets);
159
+
160
#endif
csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu
new
+468
@@ -0,0 +1,468 @@
1
+/*
2
+ * SPDX-License-Identifier: Apache-2.0
3
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+ *
5
+ * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
6
+ * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
7
+ */
8
+
9
+#include <torch/csrc/stable/library.h>
10
+#include <torch/csrc/stable/tensor.h>
11
+#include "libtorch_stable/torch_utils.h"
12
+
13
+#include <cutlass/arch/arch.h>
14
+
15
+#include "cutlass_extensions/common.hpp"
16
+
17
+#include "cute/tensor.hpp"
18
+#include "cutlass/tensor_ref.h"
19
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
20
+#include "cutlass/epilogue/thread/linear_combination.h"
21
+#include "cutlass/gemm/dispatch_policy.hpp"
22
+#include "cutlass/gemm/group_array_problem_shape.hpp"
23
+#include "cutlass/gemm/collective/collective_builder.hpp"
24
+#include "cutlass/epilogue/collective/collective_builder.hpp"
25
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
26
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
27
+
28
+#include "cutlass/util/packed_stride.hpp"
29
+#include <cassert>
30
+
31
+using namespace cute;
32
+
33
+// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
34
+template <typename ElementAB, typename ElementC, typename ElementSF,
35
+ typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
36
+__global__ void __mxfp4_get_group_gemm_starts(
37
+ ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
38
+ ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
39
+ LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
40
+ ElementAB* a_base_as_int, ElementAB* b_base_as_int,
41
+ ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
42
+ ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
43
+ const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
44
+ int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
45
+ const int64_t a_stride_val, const int64_t b_stride_val,
46
+ const int64_t c_stride_val, const int K, const int N) {
47
+ int64_t expert_id = threadIdx.x;
48
+ if (expert_id >= gridDim.x * blockDim.x) {
49
+ return;
50
+ }
51
+ int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
52
+ int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
53
+ int64_t group_size = 32;
54
+ int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
55
+ int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
56
+ int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
57
+ assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
58
+ "unexpected problem sizes");
59
+
60
+ int64_t half_k = static_cast<int64_t>(k / 2);
61
+ int64_t group_k = static_cast<int64_t>(k / group_size);
62
+ // Shape of A as uint8/byte = [M, K // 2]
63
+ a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
64
+ // Shape of B as uint8/byte = [E, N, K // 2]
65
+ b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
66
+ // Shape of C = [M, N]
67
+ out_offsets[expert_id] = out_base_as_int + expert_offset * n;
68
+ // Shape of a_scale = [sum(sf_sizes), K // group_size]
69
+ a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
70
+
71
+ assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
72
+ 0 &&
73
+ "TMA requires 128-byte alignment");
74
+
75
+ // Shape of B scale = [E, N, K // group_size]
76
+ b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
77
+ assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
78
+ 0 &&
79
+ "TMA requires 128-byte alignment");
80
+
81
+ // Initialize strides
82
+ a_strides[expert_id] = a_stride_val;
83
+ b_strides[expert_id] = b_stride_val;
84
+ c_strides[expert_id] = c_stride_val;
85
+
86
+ LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
87
+ LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
88
+
89
+ *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
90
+ static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
91
+ *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
92
+ static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
93
+}
94
+
95
+#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \
96
+ TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
97
+ LayoutSFB, ScaleConfig) \
98
+ else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \
99
+ __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
100
+ LayoutSFB, ScaleConfig> \
101
+ <<<1, num_experts, 0, stream>>>( \
102
+ static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
103
+ static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
104
+ static_cast<C_TYPE**>(out_starts.data_ptr()), \
105
+ static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
106
+ static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
107
+ reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
108
+ reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
109
+ static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
110
+ static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
111
+ static_cast<C_TYPE*>(out_tensors.data_ptr()), \
112
+ static_cast<SF_TYPE*>(a_scales.data_ptr()), \
113
+ static_cast<SF_TYPE*>(b_scales.data_ptr()), \
114
+ static_cast<int32_t*>(expert_offsets.data_ptr()), \
115
+ static_cast<int32_t*>(sf_offsets.data_ptr()), \
116
+ static_cast<int32_t*>(problem_sizes.data_ptr()), \
117
+ static_cast<int64_t*>(a_strides.data_ptr()), \
118
+ static_cast<int64_t*>(b_strides.data_ptr()), \
119
+ static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \
120
+ b_stride_val, c_stride_val, K, N); \
121
+ }
122
+
123
+template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
124
+void mxfp4_run_get_group_gemm_starts(
125
+ const torch::stable::Tensor& a_starts,
126
+ const torch::stable::Tensor& b_starts,
127
+ const torch::stable::Tensor& out_starts,
128
+ const torch::stable::Tensor& a_scales_starts,
129
+ const torch::stable::Tensor& b_scales_starts,
130
+ const torch::stable::Tensor& layout_sfa,
131
+ const torch::stable::Tensor& layout_sfb,
132
+ const torch::stable::Tensor& a_strides,
133
+ const torch::stable::Tensor& b_strides,
134
+ const torch::stable::Tensor& c_strides, int64_t a_stride_val,
135
+ int64_t b_stride_val, int64_t c_stride_val,
136
+ torch::stable::Tensor const& a_tensors,
137
+ torch::stable::Tensor const& b_tensors,
138
+ torch::stable::Tensor const& out_tensors,
139
+ torch::stable::Tensor const& a_scales,
140
+ torch::stable::Tensor const& b_scales,
141
+ torch::stable::Tensor const& expert_offsets,
142
+ torch::stable::Tensor const& sf_offsets,
143
+ torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
144
+ int num_experts = (int)expert_offsets.size(0);
145
+ auto stream = get_current_cuda_stream(a_tensors.get_device_index());
146
+
147
+ STD_TORCH_CHECK(out_tensors.size(1) == N,
148
+ "Output tensor shape doesn't match expected shape");
149
+ STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
150
+ "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
151
+ " dimension must match");
152
+ if (false) {
153
+ }
154
+ // MXFP4 uses E8M0 (float_ue8m0_t) scale factors
155
+ __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
156
+ torch::headeronly::ScalarType::BFloat16,
157
+ cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
158
+ ScaleConfig)
159
+ __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
160
+ torch::headeronly::ScalarType::Half, half,
161
+ LayoutSFA, LayoutSFB, ScaleConfig)
162
+ else {
163
+ STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
164
+ }
165
+}
166
+
167
+template <typename OutType>
168
+void run_mxfp4_blockwise_scaled_group_mm_sm100(
169
+ torch::stable::Tensor& output, const torch::stable::Tensor& a,
170
+ const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
171
+ const torch::stable::Tensor& b_blockscales,
172
+ const torch::stable::Tensor& problem_sizes,
173
+ const torch::stable::Tensor& expert_offsets,
174
+ const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
175
+ using ProblemShape =
176
+ cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
177
+ using ElementType = cutlass::float_e2m1_t;
178
+ using ElementSFType = cutlass::float_ue8m0_t;
179
+ using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
180
+ using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
181
+
182
+ using ElementC = OutType;
183
+ using ElementD = ElementC;
184
+ using ElementAccumulator = float;
185
+ // Layout definitions
186
+ using LayoutA = cutlass::layout::RowMajor;
187
+ using LayoutB = cutlass::layout::ColumnMajor;
188
+ using LayoutC = cutlass::layout::RowMajor;
189
+ using LayoutD = LayoutC;
190
+
191
+ static constexpr int AlignmentA = 32;
192
+ static constexpr int AlignmentB = 32;
193
+ static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
194
+ static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
195
+
196
+ // Architecture definitions
197
+ using ArchTag = cutlass::arch::Sm100;
198
+ using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
199
+ using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
200
+ using StageCountType = cutlass::gemm::collective::StageCountAuto;
201
+
202
+ using ClusterShape = Shape<_1, _1, _1>;
203
+ struct MMA1SMConfig {
204
+ using MmaTileShape = Shape<_128, _128, _128>;
205
+ using KernelSchedule =
206
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
207
+ using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
208
+ };
209
+
210
+ using CollectiveEpilogue =
211
+ typename cutlass::epilogue::collective::CollectiveBuilder<
212
+ ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
213
+ ClusterShape, Shape<_128, _64>, ElementAccumulator,
214
+ ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
215
+ LayoutC*, AlignmentD,
216
+ typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
217
+
218
+ using CollectiveMainloop =
219
+ typename cutlass::gemm::collective::CollectiveBuilder<
220
+ ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
221
+ ElementB, LayoutB*, AlignmentB, ElementAccumulator,
222
+ typename MMA1SMConfig::MmaTileShape, ClusterShape,
223
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
224
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
225
+ typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
226
+
227
+ using GemmKernel =
228
+ cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
229
+ CollectiveEpilogue>;
230
+
231
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
232
+ using StrideA = typename Gemm::GemmKernel::InternalStrideA;
233
+ using StrideB = typename Gemm::GemmKernel::InternalStrideB;
234
+ using StrideC = typename Gemm::GemmKernel::InternalStrideC;
235
+ using StrideD = typename Gemm::GemmKernel::InternalStrideD;
236
+
237
+ using LayoutSFA =
238
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
239
+ using LayoutSFB =
240
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
241
+ using ScaleConfig =
242
+ typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
243
+
244
+ using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
245
+ int num_experts = static_cast<int>(expert_offsets.size(0));
246
+
247
+ torch::stable::Tensor a_ptrs =
248
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
249
+ std::nullopt, a.device());
250
+ torch::stable::Tensor b_ptrs =
251
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
252
+ std::nullopt, a.device());
253
+ torch::stable::Tensor out_ptrs =
254
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
255
+ std::nullopt, a.device());
256
+ torch::stable::Tensor a_scales_ptrs =
257
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
258
+ std::nullopt, a.device());
259
+ torch::stable::Tensor b_scales_ptrs =
260
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
261
+ std::nullopt, a.device());
262
+ torch::stable::Tensor layout_sfa = torch::stable::empty(
263
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
264
+ a.device());
265
+ torch::stable::Tensor layout_sfb = torch::stable::empty(
266
+ {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
267
+ a.device());
268
+ torch::stable::Tensor a_strides1 =
269
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
270
+ std::nullopt, a.device());
271
+ torch::stable::Tensor b_strides1 =
272
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
273
+ std::nullopt, a.device());
274
+ torch::stable::Tensor c_strides1 =
275
+ torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
276
+ std::nullopt, a.device());
277
+
278
+ mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
279
+ a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
280
+ layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
281
+ b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
282
+ b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
283
+
284
+ // Create an instance of the GEMM
285
+ Gemm gemm_op;
286
+
287
+ UnderlyingProblemShape* problem_sizes_as_shapes =
288
+ static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
289
+
290
+ // Set the Scheduler info
291
+ cutlass::KernelHardwareInfo hw_info;
292
+ using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
293
+ PersistentTileSchedulerSm100GroupParams<
294
+ typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
295
+ typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
296
+ scheduler.raster_order = RasterOrderOptions::AlongM;
297
+ hw_info.device_id = a.get_device_index();
298
+ static std::unordered_map<int, int> cached_sm_counts;
299
+ if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
300
+ cached_sm_counts[hw_info.device_id] =
301
+ cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
302
+ hw_info.device_id);
303
+ }
304
+ hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
305
+
306
+ // Mainloop Arguments
307
+ typename GemmKernel::MainloopArguments mainloop_args{
308
+ static_cast<const ElementType**>(a_ptrs.data_ptr()),
309
+ static_cast<StrideA*>(a_strides1.data_ptr()),
310
+ static_cast<const ElementType**>(b_ptrs.data_ptr()),
311
+ static_cast<StrideB*>(b_strides1.data_ptr()),
312
+ static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
313
+ reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
314
+ static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
315
+ reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
316
+
317
+ // Epilogue Arguments
318
+ typename GemmKernel::EpilogueArguments epilogue_args{
319
+ {}, // epilogue.thread
320
+ nullptr,
321
+ static_cast<StrideC*>(c_strides1.data_ptr()),
322
+ static_cast<ElementD**>(out_ptrs.data_ptr()),
323
+ static_cast<StrideC*>(c_strides1.data_ptr())};
324
+ auto& fusion_args = epilogue_args.thread;
325
+ // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
326
+ fusion_args.alpha_ptr = nullptr;
327
+ fusion_args.beta_ptr = nullptr;
328
+ fusion_args.alpha = 1.0f;
329
+ fusion_args.alpha_ptr_array = nullptr;
330
+ fusion_args.dAlpha = {_0{}, _0{}, 0};
331
+ fusion_args.beta = 0.0f;
332
+ fusion_args.beta_ptr_array = nullptr;
333
+ fusion_args.dBeta = {_0{}, _0{}, 0};
334
+
335
+ // Gemm Arguments
336
+ typename GemmKernel::Arguments args{
337
+ cutlass::gemm::GemmUniversalMode::kGrouped,
338
+ {num_experts, problem_sizes_as_shapes, nullptr},
339
+ mainloop_args,
340
+ epilogue_args,
341
+ hw_info,
342
+ scheduler};
343
+
344
+ size_t workspace_size = Gemm::get_workspace_size(args);
345
+ auto workspace =
346
+ torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
347
+ std::nullopt, a.device());
348
+ const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
349
+
350
+ auto can_implement_status = gemm_op.can_implement(args);
351
+ STD_TORCH_CHECK(
352
+ can_implement_status == cutlass::Status::kSuccess,
353
+ "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
354
+
355
+ // Run the GEMM
356
+ auto status = gemm_op.initialize(args, workspace.data_ptr());
357
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
358
+ "Failed to initialize MXFP4 GEMM: status=", (int)status,
359
+ " workspace_size=", workspace_size,
360
+ " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
361
+
362
+ status = gemm_op.run(args, workspace.data_ptr(), stream);
363
+ STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
364
+ "Failed to run MXFP4 GEMM");
365
+}
366
+
367
+template <typename OutType>
368
+void run_mxfp4_blockwise_scaled_group_mm(
369
+ torch::stable::Tensor& output, const torch::stable::Tensor& a,
370
+ const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
371
+ const torch::stable::Tensor& b_blockscales,
372
+ const torch::stable::Tensor& problem_sizes,
373
+ const torch::stable::Tensor& expert_offsets,
374
+ const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
375
+ int32_t version_num = get_sm_version_num();
376
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
377
+ if (version_num >= 100 && version_num < 120) {
378
+ run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
379
+ output, a, b, a_blockscale, b_blockscales, problem_sizes,
380
+ expert_offsets, sf_offsets, M, N, K);
381
+ return;
382
+ }
383
+#endif
384
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
385
+ false,
386
+ "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
387
+ version_num, ". Required capability: 100");
388
+}
389
+
390
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
391
+constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
392
+// E8M0 scale factors stored as uint8
393
+constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
394
+#endif
395
+
396
+#define CHECK_TYPE(x, st, m) \
397
+ STD_TORCH_CHECK(x.scalar_type() == st, \
398
+ ": Inconsistency of torch::stable::Tensor type:", m)
399
+#define CHECK_TH_CUDA(x, m) \
400
+ STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
401
+#define CHECK_CONTIGUOUS(x, m) \
402
+ STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
403
+#define CHECK_INPUT(x, st, m) \
404
+ CHECK_TH_CUDA(x, m); \
405
+ CHECK_CONTIGUOUS(x, m); \
406
+ CHECK_TYPE(x, st, m)
407
+
408
+void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
409
+ const torch::stable::Tensor& a,
410
+ const torch::stable::Tensor& b,
411
+ const torch::stable::Tensor& a_blockscale,
412
+ const torch::stable::Tensor& b_blockscales,
413
+ const torch::stable::Tensor& problem_sizes,
414
+ const torch::stable::Tensor& expert_offsets,
415
+ const torch::stable::Tensor& sf_offsets) {
416
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
417
+ // Input validation
418
+ CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
419
+ CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
420
+ // MXFP4 uses E8M0 scale factors (stored as uint8)
421
+ CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
422
+ CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
423
+
424
+ STD_TORCH_CHECK(
425
+ a_blockscale.dim() == 2,
426
+ "expected a_blockscale to be of shape [num_experts, rounded_m,"
427
+ " k // group_size], observed rank: ",
428
+ a_blockscale.dim())
429
+ STD_TORCH_CHECK(b_blockscales.dim() == 3,
430
+ "expected b_blockscale to be of shape: "
431
+ " [num_experts, n, k // group_size], observed rank: ",
432
+ b_blockscales.dim())
433
+ STD_TORCH_CHECK(problem_sizes.dim() == 2,
434
+ "problem_sizes must be a 2D tensor");
435
+ STD_TORCH_CHECK(problem_sizes.size(1) == 3,
436
+ "problem_sizes must have the shape (num_experts, 3)");
437
+ STD_TORCH_CHECK(
438
+ problem_sizes.size(0) == expert_offsets.size(0),
439
+ "Number of experts in problem_sizes must match expert_offsets");
440
+ STD_TORCH_CHECK(
441
+ problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
442
+ "problem_sizes must be int32.");
443
+
444
+ int M = static_cast<int>(a.size(0));
445
+ int N = static_cast<int>(b.size(1));
446
+ int E = static_cast<int>(b.size(0));
447
+ int K = static_cast<int>(2 * b.size(2));
448
+
449
+ if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
450
+ run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
451
+ output, a, b, a_blockscale, b_blockscales, problem_sizes,
452
+ expert_offsets, sf_offsets, M, N, K);
453
+ } else {
454
+ run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
455
+ output, a, b, a_blockscale, b_blockscales, problem_sizes,
456
+ expert_offsets, sf_offsets, M, N, K);
457
+ }
458
+#else
459
+ STD_TORCH_CHECK_NOT_IMPLEMENTED(
460
+ false,
461
+ "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
462
+ "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
463
+#endif
464
+}
465
+
466
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
467
+ m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
468
+}
csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu
new
+422
@@ -0,0 +1,422 @@
1
+/*
2
+ * SPDX-License-Identifier: Apache-2.0
3
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+ *
5
+ * MXFP4 activation quantization kernel for MoE experts.
6
+ * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
7
+ * over 32-element groups.
8
+ *
9
+ * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
10
+ * - Block size 32 (2 threads per SF in PACK16 mode)
11
+ * - E8M0 (power-of-two) scale factors
12
+ * - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
13
+ */
14
+
15
+// MXFP4 requires PACK16 mode (16 elements per thread) so that
16
+// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
17
+// Must be defined before any header that (transitively) includes
18
+// nvfp4_utils.cuh.
19
+#define NVFP4_ENABLE_ELTS16 1
20
+
21
+#include <cuda.h>
22
+#include <cuda_runtime_api.h>
23
+#include <cuda_runtime.h>
24
+#include <cuda_fp8.h>
25
+
26
+#include <torch/csrc/stable/tensor.h>
27
+#include "libtorch_stable/torch_utils.h"
28
+#include "libtorch_stable/dispatch_utils.h"
29
+#include "cuda_vec_utils.cuh"
30
+#include "cuda_utils.h"
31
+
32
+#include "nvfp4_utils.cuh"
33
+static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
34
+ "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
35
+
36
+#include "launch_bounds_utils.h"
37
+
38
+namespace vllm {
39
+
40
+// MXFP4 block size constants
41
+static constexpr int MXFP4_SF_VEC_SIZE = 32;
42
+
43
+// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
44
+// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
45
+static constexpr int MXFP4_NUM_THREADS_PER_SF =
46
+ MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
47
+
48
+// MXFP4 quantization kernel for experts.
49
+// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
50
+// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
51
+// SiLU(gate)*up before quantization.
52
+template <class Type, bool FUSE_SILU_MUL = false,
53
+ bool SMALL_NUM_EXPERTS = false>
54
+__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
55
+ mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
56
+ fp4_packed_t* out, uint32_t* SFout,
57
+ uint32_t* input_offset_by_experts,
58
+ uint32_t* output_scale_offset_by_experts,
59
+ int n_experts, bool low_latency) {
60
+ using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
61
+ static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
62
+ "Vec size is not matched.");
63
+
64
+ // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
65
+ int32_t const numKTiles = (numCols + 127) / 128;
66
+
67
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
68
+ int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
69
+ int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
70
+
71
+ for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
72
+ globalIdx += gridDim.x * blockDim.x) {
73
+ int rowIdx = globalIdx / colsPerRow;
74
+ int colIdx = globalIdx % colsPerRow;
75
+
76
+ int rowIdx_in_expert = 0;
77
+ int expert_idx = 0;
78
+
79
+ if constexpr (SMALL_NUM_EXPERTS) {
80
+ for (int i = 0; i < n_experts; i++) {
81
+ uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
82
+ uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
83
+ if (rowIdx >= current_offset && rowIdx < next_offset) {
84
+ rowIdx_in_expert = rowIdx - current_offset;
85
+ expert_idx = i;
86
+ break;
87
+ }
88
+ }
89
+ } else {
90
+ uint32_t local_offsets[17];
91
+ for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
92
+ *reinterpret_cast<int4*>(local_offsets) =
93
+ __ldca(reinterpret_cast<const int4*>(
94
+ &input_offset_by_experts[chunk_start]));
95
+ *reinterpret_cast<int4*>(local_offsets + 4) =
96
+ __ldca(reinterpret_cast<const int4*>(
97
+ &input_offset_by_experts[chunk_start + 4]));
98
+ *reinterpret_cast<int4*>(local_offsets + 8) =
99
+ __ldca(reinterpret_cast<const int4*>(
100
+ &input_offset_by_experts[chunk_start + 8]));
101
+ *reinterpret_cast<int4*>(local_offsets + 12) =
102
+ __ldca(reinterpret_cast<const int4*>(
103
+ &input_offset_by_experts[chunk_start + 12]));
104
+ local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
105
+
106
+#pragma unroll
107
+ for (int i = 0; i < 16; i++) {
108
+ if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
109
+ rowIdx_in_expert = rowIdx - local_offsets[i];
110
+ expert_idx = chunk_start + i;
111
+ break;
112
+ }
113
+ }
114
+ }
115
+ }
116
+
117
+ // Load input and optionally apply fused SiLU+Mul
118
+ int64_t inOffset = rowIdx * inColsPerRow + colIdx;
119
+ PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
120
+ PackedVec quant_input;
121
+ if constexpr (FUSE_SILU_MUL) {
122
+ PackedVec in_vec_up =
123
+ reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
124
+ quant_input = compute_silu_mul(in_vec, in_vec_up);
125
+ } else {
126
+ quant_input = in_vec;
127
+ }
128
+
129
+ // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
130
+ int64_t outOffset = rowIdx * colsPerRow + colIdx;
131
+ auto& out_pos = out[outOffset];
132
+
133
+ uint32_t* SFout_in_expert =
134
+ SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
135
+
136
+ // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
137
+ auto sf_out =
138
+ cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
139
+ rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
140
+
141
+ // Block E8M0 scales only; no extra tensor-level scale in this path
142
+ constexpr float SFScaleVal = 1.0f;
143
+ // UE8M0_SF=true for MXFP4 E8M0 scale factors
144
+ out_pos =
145
+ cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
146
+ quant_input, SFScaleVal, sf_out);
147
+ }
148
+}
149
+
150
+// Large M_topk variant using shared memory for expert offsets
151
+template <class Type, bool FUSE_SILU_MUL = false,
152
+ bool SMALL_NUM_EXPERTS = false>
153
+__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
154
+ mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
155
+ fp4_packed_t* out, uint32_t* SFout,
156
+ uint32_t* input_offset_by_experts,
157
+ uint32_t* output_scale_offset_by_experts,
158
+ int n_experts) {
159
+ using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
160
+ static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
161
+ "Vec size is not matched.");
162
+
163
+ // MXFP4: numKTiles = ceil(numCols / 128)
164
+ int32_t const numKTiles = (numCols + 127) / 128;
165
+
166
+ extern __shared__ uint32_t shared_input_offsets[];
167
+
168
+ if constexpr (SMALL_NUM_EXPERTS) {
169
+ for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
170
+ shared_input_offsets[i] = input_offset_by_experts[i];
171
+ }
172
+ } else {
173
+ for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
174
+ *reinterpret_cast<int4*>(&shared_input_offsets[i]) =
175
+ *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
176
+ }
177
+ if (threadIdx.x == 0) {
178
+ shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
179
+ }
180
+ }
181
+
182
+ __syncthreads();
183
+
184
+ int tid = blockIdx.x * blockDim.x + threadIdx.x;
185
+ int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
186
+ int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
187
+
188
+ for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
189
+ globalIdx += gridDim.x * blockDim.x) {
190
+ int rowIdx = globalIdx / colsPerRow;
191
+ int colIdx = globalIdx % colsPerRow;
192
+
193
+ int rowIdx_in_expert = 0;
194
+ int expert_idx = 0;
195
+
196
+ // Binary search through experts using shared memory
197
+ int left = 0, right = n_experts - 1;
198
+ while (left <= right) {
199
+ int mid = (left + right) / 2;
200
+ uint32_t mid_offset = shared_input_offsets[mid];
201
+ uint32_t next_offset = shared_input_offsets[mid + 1];
202
+
203
+ if (rowIdx >= mid_offset && rowIdx < next_offset) {
204
+ rowIdx_in_expert = rowIdx - mid_offset;
205
+ expert_idx = mid;
206
+ break;
207
+ } else if (rowIdx < mid_offset) {
208
+ right = mid - 1;
209
+ } else {
210
+ left = mid + 1;
211
+ }
212
+ }
213
+
214
+ int64_t inOffset = rowIdx * inColsPerRow + colIdx;
215
+ PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
216
+ PackedVec quant_input;
217
+ if constexpr (FUSE_SILU_MUL) {
218
+ PackedVec in_vec_up =
219
+ reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
220
+ quant_input = compute_silu_mul(in_vec, in_vec_up);
221
+ } else {
222
+ quant_input = in_vec;
223
+ }
224
+
225
+ int64_t outOffset = rowIdx * colsPerRow + colIdx;
226
+ auto& out_pos = out[outOffset];
227
+
228
+ // MXFP4 has no global scale - only block-level E8M0 scale factors
229
+ constexpr float SFScaleVal = 1.0f;
230
+
231
+ uint32_t* SFout_in_expert =
232
+ SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
233
+
234
+ auto sf_out =
235
+ cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
236
+ rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
237
+
238
+ out_pos =
239
+ cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
240
+ quant_input, SFScaleVal, sf_out);
241
+ }
242
+}
243
+
244
+template <typename T, bool FUSE_SILU_MUL = false>
245
+void mxfp4_quant_impl(void* output, void* output_scale, void* input,
246
+ void* input_offset_by_experts,
247
+ void* output_scale_offset_by_experts, int m_topk, int k,
248
+ int n_experts, cudaStream_t stream) {
249
+ int multiProcessorCount =
250
+ get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
251
+
252
+ int const workSizePerRow = k / ELTS_PER_THREAD;
253
+ int const totalWorkSize = m_topk * workSizePerRow;
254
+ dim3 block(std::min(workSizePerRow, 512));
255
+ int const numBlocksPerSM =
256
+ vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
257
+ dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
258
+ multiProcessorCount * numBlocksPerSM));
259
+ while (grid.x <= multiProcessorCount && block.x > 64) {
260
+ grid.x *= 2;
261
+ block.x = (block.x + 1) / 2;
262
+ }
263
+
264
+ int const blockRepeat =
265
+ (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
266
+ if (blockRepeat > 1) {
267
+ size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
268
+ if (n_experts >= 4) {
269
+ mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
270
+ <<<grid, block, shared_mem_size, stream>>>(
271
+ m_topk, k, reinterpret_cast<T*>(input),
272
+ reinterpret_cast<fp4_packed_t*>(output),
273
+ reinterpret_cast<uint32_t*>(output_scale),
274
+ reinterpret_cast<uint32_t*>(input_offset_by_experts),
275
+ reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
276
+ n_experts);
277
+ } else {
278
+ mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
279
+ <<<grid, block, shared_mem_size, stream>>>(
280
+ m_topk, k, reinterpret_cast<T*>(input),
281
+ reinterpret_cast<fp4_packed_t*>(output),
282
+ reinterpret_cast<uint32_t*>(output_scale),
283
+ reinterpret_cast<uint32_t*>(input_offset_by_experts),
284
+ reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
285
+ n_experts);
286
+ }
287
+ } else {
288
+ if (n_experts >= 16) {
289
+ mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
290
+ <<<grid, block, 0, stream>>>(
291
+ m_topk, k, reinterpret_cast<T*>(input),
292
+ reinterpret_cast<fp4_packed_t*>(output),
293
+ reinterpret_cast<uint32_t*>(output_scale),
294
+ reinterpret_cast<uint32_t*>(input_offset_by_experts),
295
+ reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
296
+ n_experts, /* bool low_latency */ true);
297
+ } else {
298
+ mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
299
+ m_topk, k, reinterpret_cast<T*>(input),
300
+ reinterpret_cast<fp4_packed_t*>(output),
301
+ reinterpret_cast<uint32_t*>(output_scale),
302
+ reinterpret_cast<uint32_t*>(input_offset_by_experts),
303
+ reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
304
+ n_experts, /* bool low_latency */ true);
305
+ }
306
+ }
307
+}
308
+
309
+} // namespace vllm
310
+
311
+/*Quantization entry for mxfp4 experts quantization*/
312
+#define CHECK_TH_CUDA(x, m) \
313
+ STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
314
+#define CHECK_CONTIGUOUS(x, m) \
315
+ STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
316
+#define CHECK_INPUT(x, m) \
317
+ CHECK_TH_CUDA(x, m); \
318
+ CHECK_CONTIGUOUS(x, m);
319
+
320
+constexpr auto HALF = torch::headeronly::ScalarType::Half;
321
+constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
322
+constexpr auto INT = torch::headeronly::ScalarType::Int;
323
+constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
324
+
325
+static constexpr int MXFP4_BLOCK_SIZE = 32;
326
+
327
+static void validate_mxfp4_experts_quant_inputs(
328
+ torch::stable::Tensor const& output,
329
+ torch::stable::Tensor const& output_scale,
330
+ torch::stable::Tensor const& input,
331
+ torch::stable::Tensor const& input_offset_by_experts,
332
+ torch::stable::Tensor const& output_scale_offset_by_experts,
333
+ int64_t n_experts, int64_t m_topk, int64_t k) {
334
+ CHECK_INPUT(output, "output");
335
+ CHECK_INPUT(output_scale, "output_scale");
336
+ CHECK_INPUT(input, "input");
337
+ CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
338
+ CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
339
+
340
+ STD_TORCH_CHECK(output.dim() == 2);
341
+ STD_TORCH_CHECK(output_scale.dim() == 2);
342
+ STD_TORCH_CHECK(input.dim() == 2);
343
+ STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
344
+ STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
345
+
346
+ STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
347
+ STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
348
+ STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
349
+ // output is uint8 (two mxfp4 values packed into one uint8)
350
+ // output_scale is int32 (four E8M0 values packed into one int32)
351
+ STD_TORCH_CHECK(output.scalar_type() == UINT8);
352
+ STD_TORCH_CHECK(output_scale.scalar_type() == INT);
353
+
354
+ STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
355
+ STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
356
+ STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
357
+ STD_TORCH_CHECK(output.size(0) == m_topk);
358
+ STD_TORCH_CHECK(output.size(1) == k / 2);
359
+ int scales_k = k / MXFP4_BLOCK_SIZE;
360
+ // K-dimension scale columns padded to a multiple of 4 for swizzle layout
361
+ int padded_k = (scales_k + (4 - 1)) / 4 * 4;
362
+ // 4 = 4 E8M0 values packed into one int32
363
+ STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
364
+}
365
+
366
+void mxfp4_experts_quant(
367
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
368
+ torch::stable::Tensor const& input,
369
+ torch::stable::Tensor const& input_offset_by_experts,
370
+ torch::stable::Tensor const& output_scale_offset_by_experts,
371
+ int64_t n_experts) {
372
+ auto m_topk = input.size(0);
373
+ auto k = input.size(1);
374
+
375
+ validate_mxfp4_experts_quant_inputs(
376
+ output, output_scale, input, input_offset_by_experts,
377
+ output_scale_offset_by_experts, n_experts, m_topk, k);
378
+
379
+ const torch::stable::accelerator::DeviceGuard device_guard(
380
+ input.get_device_index());
381
+ const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
382
+
383
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
384
+ input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
385
+ using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
386
+ vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
387
+ output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
388
+ input_offset_by_experts.data_ptr(),
389
+ output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
390
+ stream);
391
+ });
392
+}
393
+
394
+void silu_and_mul_mxfp4_experts_quant(
395
+ torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
396
+ torch::stable::Tensor const& input,
397
+ torch::stable::Tensor const& input_offset_by_experts,
398
+ torch::stable::Tensor const& output_scale_offset_by_experts,
399
+ int64_t n_experts) {
400
+ auto m_topk = input.size(0);
401
+ auto k_times_2 = input.size(1);
402
+ STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
403
+ auto k = k_times_2 / 2;
404
+
405
+ validate_mxfp4_experts_quant_inputs(
406
+ output, output_scale, input, input_offset_by_experts,
407
+ output_scale_offset_by_experts, n_experts, m_topk, k);
408
+
409
+ const torch::stable::accelerator::DeviceGuard device_guard(
410
+ input.get_device_index());
411
+ const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
412
+
413
+ VLLM_STABLE_DISPATCH_HALF_TYPES(
414
+ input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
415
+ using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
416
+ vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
417
+ output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
418
+ input_offset_by_experts.data_ptr(),
419
+ output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
420
+ stream);
421
+ });
422
+}
csrc/libtorch_stable/torch_bindings.cpp
+22
@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
116
" Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
117
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
118
119
+ // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
120
+ ops.def(
121
+ "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
122
+ " Tensor a_blockscale, Tensor b_blockscales,"
123
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
124
+
125
// Compute NVFP4 block quantized tensor.
126
ops.def(
127
"scaled_fp4_quant(Tensor input,"
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
155
"Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
156
"Tensor output_scale_offset_by_experts) -> ()");
157
158
+ // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
159
+ ops.def(
160
+ "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
161
+ "Tensor input, Tensor input_offset_by_experts,"
162
+ "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
163
+
164
+ // Fused SiLU+Mul+MXFP4 experts quantization.
165
+ ops.def(
166
+ "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
167
+ "output_scale,"
168
+ "Tensor input, Tensor input_offset_by_experts,"
169
+ "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
170
+
171
// Fused SiLU+Mul+NVFP4 quantization.
172
ops.def(
173
"silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -233,6 +252,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
252
ops.impl("silu_and_mul_scaled_fp4_experts_quant",
253
TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
254
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
255
+ ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
256
+ ops.impl("silu_and_mul_mxfp4_experts_quant",
257
+ TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
258
259
// W4A8 ops: impl registrations are in the source files
260
// (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)
tests/kernels/moe/test_mxfp4_moe.py
new
+248
@@ -0,0 +1,248 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
5
+
6
+import random
7
+
8
+import pytest
9
+import torch
10
+
11
+from tests.kernels.utils import torch_moe_single
12
+from vllm import _custom_ops as ops
13
+from vllm.platforms import current_platform
14
+from vllm.utils.torch_utils import set_random_seed
15
+
16
+random.seed(42)
17
+set_random_seed(42)
18
+
19
+MXFP4_BLOCK_SIZE = 32
20
+
21
+
22
+def align(val: int, alignment: int = 128) -> int:
23
+ return int((val + alignment - 1) // alignment * alignment)
24
+
25
+
26
+def calc_diff(x, y):
27
+ x, y = x.double(), y.double()
28
+ denominator = (x * x + y * y).sum()
29
+ sim = 2 * (x * y).sum() / denominator
30
+ return 1 - sim
31
+
32
+
33
+def is_sm100_supported() -> bool:
34
+ return current_platform.is_cuda() and current_platform.is_device_capability_family(
35
+ 100
36
+ )
37
+
38
+
39
+def compute_ref_output(
40
+ input_tensor: torch.Tensor,
41
+ weight_list: list[torch.Tensor],
42
+ expert_offsets: list[int],
43
+ expert_offset: int,
44
+ num_experts: int,
45
+) -> torch.Tensor:
46
+ """Reference output using torch_moe_single with top-1 routing."""
47
+ score = torch.full(
48
+ (expert_offset, num_experts),
49
+ -1e9,
50
+ device=input_tensor.device,
51
+ dtype=torch.float32,
52
+ )
53
+ for g in range(num_experts):
54
+ start = expert_offsets[g]
55
+ end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
56
+ score[start:end, g] = 0.0
57
+
58
+ return torch_moe_single(
59
+ input_tensor, torch.stack(weight_list, dim=0), score, topk=1
60
+ )
61
+
62
+
63
+@pytest.mark.skipif(
64
+ not is_sm100_supported(),
65
+ reason="cutlass_mxfp4_group_mm requires CUDA SM100",
66
+)
67
+@pytest.mark.parametrize("num_experts", [8, 16, 32])
68
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
69
+def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
70
+ """
71
+ Test the MXFP4 grouped GEMM kernel by:
72
+ 1. Creating random per-expert inputs and weights
73
+ 2. Quantizing both to MXFP4 using the CUDA kernel
74
+ 3. Running the CUTLASS grouped GEMM
75
+ 4. Comparing against BF16 reference
76
+ """
77
+ device = "cuda"
78
+ alignment = 128
79
+ # N and K must be multiples of 128 for clean swizzle layout
80
+ n_g = random.randint(1, 16) * alignment
81
+ k_g = random.randint(1, 16) * alignment
82
+
83
+ expert_offset = 0
84
+ expert_offsets_input = []
85
+ problem_sizes = []
86
+ input_list = []
87
+ weight_list = []
88
+
89
+ for g in range(num_experts):
90
+ m_g = random.randint(1, 256)
91
+ expert_offsets_input.append(expert_offset)
92
+ expert_offset += m_g
93
+ problem_sizes.append([m_g, n_g, k_g])
94
+
95
+ input_list.append(
96
+ torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
97
+ )
98
+ weight_list.append(
99
+ torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
100
+ )
101
+
102
+ input_tensor = torch.concat(input_list, dim=0) # [M_total, K]
103
+
104
+ # --- Quantize INPUTS via mxfp4_experts_quant ---
105
+ input_bs_offsets = []
106
+ tot = 0
107
+ for g in range(num_experts):
108
+ input_bs_offsets.append(tot)
109
+ tot += align(problem_sizes[g][0], 128)
110
+ input_bs_offsets.append(tot)
111
+
112
+ _inp_expert_offsets = torch.tensor(
113
+ expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
114
+ )
115
+ _inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
116
+
117
+ input_quant, input_sf = ops.mxfp4_experts_quant(
118
+ input_tensor,
119
+ _inp_expert_offsets,
120
+ _inp_bs_offsets,
121
+ num_experts,
122
+ topk=1,
123
+ )
124
+
125
+ # --- Quantize WEIGHTS via mxfp4_experts_quant ---
126
+ # Treat each expert's N weight rows as an "expert" with N tokens
127
+ weight_tensor = torch.concat(weight_list, dim=0) # [E*N, K]
128
+ weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
129
+ # N is always multiple of 128, so blockscale offsets are clean
130
+ weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
131
+
132
+ _wt_expert_offsets = torch.tensor(
133
+ weight_expert_offsets, device=device, dtype=torch.int32
134
+ )
135
+ _wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
136
+
137
+ weight_quant, weight_sf = ops.mxfp4_experts_quant(
138
+ weight_tensor,
139
+ _wt_expert_offsets,
140
+ _wt_bs_offsets,
141
+ num_experts,
142
+ topk=1,
143
+ )
144
+
145
+ # Reshape weight quantized data to [E, N, K//2]
146
+ weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
147
+
148
+ # Reshape weight scale factors to [E, N, K//32]
149
+ # The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
150
+ scales_per_row = k_g // MXFP4_BLOCK_SIZE
151
+ weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
152
+ weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
153
+
154
+ # Output
155
+ output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
156
+
157
+ _problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
158
+ _expert_offsets = torch.tensor(
159
+ expert_offsets_input, device=device, dtype=torch.int32
160
+ )
161
+ _input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
162
+
163
+ # Run the MXFP4 grouped GEMM
164
+ ops.cutlass_mxfp4_moe_mm(
165
+ output,
166
+ input_quant,
167
+ weight_quant,
168
+ input_sf,
169
+ weight_sf_3d,
170
+ _problem_sizes,
171
+ _expert_offsets,
172
+ _input_bs,
173
+ )
174
+
175
+ # Reference: BF16 matmul
176
+ ref_output = compute_ref_output(
177
+ input_tensor=input_tensor,
178
+ weight_list=weight_list,
179
+ expert_offsets=expert_offsets_input,
180
+ expert_offset=expert_offset,
181
+ num_experts=num_experts,
182
+ )
183
+
184
+ # Compare per-expert
185
+ for g in range(num_experts):
186
+ start = expert_offsets_input[g]
187
+ end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
188
+ if start == end:
189
+ continue
190
+ baseline = ref_output[start:end]
191
+ actual = output[start:end]
192
+ diff = calc_diff(actual, baseline)
193
+ print(
194
+ f"m_g={end - start} n_g={n_g} k_g={k_g} "
195
+ f"num_experts={num_experts}, "
196
+ f"out_dtype={out_dtype}, diff={diff:.5f}"
197
+ )
198
+ # FP4 quantization is very lossy (~4 bits precision)
199
+ # Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
200
+ assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
201
+
202
+
203
+@pytest.mark.skipif(
204
+ not is_sm100_supported(),
205
+ reason="mxfp4_experts_quant requires CUDA SM100",
206
+)
207
+def test_mxfp4_experts_quant_basic():
208
+ """
209
+ Basic smoke test for the MXFP4 experts quantization kernel.
210
+ """
211
+ device = "cuda"
212
+ num_experts = 4
213
+ k = 256
214
+ tokens_per_expert = 16
215
+
216
+ total_tokens = tokens_per_expert * num_experts
217
+ input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
218
+
219
+ expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
220
+ blockscale_offsets = [
221
+ align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
222
+ ]
223
+
224
+ _expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
225
+ _blockscale_offsets = torch.tensor(
226
+ blockscale_offsets, device=device, dtype=torch.int32
227
+ )
228
+
229
+ output, output_sf = ops.mxfp4_experts_quant(
230
+ input_tensor,
231
+ _expert_offsets,
232
+ _blockscale_offsets,
233
+ num_experts,
234
+ topk=1,
235
+ )
236
+
237
+ assert output.shape == (total_tokens, k // 2)
238
+ assert output.dtype == torch.uint8
239
+ assert output_sf.dtype == torch.uint8
240
+ assert output.any(), "Quantized output is all zeros"
241
+ print(
242
+ f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
243
+ )
244
+ print("PASSED")
245
+
246
+
247
+if __name__ == "__main__":
248
+ pytest.main([__file__, "-v", "-s"])
vllm/_custom_ops.py
+135
@@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm(
1150
)
1151
1152
1153
+def cutlass_mxfp4_moe_mm(
1154
+ out_tensors: torch.Tensor,
1155
+ a_tensors: torch.Tensor,
1156
+ b_tensors: torch.Tensor,
1157
+ a_scales: torch.Tensor,
1158
+ b_scales: torch.Tensor,
1159
+ problem_sizes: torch.Tensor,
1160
+ expert_offsets: torch.Tensor,
1161
+ sf_offsets: torch.Tensor,
1162
+):
1163
+ """
1164
+ An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
1165
+
1166
+ Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
1167
+ - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
1168
+ - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
1169
+ - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
1170
+ - expert_offsets/sf_offsets: expert boundary indices
1171
+ - problem_sizes: (num_experts, 3) with (M, N, K) per expert
1172
+ """
1173
+ return torch.ops._C.cutlass_mxfp4_group_mm(
1174
+ out_tensors,
1175
+ a_tensors,
1176
+ b_tensors,
1177
+ a_scales,
1178
+ b_scales,
1179
+ problem_sizes,
1180
+ expert_offsets,
1181
+ sf_offsets,
1182
+ )
1183
+
1184
+
1185
def mxfp8_experts_quant(
1186
input_tensor: torch.Tensor,
1187
problem_sizes: torch.Tensor,
@@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant(
1880
return output, output_scales
1881
1882
1883
+def mxfp4_experts_quant(
1884
+ input_tensor: torch.Tensor,
1885
+ expert_offsets: torch.Tensor,
1886
+ blockscale_offsets: torch.Tensor,
1887
+ n_experts: int,
1888
+ topk: int,
1889
+) -> tuple[torch.Tensor, torch.Tensor]:
1890
+ """
1891
+ Quantize input tensor to MXFP4 for packed MoE inputs.
1892
+ Uses 32-element blocks with E8M0 (power-of-two) scale factors.
1893
+ MXFP4 has no global scale - only block-level E8M0 scale factors.
1894
+
1895
+ Args:
1896
+ input_tensor: [m_topk, k] BF16/FP16 activations
1897
+ expert_offsets: [n_experts+1] token boundaries per expert
1898
+ blockscale_offsets: [n_experts+1] SF row boundaries per expert
1899
+ n_experts: number of experts
1900
+ topk: number of top-k experts
1901
+ Returns:
1902
+ output: [m_topk, k//2] packed E2M1 values (uint8)
1903
+ output_scales: E8M0 blockscales in swizzled layout (uint8 view)
1904
+ """
1905
+ assert not current_platform.is_rocm()
1906
+ assert input_tensor.ndim == 2
1907
+
1908
+ MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1909
+ m_numtopk, k = input_tensor.shape
1910
+
1911
+ assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
1912
+ f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1913
+ f"{MAX_TOKENS_PER_EXPERT})"
1914
+ f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
1915
+ f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
1916
+ )
1917
+ scales_k = k // 32
1918
+ padded_k = (scales_k + (4 - 1)) // 4
1919
+
1920
+ output = torch.empty(
1921
+ m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1922
+ )
1923
+ output_scales = torch.empty(
1924
+ MAX_TOKENS_PER_EXPERT * topk,
1925
+ padded_k,
1926
+ dtype=torch.int32,
1927
+ device=input_tensor.device,
1928
+ )
1929
+ torch.ops._C.mxfp4_experts_quant(
1930
+ output,
1931
+ output_scales,
1932
+ input_tensor,
1933
+ expert_offsets,
1934
+ blockscale_offsets,
1935
+ n_experts,
1936
+ )
1937
+ # E8M0 SFs are stored as uint8
1938
+ output_scales = output_scales.view(torch.uint8)
1939
+ return output, output_scales
1940
+
1941
+
1942
+def silu_and_mul_mxfp4_experts_quant(
1943
+ input_tensor: torch.Tensor,
1944
+ expert_offsets: torch.Tensor,
1945
+ blockscale_offsets: torch.Tensor,
1946
+ n_experts: int,
1947
+ topk: int,
1948
+) -> tuple[torch.Tensor, torch.Tensor]:
1949
+ """
1950
+ Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
1951
+ MXFP4 has no global scale - only block-level E8M0 scale factors.
1952
+ """
1953
+ assert not current_platform.is_rocm()
1954
+ assert input_tensor.ndim == 2
1955
+
1956
+ MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1957
+ m_numtopk, k_times_2 = input_tensor.shape
1958
+ assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
1959
+ k = k_times_2 // 2
1960
+
1961
+ assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
1962
+ scales_k = k // 32
1963
+ padded_k = (scales_k + (4 - 1)) // 4
1964
+
1965
+ output = torch.empty(
1966
+ m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1967
+ )
1968
+ output_scales = torch.empty(
1969
+ MAX_TOKENS_PER_EXPERT * topk,
1970
+ padded_k,
1971
+ dtype=torch.int32,
1972
+ device=input_tensor.device,
1973
+ )
1974
+ torch.ops._C.silu_and_mul_mxfp4_experts_quant(
1975
+ output,
1976
+ output_scales,
1977
+ input_tensor,
1978
+ expert_offsets,
1979
+ blockscale_offsets,
1980
+ n_experts,
1981
+ )
1982
+ output_scales = output_scales.view(torch.uint8)
1983
+ return output, output_scales
1984
+
1985
+
1986
# fp8
1987
def scaled_fp8_quant(
1988
input: torch.Tensor,
vllm/model_executor/layers/fused_moe/config.py
+19
@@ -762,6 +762,25 @@ def nvfp4_moe_quant_config(
762
)
763
764
765
+def mxfp4_moe_quant_config(
766
+ w1_scale: torch.Tensor,
767
+ w2_scale: torch.Tensor,
768
+) -> FusedMoEQuantConfig:
769
+ """
770
+ Construct a quant config for MXFP4 x MXFP4 MoE.
771
+ MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
772
+ separate alphas / global activation scales in this config.
773
+ """
774
+ return FusedMoEQuantConfig.make(
775
+ "mxfp4",
776
+ w1_scale=w1_scale,
777
+ w2_scale=w2_scale,
778
+ per_act_token_quant=False,
779
+ per_out_ch_quant=False,
780
+ block_shape=None,
781
+ )
782
+
783
+
784
def nvfp4_w4a16_moe_quant_config(
785
g1_alphas: torch.Tensor,
786
g2_alphas: torch.Tensor,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+295
@@ -36,6 +36,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
kFp8DynamicTokenSym,
37
kFp8StaticChannelSym,
38
kFp8StaticTensorSym,
39
+ kMxfp4Dynamic,
40
+ kMxfp4Static,
41
kNvfp4Dynamic,
42
kNvfp4Static,
43
)
@@ -795,6 +797,299 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
797
)
798
799
800
+def run_cutlass_moe_mxfp4(
801
+ output: torch.Tensor,
802
+ a: torch.Tensor,
803
+ w1_fp4: torch.Tensor,
804
+ w1_blockscale: torch.Tensor,
805
+ w2_fp4: torch.Tensor,
806
+ w2_blockscale: torch.Tensor,
807
+ topk_weights: torch.Tensor,
808
+ topk_ids: torch.Tensor,
809
+ activation: MoEActivation,
810
+ workspace13: torch.Tensor,
811
+ workspace2: torch.Tensor,
812
+ m: int,
813
+ n: int,
814
+ k: int,
815
+ e: int,
816
+ device: torch.device,
817
+ apply_router_weight_on_input: bool = False,
818
+) -> None:
819
+ """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
820
+ is_gated = activation.is_gated
821
+ w1_n = n * 2 if is_gated else n
822
+
823
+ assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
824
+ assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
825
+ assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
826
+ assert (
827
+ w1_fp4.ndim == 3
828
+ and w2_fp4.ndim == 3
829
+ and w1_blockscale.ndim == 3
830
+ and w2_blockscale.ndim == 3
831
+ ), "All Weights must be of rank 3 for cutlass_moe_mxfp4"
832
+ m_a, k_a = a.shape
833
+ e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape
834
+ e_w2, k_w2, half_n_w2 = w2_fp4.shape
835
+
836
+ assert e_w1 == e_w2 and e_w1 == e
837
+ assert k_a == half_k_w1 * 2 and k == k_w2
838
+ assert w1_n_actual == w1_n and half_n_w2 * 2 == n
839
+ assert m == m_a
840
+ assert 2 * half_k_w1 == k_w2
841
+ assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
842
+ assert topk_weights.size(0) == m and topk_ids.size(0) == m
843
+
844
+ topk = topk_ids.size(1)
845
+ out_dtype = a.dtype
846
+ num_topk = topk_ids.size(1)
847
+
848
+ expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
849
+ blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
850
+ problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
851
+ problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
852
+
853
+ a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
854
+ c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
855
+
856
+ if apply_router_weight_on_input:
857
+ assert num_topk == 1, (
858
+ "apply_router_weight_on_input is only implemented for topk=1"
859
+ )
860
+ a.mul_(topk_weights.to(out_dtype))
861
+
862
+ ops.get_cutlass_moe_mm_data(
863
+ topk_ids,
864
+ expert_offsets,
865
+ problem_sizes1,
866
+ problem_sizes2,
867
+ a_map,
868
+ c_map,
869
+ e,
870
+ n,
871
+ k,
872
+ blockscale_offsets,
873
+ is_gated=is_gated,
874
+ )
875
+
876
+ a = ops.shuffle_rows(a, a_map)
877
+ rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant(
878
+ a,
879
+ expert_offsets,
880
+ blockscale_offsets,
881
+ e,
882
+ num_topk,
883
+ )
884
+ c1 = _resize_cache(workspace13, (m * topk, w1_n))
885
+ c2 = _resize_cache(workspace2, (m * topk, n))
886
+ c3 = _resize_cache(workspace13, (m * topk, k))
887
+
888
+ ops.cutlass_mxfp4_moe_mm(
889
+ c1,
890
+ rep_a_fp4,
891
+ w1_fp4,
892
+ rep_a_blockscale,
893
+ w1_blockscale,
894
+ problem_sizes1,
895
+ expert_offsets[:-1],
896
+ blockscale_offsets[:-1],
897
+ )
898
+ del rep_a_fp4, rep_a_blockscale
899
+ if activation == MoEActivation.SILU:
900
+ int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
901
+ c1, expert_offsets, blockscale_offsets, e, num_topk
902
+ )
903
+ else:
904
+ apply_moe_activation(activation, c2, c1)
905
+ int_fp4, int_blockscale = ops.mxfp4_experts_quant(
906
+ c2, expert_offsets, blockscale_offsets, e, num_topk
907
+ )
908
+
909
+ ops.cutlass_mxfp4_moe_mm(
910
+ c3,
911
+ int_fp4,
912
+ w2_fp4,
913
+ int_blockscale,
914
+ w2_blockscale,
915
+ problem_sizes2,
916
+ expert_offsets[:-1],
917
+ blockscale_offsets[:-1],
918
+ )
919
+ del int_fp4, int_blockscale
920
+
921
+ c3 = ops.shuffle_rows(c3, c_map)
922
+
923
+ assert output.dtype == out_dtype
924
+ if not apply_router_weight_on_input:
925
+ output.copy_(
926
+ (
927
+ c3.view(m, num_topk, k)
928
+ * topk_weights.view(m, num_topk, 1).to(out_dtype)
929
+ ).sum(dim=1),
930
+ non_blocking=True,
931
+ )
932
+ else:
933
+ output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
934
+ return
935
+
936
+
937
+def swizzle_mxfp4_scales(
938
+ scales: torch.Tensor,
939
+ N: int,
940
+ K: int,
941
+) -> torch.Tensor:
942
+ """Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout.
943
+
944
+ CUTLASS expects MX scale factors in a tiled layout:
945
+ [numMTiles, numKTiles, 32, 4, 4]
946
+ where numMTiles = ceil(N/128), numKTiles = ceil(K/128),
947
+ and the inner dimensions correspond to the swizzle pattern:
948
+ mTileIdx = mIdx / 128
949
+ outerMIdx = mIdx % 32
950
+ innerMIdx = (mIdx / 32) % 4
951
+ kTileIdx = kIdx / 4
952
+ innerKIdx = kIdx % 4
953
+ with kIdx = col_in_scale_space (i.e., index into K//32).
954
+ """
955
+ assert scales.dtype == torch.uint8
956
+ num_scale_cols = K // 32 # number of E8M0 scale values per row
957
+
958
+ num_m_tiles = (N + 127) // 128
959
+ num_k_tiles = (num_scale_cols + 3) // 4
960
+
961
+ # Pad N to multiple of 128 and scale_cols to multiple of 4
962
+ padded_N = num_m_tiles * 128
963
+ padded_scale_cols = num_k_tiles * 4
964
+
965
+ # Start with flat scales, pad if needed
966
+ padded = torch.zeros(
967
+ padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device
968
+ )
969
+ padded[:N, :num_scale_cols] = scales
970
+
971
+ # Reshape to tile structure:
972
+ # [numMTiles, 4, 32, numKTiles, 4]
973
+ # mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx
974
+ tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4)
975
+ # Permute to [numMTiles, numKTiles, 32, 4, 4]
976
+ # (outerMIdx, innerMIdx, innerKIdx)
977
+ tiled = tiled.permute(0, 3, 2, 1, 4).contiguous()
978
+ return tiled.reshape(-1)
979
+
980
+
981
+class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
982
+ """CUTLASS MXFP4 x MXFP4 fused MoE expert implementation."""
983
+
984
+ @property
985
+ def expects_unquantized_inputs(self) -> bool:
986
+ return True
987
+
988
+ @staticmethod
989
+ def _supports_current_device() -> bool:
990
+ p = current_platform
991
+ return p.is_cuda() and p.is_device_capability_family(100)
992
+
993
+ @staticmethod
994
+ def _supports_no_act_and_mul() -> bool:
995
+ return True
996
+
997
+ @staticmethod
998
+ def _supports_quant_scheme(
999
+ weight_key: QuantKey | None,
1000
+ activation_key: QuantKey | None,
1001
+ ) -> bool:
1002
+ return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
1003
+
1004
+ @staticmethod
1005
+ def _supports_activation(activation: MoEActivation) -> bool:
1006
+ return activation in [
1007
+ MoEActivation.SILU,
1008
+ MoEActivation.GELU,
1009
+ MoEActivation.SWIGLUOAI,
1010
+ MoEActivation.SWIGLUSTEP,
1011
+ MoEActivation.SILU_NO_MUL,
1012
+ MoEActivation.GELU_NO_MUL,
1013
+ MoEActivation.RELU2_NO_MUL,
1014
+ ]
1015
+
1016
+ @staticmethod
1017
+ def _supports_parallel_config(
1018
+ moe_parallel_config: FusedMoEParallelConfig,
1019
+ ) -> bool:
1020
+ return moe_parallel_config.ep_size == 1
1021
+
1022
+ @staticmethod
1023
+ def activation_format() -> mk.FusedMoEActivationFormat:
1024
+ return mk.FusedMoEActivationFormat.Standard
1025
+
1026
+ def supports_expert_map(self) -> bool:
1027
+ return False
1028
+
1029
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1030
+ return TopKWeightAndReduceNoOP()
1031
+
1032
+ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
1033
+ return act_dtype
1034
+
1035
+ def workspace_shapes(
1036
+ self,
1037
+ M: int,
1038
+ N: int,
1039
+ K: int,
1040
+ topk: int,
1041
+ global_num_experts: int,
1042
+ local_num_experts: int,
1043
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1044
+ activation: MoEActivation,
1045
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1046
+ workspace1 = (M * topk, max(2 * N, K))
1047
+ workspace2 = (M * topk, N)
1048
+ output = (M, K)
1049
+ return (workspace1, workspace2, output)
1050
+
1051
+ def apply(
1052
+ self,
1053
+ output: torch.Tensor,
1054
+ hidden_states: torch.Tensor,
1055
+ w1: torch.Tensor,
1056
+ w2: torch.Tensor,
1057
+ topk_weights: torch.Tensor,
1058
+ topk_ids: torch.Tensor,
1059
+ activation: MoEActivation,
1060
+ global_num_experts: int,
1061
+ expert_map: torch.Tensor | None,
1062
+ a1q_scale: torch.Tensor | None,
1063
+ a2_scale: torch.Tensor | None,
1064
+ workspace13: torch.Tensor | None,
1065
+ workspace2: torch.Tensor | None,
1066
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1067
+ apply_router_weight_on_input: bool,
1068
+ ):
1069
+ e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
1070
+ n = w2.shape[2] * 2
1071
+
1072
+ run_cutlass_moe_mxfp4(
1073
+ output=output,
1074
+ a=hidden_states,
1075
+ w1_fp4=w1,
1076
+ w1_blockscale=self.w1_scale,
1077
+ w2_fp4=w2,
1078
+ w2_blockscale=self.w2_scale,
1079
+ topk_weights=topk_weights,
1080
+ topk_ids=topk_ids,
1081
+ activation=activation,
1082
+ workspace13=workspace13,
1083
+ workspace2=workspace2,
1084
+ m=m,
1085
+ n=n,
1086
+ k=k,
1087
+ e=e,
1088
+ device=hidden_states.device,
1089
+ apply_router_weight_on_input=apply_router_weight_on_input,
1090
+ )
1091
+
1092
+
1093
# W4A8
1094
def run_cutlass_moe_w4a8_fp8(
1095
output: torch.Tensor,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+65
-13
@@ -4,6 +4,7 @@
4
5
import torch
6
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.fused_moe import (
10
FusedMoE,
@@ -11,6 +12,10 @@ from vllm.model_executor.layers.fused_moe import (
12
)
13
from vllm.model_executor.layers.fused_moe.config import (
14
FusedMoEQuantConfig,
15
+ mxfp4_moe_quant_config,
16
+)
17
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
18
+ CutlassExpertsMxfp4,
19
)
20
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
21
MarlinExperts,
@@ -36,7 +41,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
41
super().__init__(moe)
42
self.group_size = 32
43
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
39
- self.experts_cls = MarlinExperts
44
+ self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
45
+ self.experts_cls: type[mk.FusedMoEExperts]
46
+ if self.use_cutlass_mxfp4:
47
+ logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local")
48
+ self.experts_cls = CutlassExpertsMxfp4
49
+ else:
50
+ logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local")
51
+ self.experts_cls = MarlinExperts
52
53
def create_weights(
54
self,
@@ -109,11 +121,19 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
121
def get_fused_moe_quant_config(
122
self, layer: torch.nn.Module
123
) -> FusedMoEQuantConfig | None:
112
- return make_mxfp4_moe_quant_config(
113
- mxfp4_backend=self.mxfp4_backend,
114
- w1_scale=layer.w13_weight_scale,
115
- w2_scale=layer.w2_weight_scale,
116
- )
124
+ if self.use_cutlass_mxfp4:
125
+ # W4A4: both weights and activations quantized to MXFP4
126
+ return mxfp4_moe_quant_config(
127
+ w1_scale=layer.w13_weight_scale,
128
+ w2_scale=layer.w2_weight_scale,
129
+ )
130
+ else:
131
+ # W4A16: weight-only via Marlin
132
+ return make_mxfp4_moe_quant_config(
133
+ mxfp4_backend=self.mxfp4_backend,
134
+ w1_scale=layer.w13_weight_scale,
135
+ w2_scale=layer.w2_weight_scale,
136
+ )
137
138
def process_weights_after_loading(self, layer: FusedMoE) -> None:
139
layer.w13_weight = torch.nn.Parameter(
@@ -126,13 +146,45 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
146
)
147
delattr(layer, "w2_weight_packed")
148
129
- logger.warning_once(
130
- "Your GPU does not have native support for FP4 computation but "
131
- "FP4 quantization is being used. Weight-only FP4 compression "
132
- "will be used leveraging the Marlin kernel. This may degrade "
133
- "performance for compute-heavy workloads."
134
- )
135
- prepare_moe_fp4_layer_for_marlin(layer)
149
+ if self.use_cutlass_mxfp4:
150
+ # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
151
+ # to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
152
+ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
153
+ swizzle_mxfp4_scales,
154
+ )
155
+
156
+ E = layer.w13_weight_scale.shape[0]
157
+ w13_N = layer.w13_weight_scale.shape[1]
158
+ w13_scale_K = layer.w13_weight_scale.shape[2]
159
+ w13_K = w13_scale_K * 32
160
+
161
+ w2_M = layer.w2_weight_scale.shape[1]
162
+ w2_scale_N = layer.w2_weight_scale.shape[2]
163
+ w2_N = w2_scale_N * 32
164
+
165
+ swizzled_w13 = []
166
+ swizzled_w2 = []
167
+ for e_idx in range(E):
168
+ s13 = layer.w13_weight_scale[e_idx]
169
+ sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
170
+ swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
171
+ s2 = layer.w2_weight_scale[e_idx]
172
+ sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
173
+ swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
174
+ layer.w13_weight_scale = torch.nn.Parameter(
175
+ torch.stack(swizzled_w13), requires_grad=False
176
+ )
177
+ layer.w2_weight_scale = torch.nn.Parameter(
178
+ torch.stack(swizzled_w2), requires_grad=False
179
+ )
180
+ else:
181
+ logger.warning_once(
182
+ "Your GPU does not have native support for FP4 computation "
183
+ "but FP4 quantization is being used. Weight-only FP4 "
184
+ "compression will be used leveraging the Marlin kernel. "
185
+ "This may degrade performance for compute-heavy workloads."
186
+ )
187
+ prepare_moe_fp4_layer_for_marlin(layer)
188
189
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
190
if self.moe_quant_config is not None: