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 11, 20:00 UTC
[Kernel]Support W4A8 Grouped GEMM on Hopper (#29691)
Signed-off-by: czhu-cohere <conway.zhu@cohere.com>
czhu-cohere committed
Dec 8, 2025 at 22:29 UTC
f6227c22ab8976a24913122874c24624102da1b4
22 files changed
+2045
-101
CMakeLists.txt
+4
-1
@@ -874,7 +874,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
874
cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
875
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
876
set(SRCS
877
- "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu")
877
+ "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu"
878
+ "csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu"
879
+ "csrc/quantization/cutlass_w4a8/w4a8_utils.cu"
880
+ )
881
882
set_gencode_flags_for_srcs(
883
SRCS "${SRCS}"
csrc/ops.h
+2
-1
@@ -262,7 +262,8 @@ void get_cutlass_moe_mm_data(
262
void get_cutlass_moe_mm_problem_sizes(
263
const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
264
torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
265
- const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets);
265
+ const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
266
+ std::optional<bool> force_swap_ab = std::nullopt);
267
268
void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
269
torch::Tensor& problem_sizes1,
csrc/quantization/cutlass_w4a8/get_group_starts.cuh
new
+104
@@ -0,0 +1,104 @@
1
+// see csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh
2
+#pragma once
3
+
4
+#include <cuda.h>
5
+#include <torch/all.h>
6
+#include <c10/cuda/CUDAStream.h>
7
+
8
+#include "core/scalar_type.hpp"
9
+#include "cutlass/bfloat16.h"
10
+#include "cutlass/float8.h"
11
+
12
+// ElementB is int32 (packed int4)
13
+// ElementGroupScale is cutlass::Array<cutlass::float_e4m3_t, 8> (packed fp8)
14
+template <typename ElementA, typename ElementB, typename ElementC,
15
+ typename ElementAccumulator, typename ElementGroupScale>
16
+__global__ void get_group_gemm_starts(
17
+ int64_t* expert_offsets, ElementA** a_offsets, ElementB** b_offsets,
18
+ ElementC** out_offsets, ElementAccumulator** a_scales_offsets,
19
+ ElementAccumulator** b_scales_offsets,
20
+ ElementGroupScale** b_group_scales_offsets, ElementA* a_base_as_int,
21
+ ElementB* b_base_as_int, ElementC* out_base_as_int,
22
+ ElementAccumulator* a_scales_base_as_int,
23
+ ElementAccumulator* b_scales_base_as_int,
24
+ ElementGroupScale* b_group_scales_base_as_int, int64_t n, int64_t k,
25
+ int64_t scale_k) {
26
+ int expert_id = threadIdx.x;
27
+
28
+ int64_t expert_offset = expert_offsets[expert_id];
29
+
30
+ // same as w8a8
31
+ a_offsets[expert_id] = a_base_as_int + expert_offset * k;
32
+ out_offsets[expert_id] = out_base_as_int + expert_offset * n;
33
+ a_scales_offsets[expert_id] = a_scales_base_as_int + expert_offset;
34
+ b_scales_offsets[expert_id] = b_scales_base_as_int + (n * expert_id);
35
+
36
+ // w4a8 specific
37
+ constexpr int pack_factor = 8; // pack 8 int4 into int32
38
+ b_offsets[expert_id] = b_base_as_int + (expert_id * k * n / pack_factor);
39
+ b_group_scales_offsets[expert_id] =
40
+ b_group_scales_base_as_int + (expert_id * scale_k * n);
41
+}
42
+
43
+#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
44
+ else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
45
+ get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \
46
+ cutlass::Array<cutlass::float_e4m3_t, 8>> \
47
+ <<<1, num_experts, 0, stream>>>( \
48
+ static_cast<int64_t*>(expert_offsets.data_ptr()), \
49
+ static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \
50
+ static_cast<int32_t**>(b_ptrs.data_ptr()), \
51
+ static_cast<C_TYPE**>(out_ptrs.data_ptr()), \
52
+ static_cast<float**>(a_scales_ptrs.data_ptr()), \
53
+ static_cast<float**>(b_scales_ptrs.data_ptr()), \
54
+ static_cast<cutlass::Array<cutlass::float_e4m3_t, 8>**>( \
55
+ b_group_scales_ptrs.data_ptr()), \
56
+ static_cast<cutlass::float_e4m3_t*>(a_tensors.data_ptr()), \
57
+ static_cast<int32_t*>(b_tensors.data_ptr()), \
58
+ static_cast<C_TYPE*>(out_tensors.data_ptr()), \
59
+ static_cast<float*>(a_scales.data_ptr()), \
60
+ static_cast<float*>(b_scales.data_ptr()), \
61
+ static_cast<cutlass::Array<cutlass::float_e4m3_t, 8>*>( \
62
+ b_group_scales.data_ptr()), \
63
+ n, k, scale_k); \
64
+ }
65
+
66
+namespace {
67
+
68
+void run_get_group_gemm_starts(
69
+ torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
70
+ torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
71
+ torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
72
+ torch::Tensor& b_group_scales_ptrs, torch::Tensor const& a_tensors,
73
+ torch::Tensor const& b_tensors, torch::Tensor& out_tensors,
74
+ torch::Tensor const& a_scales, torch::Tensor const& b_scales,
75
+ torch::Tensor const& b_group_scales, const int64_t b_group_size) {
76
+ TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
77
+ TORCH_CHECK(b_tensors.dtype() == torch::kInt32); // int4 8x packed into int32
78
+ TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
79
+ TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
80
+ TORCH_CHECK(b_group_scales.dtype() ==
81
+ torch::kFloat8_e4m3fn); // the underlying torch type is e4m3
82
+ TORCH_CHECK(out_tensors.dtype() ==
83
+ torch::kBFloat16); // only support bf16 for now
84
+ // expect int64_t to avoid overflow during offset calculations
85
+ TORCH_CHECK(expert_offsets.dtype() == torch::kInt64);
86
+
87
+ int num_experts = static_cast<int>(expert_offsets.size(0));
88
+ // logical k, n
89
+ int64_t n = out_tensors.size(1);
90
+ int64_t k = a_tensors.size(1);
91
+ int64_t scale_k = cutlass::ceil_div(k, b_group_size);
92
+
93
+ auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
94
+
95
+ if (false) {
96
+ }
97
+ __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
98
+ __CALL_GET_STARTS_KERNEL(torch::kFloat16, half)
99
+ else {
100
+ TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
101
+ }
102
+}
103
+
104
+} // namespace
\ No newline at end of file
csrc/quantization/cutlass_w4a8/w4a8_grouped_mm_entry.cu
new
+483
@@ -0,0 +1,483 @@
1
+#include <vector>
2
+#include <tuple>
3
+
4
+#include "cutlass/cutlass.h"
5
+
6
+#include "cute/tensor.hpp"
7
+#include "cutlass/gemm/dispatch_policy.hpp"
8
+#include "cutlass/gemm/group_array_problem_shape.hpp"
9
+#include "cutlass/gemm/collective/collective_builder.hpp"
10
+#include "cutlass/epilogue/collective/collective_builder.hpp"
11
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
12
+
13
+#include "cutlass/util/packed_stride.hpp"
14
+#include "cutlass/util/mixed_dtype_utils.hpp"
15
+
16
+// vllm includes
17
+#include <ATen/cuda/CUDAContext.h>
18
+#include <c10/cuda/CUDAGuard.h>
19
+#include <torch/all.h>
20
+#include "cutlass_extensions/torch_utils.hpp"
21
+#include "cutlass_extensions/common.hpp"
22
+
23
+#include "core/registration.h"
24
+#include "get_group_starts.cuh"
25
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
26
+#include "w4a8_utils.cuh"
27
+
28
+namespace vllm::cutlass_w4a8_moe {
29
+
30
+using namespace cute;
31
+
32
+// -------------------------------------------------------------------------------------
33
+// Static configuration shared across all instantiations
34
+// -------------------------------------------------------------------------------------
35
+using ProblemShape =
36
+ cutlass::gemm::GroupProblemShape<Shape<int, int, int>>; // <M,N,K> per
37
+ // group
38
+using MmaType = cutlass::float_e4m3_t;
39
+using QuantType = cutlass::int4b_t;
40
+
41
+constexpr int TileShapeK = 128 * 8 / sizeof_bits<MmaType>::value;
42
+static int constexpr PackFactor = 8; // 8 int4 packed into int32
43
+
44
+// A matrix configuration
45
+using ElementA = MmaType;
46
+using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand
47
+constexpr int AlignmentA =
48
+ 128 /
49
+ cutlass::sizeof_bits<ElementA>::value; // Alignment of A matrix in units of
50
+ // elements (up to 16 bytes)
51
+
52
+// B matrix configuration
53
+using ElementB = QuantType; // Element type for B matrix operand
54
+using LayoutB =
55
+ cutlass::layout::ColumnMajor; // Layout type for B matrix operand
56
+constexpr int AlignmentB =
57
+ 128 / cutlass::sizeof_bits<
58
+ ElementB>::value; // Memory access granularity/alignment of B
59
+ // matrix in units of elements (up to 16 bytes)
60
+
61
+// This example manually swaps and transposes, so keep transpose of input
62
+// layouts
63
+using LayoutA_Transpose =
64
+ typename cutlass::layout::LayoutTranspose<LayoutA>::type;
65
+using LayoutB_Transpose =
66
+ typename cutlass::layout::LayoutTranspose<LayoutB>::type;
67
+
68
+// Need to pass a pointer type to make the 3rd dimension of Stride be _0
69
+using StrideA =
70
+ cute::remove_pointer_t<cutlass::detail::TagToStrideA_t<LayoutA*>>;
71
+using StrideB =
72
+ cute::remove_pointer_t<cutlass::detail::TagToStrideB_t<LayoutB*>>;
73
+
74
+// Define the CuTe layout for reoredered quantized tensor B
75
+// LayoutAtomQuant places values that will be read by the same thread in
76
+// contiguous locations in global memory. It specifies the reordering within a
77
+// single warp's fragment
78
+using LayoutAtomQuant =
79
+ decltype(cutlass::compute_memory_reordering_atom<MmaType>());
80
+using LayoutB_Reordered = decltype(cute::tile_to_shape(
81
+ LayoutAtomQuant{}, Layout<Shape<int, int, Int<1>>, StrideB>{}));
82
+
83
+using ElementScale = cutlass::float_e4m3_t;
84
+using LayoutScale = cutlass::layout::RowMajor;
85
+
86
+// C/D matrix configuration
87
+using ElementC =
88
+ cutlass::bfloat16_t; // Element type for C and D matrix operands
89
+using LayoutC =
90
+ cutlass::layout::RowMajor; // Layout type for C and D matrix operands
91
+constexpr int AlignmentC =
92
+ 128 / cutlass::sizeof_bits<
93
+ ElementC>::value; // Memory access granularity/alignment of C
94
+ // matrix in units of elements (up to 16 bytes)
95
+
96
+// D matrix configuration
97
+using ElementD = ElementC;
98
+using LayoutD = LayoutC;
99
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
100
+
101
+// Core kernel configurations
102
+using ElementAccumulator = float; // Element type for internal accumulation
103
+using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that
104
+ // supports the intended feature
105
+using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag
106
+using StageCountType =
107
+ cutlass::gemm::collective::StageCountAuto; // Stage count maximized based
108
+ // on the tile size
109
+
110
+// per-channel and per-token scales for epilogue
111
+using ElementSChannel = float;
112
+
113
+template <class TileShape_MN, class ClusterShape_MNK, class KernelSchedule,
114
+ class EpilogueSchedule>
115
+struct W4A8GroupedGemmKernel {
116
+ using TileShape =
117
+ decltype(cute::append(TileShape_MN{}, cute::Int<TileShapeK>{}));
118
+ using ClusterShape = ClusterShape_MNK;
119
+
120
+ // per-channel, per-token scales epilogue
121
+ using ChTokScalesEpilogue =
122
+ typename vllm::c3x::ScaledEpilogueArray<ElementAccumulator, ElementD,
123
+ TileShape>;
124
+ using EVTCompute = typename ChTokScalesEpilogue::EVTCompute;
125
+ using CollectiveEpilogue =
126
+ typename cutlass::epilogue::collective::CollectiveBuilder<
127
+ ArchTag, OperatorClass, TileShape, ClusterShape,
128
+ cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
129
+ ElementSChannel, ElementC,
130
+ typename cutlass::layout::LayoutTranspose<LayoutC>::type*, AlignmentC,
131
+ ElementD, typename cutlass::layout::LayoutTranspose<LayoutD>::type*,
132
+ AlignmentD, EpilogueSchedule, EVTCompute>::CollectiveOp;
133
+
134
+ // =========================================================== MIXED INPUT
135
+ // WITH SCALES
136
+ // ===========================================================================
137
+ // The Scale information must get paired with the operand that will be scaled.
138
+ // In this example, B is scaled so we make a tuple of B's information and the
139
+ // scale information.
140
+ using CollectiveMainloopShuffled =
141
+ typename cutlass::gemm::collective::CollectiveBuilder<
142
+ ArchTag, OperatorClass,
143
+ cute::tuple<ElementB, cutlass::Array<ElementScale, 8>>,
144
+ LayoutB_Reordered*, AlignmentB, ElementA, LayoutA_Transpose*,
145
+ AlignmentA, ElementAccumulator, TileShape, ClusterShape,
146
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
147
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
148
+ KernelSchedule>::CollectiveOp;
149
+
150
+ using GemmKernelShuffled = cutlass::gemm::kernel::GemmUniversal<
151
+ ProblemShape, CollectiveMainloopShuffled, CollectiveEpilogue>;
152
+
153
+ using GemmShuffled =
154
+ cutlass::gemm::device::GemmUniversalAdapter<GemmKernelShuffled>;
155
+
156
+ using StrideC = typename GemmKernelShuffled::InternalStrideC;
157
+ using StrideD = typename GemmKernelShuffled::InternalStrideD;
158
+
159
+ using StrideC_ref = cutlass::detail::TagToStrideC_t<LayoutC>;
160
+ using StrideD_ref = cutlass::detail::TagToStrideC_t<LayoutD>;
161
+ using StrideS = typename CollectiveMainloopShuffled::StrideScale;
162
+ using StrideS_ref = cutlass::detail::TagToStrideB_t<LayoutScale>;
163
+
164
+ // static asserts for passing in strides/layouts
165
+ // pack to 2x int64
166
+ static_assert(sizeof(StrideS) == 2 * sizeof(int64_t));
167
+ // pack to 3xint32,
168
+ static_assert(sizeof(LayoutB_Reordered) % sizeof(int32_t) == 0,
169
+ "LayoutB_Reordered size must be divisible by 4 bytes");
170
+
171
+ static void grouped_mm(
172
+ torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
173
+ const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
174
+ const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
175
+ const int64_t b_group_size, const torch::Tensor& expert_offsets,
176
+ const torch::Tensor& problem_sizes_torch, const torch::Tensor& a_strides,
177
+ const torch::Tensor& b_strides, const torch::Tensor& c_strides,
178
+ const torch::Tensor& group_scale_strides) {
179
+ auto device = a_tensors.device();
180
+ auto device_id = device.index();
181
+ const at::cuda::OptionalCUDAGuard device_guard(device);
182
+ auto stream = at::cuda::getCurrentCUDAStream(device_id);
183
+
184
+ int num_experts = static_cast<int>(expert_offsets.size(0));
185
+ int n = static_cast<int>(b_tensors.size(1));
186
+ int k = static_cast<int>(b_tensors.size(2)) * PackFactor;
187
+
188
+ auto options_int =
189
+ torch::TensorOptions().dtype(torch::kInt64).device(device);
190
+ torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
191
+ torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
192
+ torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
193
+ torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
194
+ torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
195
+ torch::Tensor b_group_scales_ptrs = torch::empty(num_experts, options_int);
196
+
197
+ // get the correct offsets to pass to gemm
198
+ run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
199
+ a_scales_ptrs, b_scales_ptrs, b_group_scales_ptrs,
200
+ a_tensors, b_tensors, out_tensors, a_scales,
201
+ b_scales, b_group_scales, b_group_size);
202
+
203
+ // construct args
204
+ using Args = typename GemmShuffled::Arguments;
205
+ using MainloopArguments = typename GemmKernelShuffled::MainloopArguments;
206
+ using EpilogueArguments = typename GemmKernelShuffled::EpilogueArguments;
207
+ Args arguments;
208
+
209
+ ProblemShape::UnderlyingProblemShape* problem_sizes_as_shapes =
210
+ static_cast<ProblemShape::UnderlyingProblemShape*>(
211
+ problem_sizes_torch.data_ptr());
212
+ ProblemShape prob_shape{num_experts, problem_sizes_as_shapes, nullptr};
213
+
214
+ // SwapAB so B operands come first
215
+ MainloopArguments mainloop_arguments{
216
+ static_cast<const QuantType**>(b_ptrs.data_ptr()),
217
+ static_cast<LayoutB_Reordered*>(b_strides.data_ptr()),
218
+ static_cast<const MmaType**>(a_ptrs.data_ptr()),
219
+ static_cast<StrideA*>(a_strides.data_ptr()),
220
+ static_cast<const cutlass::Array<ElementScale, 8>**>(
221
+ b_group_scales_ptrs.data_ptr()),
222
+ static_cast<StrideS*>(group_scale_strides.data_ptr()),
223
+ static_cast<int>(b_group_size)};
224
+
225
+ EpilogueArguments epilogue_arguments{
226
+ // since we are doing SwapAB the channel scales comes first, then token
227
+ // scales
228
+ ChTokScalesEpilogue::prepare_args( // see ScaledEpilogueArray
229
+ static_cast<const ElementAccumulator**>(
230
+ b_scales_ptrs.data_ptr()), // per-channel
231
+ static_cast<const ElementAccumulator**>(
232
+ a_scales_ptrs.data_ptr()), // per-token
233
+ true, true),
234
+ nullptr, // C
235
+ static_cast<StrideC*>(c_strides.data_ptr()), // C
236
+ static_cast<ElementD**>(out_ptrs.data_ptr()), // D
237
+ static_cast<StrideC*>(c_strides.data_ptr()) // D
238
+ };
239
+
240
+ static const cutlass::KernelHardwareInfo hw_info{
241
+ device_id,
242
+ cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
243
+ device_id)};
244
+
245
+ arguments = Args{cutlass::gemm::GemmUniversalMode::kGrouped, prob_shape,
246
+ mainloop_arguments, epilogue_arguments, hw_info};
247
+
248
+ // Allocate workspace
249
+ size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
250
+ torch::Tensor workspace =
251
+ torch::empty(workspace_size,
252
+ torch::TensorOptions().dtype(torch::kU8).device(device));
253
+
254
+ // Run GEMM
255
+ GemmShuffled gemm;
256
+ CUTLASS_CHECK(gemm.can_implement(arguments));
257
+ CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
258
+ CUTLASS_CHECK(gemm.run(stream));
259
+ }
260
+};
261
+
262
+// ----------------------------------------------------------------------------
263
+// Kernel instantiations and dispatch logic
264
+// ----------------------------------------------------------------------------
265
+using Coop = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative;
266
+using CoopEpi = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative;
267
+
268
+// Kernel_TileShape_ClusterShape_Schedule
269
+using Kernel_128x16_1x1x1_Coop =
270
+ W4A8GroupedGemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>, Coop, CoopEpi>;
271
+using Kernel_128x16_2x1x1_Coop =
272
+ W4A8GroupedGemmKernel<Shape<_128, _16>, Shape<_2, _1, _1>, Coop, CoopEpi>;
273
+
274
+using Kernel_256x16_1x1x1_Coop =
275
+ W4A8GroupedGemmKernel<Shape<_256, _16>, Shape<_1, _1, _1>, Coop, CoopEpi>;
276
+using Kernel_256x16_2x1x1_Coop =
277
+ W4A8GroupedGemmKernel<Shape<_256, _16>, Shape<_2, _1, _1>, Coop, CoopEpi>;
278
+
279
+using Kernel_256x32_1x1x1_Coop =
280
+ W4A8GroupedGemmKernel<Shape<_256, _32>, Shape<_1, _1, _1>, Coop, CoopEpi>;
281
+using Kernel_256x32_2x1x1_Coop =
282
+ W4A8GroupedGemmKernel<Shape<_256, _32>, Shape<_2, _1, _1>, Coop, CoopEpi>;
283
+
284
+using Kernel_256x64_1x1x1_Coop =
285
+ W4A8GroupedGemmKernel<Shape<_256, _64>, Shape<_1, _1, _1>, Coop, CoopEpi>;
286
+using Kernel_256x64_2x1x1_Coop =
287
+ W4A8GroupedGemmKernel<Shape<_256, _64>, Shape<_2, _1, _1>, Coop, CoopEpi>;
288
+
289
+using Kernel_256x128_1x1x1_Coop =
290
+ W4A8GroupedGemmKernel<Shape<_256, _128>, Shape<_1, _1, _1>, Coop, CoopEpi>;
291
+using Kernel_256x128_2x1x1_Coop =
292
+ W4A8GroupedGemmKernel<Shape<_256, _128>, Shape<_2, _1, _1>, Coop, CoopEpi>;
293
+
294
+using Kernel_128x256_2x1x1_Coop =
295
+ W4A8GroupedGemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>, Coop, CoopEpi>;
296
+
297
+void mm_dispatch(
298
+ torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
299
+ const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
300
+ const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
301
+ const int64_t b_group_size, const torch::Tensor& expert_offsets,
302
+ const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
303
+ const torch::Tensor& b_strides, const torch::Tensor& c_strides,
304
+ const torch::Tensor& group_scale_strides, const std::string& schedule) {
305
+ if (schedule == "Kernel_128x16_1x1x1_Coop") {
306
+ Kernel_128x16_1x1x1_Coop::grouped_mm(
307
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
308
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
309
+ c_strides, group_scale_strides);
310
+ } else if (schedule == "Kernel_128x16_2x1x1_Coop") {
311
+ Kernel_128x16_2x1x1_Coop::grouped_mm(
312
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
313
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
314
+ c_strides, group_scale_strides);
315
+ } else if (schedule == "Kernel_256x16_1x1x1_Coop") {
316
+ Kernel_256x16_1x1x1_Coop::grouped_mm(
317
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
318
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
319
+ c_strides, group_scale_strides);
320
+ } else if (schedule == "Kernel_256x16_2x1x1_Coop") {
321
+ Kernel_256x16_2x1x1_Coop::grouped_mm(
322
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
323
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
324
+ c_strides, group_scale_strides);
325
+ } else if (schedule == "Kernel_256x32_1x1x1_Coop") {
326
+ Kernel_256x32_1x1x1_Coop::grouped_mm(
327
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
328
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
329
+ c_strides, group_scale_strides);
330
+ } else if (schedule == "Kernel_256x32_2x1x1_Coop") {
331
+ Kernel_256x32_2x1x1_Coop::grouped_mm(
332
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
333
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
334
+ c_strides, group_scale_strides);
335
+ } else if (schedule == "Kernel_256x64_1x1x1_Coop") {
336
+ Kernel_256x64_1x1x1_Coop::grouped_mm(
337
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
338
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
339
+ c_strides, group_scale_strides);
340
+ } else if (schedule == "Kernel_256x64_2x1x1_Coop") {
341
+ Kernel_256x64_2x1x1_Coop::grouped_mm(
342
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
343
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
344
+ c_strides, group_scale_strides);
345
+ } else if (schedule == "Kernel_256x128_1x1x1_Coop") {
346
+ Kernel_256x128_1x1x1_Coop::grouped_mm(
347
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
348
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
349
+ c_strides, group_scale_strides);
350
+ } else if (schedule == "Kernel_256x128_2x1x1_Coop") {
351
+ Kernel_256x128_2x1x1_Coop::grouped_mm(
352
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
353
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
354
+ c_strides, group_scale_strides);
355
+ } else if (schedule == "Kernel_128x256_2x1x1_Coop") {
356
+ Kernel_128x256_2x1x1_Coop::grouped_mm(
357
+ out_tensors, a_tensors, b_tensors, a_scales, b_scales, b_group_scales,
358
+ b_group_size, expert_offsets, problem_sizes, a_strides, b_strides,
359
+ c_strides, group_scale_strides);
360
+ } else {
361
+ TORCH_CHECK(false,
362
+ "cutlass_w4a8_moe_mm: unknown schedule string: ", schedule);
363
+ }
364
+}
365
+
366
+void mm(torch::Tensor& out_tensors, const torch::Tensor& a_tensors,
367
+ const torch::Tensor& b_tensors, const torch::Tensor& a_scales,
368
+ const torch::Tensor& b_scales, const torch::Tensor& b_group_scales,
369
+ const int64_t b_group_size, const torch::Tensor& expert_offsets,
370
+ const torch::Tensor& problem_sizes, const torch::Tensor& a_strides,
371
+ const torch::Tensor& b_strides, const torch::Tensor& c_strides,
372
+ const torch::Tensor& group_scale_strides,
373
+ std::optional<std::string> maybe_schedule) {
374
+ // user has specified a schedule
375
+ if (maybe_schedule) {
376
+ mm_dispatch(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
377
+ b_group_scales, b_group_size, expert_offsets, problem_sizes,
378
+ a_strides, b_strides, c_strides, group_scale_strides,
379
+ *maybe_schedule);
380
+ return;
381
+ }
382
+
383
+ // use heuristic
384
+ int m_full = a_tensors.size(0);
385
+ int n = b_tensors.size(1);
386
+ int k = b_tensors.size(2) * PackFactor; // logical k
387
+ int num_experts = b_tensors.size(0);
388
+ // per-expert batch size assuming uniform distribution
389
+ int m_expert = m_full / num_experts;
390
+
391
+ std::string schedule;
392
+ if (m_expert <= 16) {
393
+ schedule = "Kernel_128x16_2x1x1_Coop";
394
+ } else if (m_expert <= 32) {
395
+ schedule = "Kernel_256x32_1x1x1_Coop";
396
+ } else if (m_expert <= 64) {
397
+ schedule = "Kernel_256x64_1x1x1_Coop";
398
+ } else if (m_expert <= 128) {
399
+ schedule = "Kernel_256x128_2x1x1_Coop";
400
+ } else { // m_expert > 128
401
+ schedule = "Kernel_128x256_2x1x1_Coop";
402
+ }
403
+
404
+ mm_dispatch(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
405
+ b_group_scales, b_group_size, expert_offsets, problem_sizes,
406
+ a_strides, b_strides, c_strides, group_scale_strides, schedule);
407
+}
408
+
409
+std::tuple<torch::Tensor, torch::Tensor> encode_and_reorder_int4b(
410
+ torch::Tensor const& b_tensors) {
411
+ TORCH_CHECK(b_tensors.dtype() == torch::kInt32);
412
+ TORCH_CHECK(b_tensors.dim() == 3); // (experts, n, k)
413
+ TORCH_CHECK(b_tensors.is_contiguous());
414
+ TORCH_CHECK(b_tensors.is_cuda());
415
+
416
+ int n = static_cast<int>(b_tensors.size(1));
417
+ int k = static_cast<int>(b_tensors.size(2)) * PackFactor; // logical k
418
+
419
+ // CUTLASS reorder_tensor requires k % 256 == 0 and n % 16 == 0.
420
+ // These misalignments cause silent OOB unless run under Compute Sanitizer.
421
+ TORCH_CHECK(k % 256 == 0, "logical k must be divisible by 256");
422
+ TORCH_CHECK(n % 16 == 0, "n must be divisible by 16");
423
+
424
+ // we will store the layout to an int32 tensor;
425
+ // this is the number of elements we need per layout
426
+ constexpr size_t layout_width = sizeof(LayoutB_Reordered) / sizeof(int32_t);
427
+
428
+ torch::Tensor b_tensors_packed = torch::empty_like(b_tensors);
429
+ int num_experts = static_cast<int>(b_tensors.size(0));
430
+
431
+ auto b_ptr = static_cast<QuantType const*>(b_tensors.const_data_ptr());
432
+ auto b_packed_ptr = static_cast<QuantType*>(b_tensors_packed.data_ptr());
433
+
434
+ // multiply by ull so result does not overflow int32
435
+ size_t num_int4_elems = 1ull * num_experts * n * k;
436
+ bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(b_ptr, b_packed_ptr,
437
+ num_int4_elems);
438
+ TORCH_CHECK(ok, "unified_encode_int4b failed");
439
+
440
+ // construct the layout once; assumes each expert has the same layout
441
+ using LayoutType = LayoutB_Reordered;
442
+ std::vector<LayoutType> layout_B_reordered_host(num_experts);
443
+ auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, Int<1>{}});
444
+ auto shape_B = cute::make_shape(n, k, Int<1>{});
445
+ auto layout_B = make_layout(shape_B, stride_B);
446
+ LayoutType layout_B_reordered = tile_to_shape(LayoutAtomQuant{}, shape_B);
447
+
448
+ // reorder weights for each expert
449
+ for (int i = 0; i < num_experts; i++) {
450
+ // since the storage type of int4b is 1 byte but one element is 4 bits
451
+ // we need to adjust the offset
452
+ int64_t offset =
453
+ 1ull * i * n * k * cutlass::sizeof_bits<QuantType>::value / 8;
454
+ cutlass::reorder_tensor(b_packed_ptr + offset, layout_B,
455
+ layout_B_reordered);
456
+ }
457
+
458
+ // save the packed layout to torch tensor so we can re-use it
459
+ auto cpu_opts =
460
+ torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
461
+ torch::Tensor layout_cpu =
462
+ torch::empty({num_experts, layout_width}, cpu_opts);
463
+
464
+ int32_t* layout_data = layout_cpu.data_ptr<int32_t>();
465
+ for (int i = 0; i < num_experts; ++i) {
466
+ std::memcpy(layout_data + i * layout_width, // dst (int32*)
467
+ &layout_B_reordered, // src (LayoutType*)
468
+ sizeof(LayoutType)); // number of bytes
469
+ }
470
+
471
+ torch::Tensor packed_layout =
472
+ layout_cpu.to(b_tensors.device(), /*non_blocking=*/false);
473
+
474
+ return {b_tensors_packed, packed_layout};
475
+}
476
+
477
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
478
+ m.impl("cutlass_w4a8_moe_mm", &mm);
479
+ m.impl("cutlass_encode_and_reorder_int4b_grouped", &encode_and_reorder_int4b);
480
+}
481
+
482
+} // namespace vllm::cutlass_w4a8_moe
483
+/////////////////////////////////////////////////////////////////////////////////////////////////
\ No newline at end of file
csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu
+3
-67
@@ -7,6 +7,7 @@
7
#include <c10/cuda/CUDAGuard.h>
8
#include <torch/all.h>
9
#include "cutlass_extensions/torch_utils.hpp"
10
+#include "w4a8_utils.cuh"
11
12
#include "core/registration.h"
13
@@ -395,71 +396,6 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
396
return packed_scales;
397
}
398
398
-/*
399
- GPU-accelerated implementation of cutlass::unified_encode_int4b.
400
- Constructs a lookup table in constant memory to map 8 bits
401
- (two 4-bit values) at a time. Assumes memory is contiguous
402
- and pointers are 16-byte aligned.
403
-*/
404
-__constant__ uint8_t kNibbleLUT[256];
405
-
406
-__global__ void unified_encode_int4b_device(const uint8_t* in, uint8_t* out,
407
- size_t nbytes) {
408
- constexpr size_t V = sizeof(uint4); // 16 bytes
409
- const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
410
- const size_t nthreads = size_t(gridDim.x) * blockDim.x;
411
- const size_t nvec = nbytes / V;
412
-
413
- // 1-D grid-stride loop over 16-byte chunks
414
- for (size_t vec = tid; vec < nvec; vec += nthreads) {
415
- uint4 v = reinterpret_cast<const uint4*>(in)[vec];
416
- uint8_t* b = reinterpret_cast<uint8_t*>(&v);
417
-#pragma unroll
418
- for (int i = 0; i < int(V); ++i) b[i] = kNibbleLUT[b[i]];
419
- reinterpret_cast<uint4*>(out)[vec] = v;
420
- }
421
-}
422
-
423
-static bool upload_lut() {
424
- std::array<uint8_t, 256> lut{};
425
- auto map_nib = [](uint8_t v) -> uint8_t {
426
- // 1..7 -> (8 - v); keep 0 and 8..15
427
- return (v == 0 || (v & 0x8)) ? v : uint8_t(8 - v);
428
- };
429
- for (int b = 0; b < 256; ++b) {
430
- uint8_t lo = b & 0xF;
431
- uint8_t hi = (b >> 4) & 0xF;
432
- lut[b] = uint8_t((map_nib(hi) << 4) | map_nib(lo));
433
- }
434
- cudaError_t e = cudaMemcpyToSymbol(kNibbleLUT, lut.data(), lut.size(),
435
- /*offset=*/0, cudaMemcpyHostToDevice);
436
-
437
- return (e == cudaSuccess);
438
-}
439
-
440
-static bool unified_encode_int4b(cutlass::int4b_t const* in,
441
- cutlass::int4b_t* out, size_t num_int4_elems) {
442
- // Build/upload LUT
443
- if (!upload_lut()) return false;
444
-
445
- static_assert(sizeof(typename cutlass::int4b_t::Storage) == 1,
446
- "int4 storage must be 1 byte");
447
- const size_t nbytes = num_int4_elems >> 1;
448
-
449
- auto* in_bytes = reinterpret_cast<uint8_t const*>(in);
450
- auto* out_bytes = reinterpret_cast<uint8_t*>(out);
451
-
452
- // kernel launch params
453
- constexpr int block = 256;
454
- const size_t nvec = nbytes / sizeof(uint4); // # of 16B vectors
455
- int grid = int((nvec + block - 1) / block);
456
- if (grid == 0) grid = 1; // ensure we still cover the tail in the kernel
457
-
458
- unified_encode_int4b_device<<<grid, block>>>(in_bytes, out_bytes, nbytes);
459
- cudaError_t err = cudaGetLastError();
460
- return (err == cudaSuccess);
461
-}
462
-
399
torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
400
TORCH_CHECK(B.dtype() == torch::kInt32);
401
TORCH_CHECK(B.dim() == 2);
@@ -477,8 +413,8 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
413
LayoutB_Reordered layout_B_reordered =
414
cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
415
480
- bool ok =
481
- vllm::cutlass_w4a8::unified_encode_int4b(B_ptr, B_packed_ptr, n * k);
416
+ bool ok = vllm::cutlass_w4a8_utils::unified_encode_int4b(B_ptr, B_packed_ptr,
417
+ n * k);
418
TORCH_CHECK(ok, "unified_encode_int4b failed");
419
cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
420
csrc/quantization/cutlass_w4a8/w4a8_utils.cu
new
+90
@@ -0,0 +1,90 @@
1
+#include "w4a8_utils.cuh"
2
+
3
+#include <array>
4
+#include <cuda_runtime.h>
5
+#include <cstdio>
6
+
7
+namespace vllm::cutlass_w4a8_utils {
8
+
9
+/*
10
+ GPU-accelerated implementation of cutlass::unified_encode_int4b.
11
+ Constructs a lookup table in constant memory to map 8 bits
12
+ (two 4-bit values) at a time. Assumes memory is contiguous
13
+ and pointers are 16-byte aligned.
14
+*/
15
+__constant__ uint8_t kNibbleLUT[256];
16
+
17
+__global__ void unified_encode_int4b_device(const uint8_t* in, uint8_t* out,
18
+ size_t nbytes) {
19
+ constexpr size_t V = sizeof(uint4); // 16 bytes
20
+ const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
21
+ const size_t nthreads = size_t(gridDim.x) * blockDim.x;
22
+ const size_t nvec = nbytes / V;
23
+
24
+ // 1-D grid-stride loop over 16-byte chunks
25
+ for (size_t vec = tid; vec < nvec; vec += nthreads) {
26
+ uint4 v = reinterpret_cast<const uint4*>(in)[vec];
27
+ uint8_t* b = reinterpret_cast<uint8_t*>(&v);
28
+#pragma unroll
29
+ for (int i = 0; i < int(V); ++i) b[i] = kNibbleLUT[b[i]];
30
+ reinterpret_cast<uint4*>(out)[vec] = v;
31
+ }
32
+}
33
+
34
+static bool upload_lut() {
35
+ std::array<uint8_t, 256> lut{};
36
+ auto map_nib = [](uint8_t v) -> uint8_t {
37
+ // 1..7 -> (8 - v); keep 0 and 8..15
38
+ return (v == 0 || (v & 0x8)) ? v : uint8_t(8 - v);
39
+ };
40
+ for (int b = 0; b < 256; ++b) {
41
+ uint8_t lo = b & 0xF;
42
+ uint8_t hi = (b >> 4) & 0xF;
43
+ lut[b] = uint8_t((map_nib(hi) << 4) | map_nib(lo));
44
+ }
45
+ cudaError_t e = cudaMemcpyToSymbol(kNibbleLUT, lut.data(), lut.size(),
46
+ /*offset=*/0, cudaMemcpyHostToDevice);
47
+
48
+ return (e == cudaSuccess);
49
+}
50
+
51
+bool unified_encode_int4b(cutlass::int4b_t const* in, cutlass::int4b_t* out,
52
+ size_t num_int4_elems) {
53
+ // Build/upload LUT
54
+ if (!upload_lut()) return false;
55
+
56
+ static_assert(sizeof(typename cutlass::int4b_t::Storage) == 1,
57
+ "int4 storage must be 1 byte");
58
+ const size_t nbytes = num_int4_elems >> 1;
59
+
60
+ auto* in_bytes = reinterpret_cast<uint8_t const*>(in);
61
+ auto* out_bytes = reinterpret_cast<uint8_t*>(out);
62
+
63
+ // kernel launch params
64
+ constexpr int block = 256;
65
+ const size_t nvec = nbytes / sizeof(uint4); // # of 16B vectors
66
+ int grid = int((nvec + block - 1) / block);
67
+ if (grid == 0) grid = 1; // ensure we still cover the tail in the kernel
68
+
69
+ unified_encode_int4b_device<<<grid, block>>>(in_bytes, out_bytes, nbytes);
70
+
71
+ // launch errors
72
+ cudaError_t err = cudaGetLastError();
73
+ if (err != cudaSuccess) {
74
+ printf("unified_encode_int4b_device launch error: %s (%d)\n",
75
+ cudaGetErrorString(err), err);
76
+ return false;
77
+ }
78
+
79
+ // runtime errors
80
+ err = cudaDeviceSynchronize();
81
+ if (err != cudaSuccess) {
82
+ printf("unified_encode_int4b_device runtime error: %s (%d)\n",
83
+ cudaGetErrorString(err), err);
84
+ return false;
85
+ }
86
+
87
+ return true;
88
+}
89
+
90
+} // namespace vllm::cutlass_w4a8_utils
\ No newline at end of file
csrc/quantization/cutlass_w4a8/w4a8_utils.cuh
new
+11
@@ -0,0 +1,11 @@
1
+#pragma once
2
+
3
+#include <cstddef>
4
+#include "cutlass/numeric_types.h"
5
+
6
+namespace vllm::cutlass_w4a8_utils {
7
+
8
+bool unified_encode_int4b(cutlass::int4b_t const* in, cutlass::int4b_t* out,
9
+ size_t num_int4_elems);
10
+
11
+} // namespace vllm::cutlass_w4a8_utils
\ No newline at end of file
csrc/quantization/w8a8/cutlass/moe/moe_data.cu
+5
-3
@@ -136,15 +136,17 @@ inline void launch_compute_problem_sizes(const torch::Tensor& topk_ids,
136
void get_cutlass_moe_mm_problem_sizes_caller(
137
const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
138
torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
139
- const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {
139
+ const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
140
+ std::optional<bool> force_swap_ab = std::nullopt) {
141
auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
142
auto options_int32 =
143
torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
144
torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);
145
146
// Swap-AB should be disabled for FP4 path
146
- bool may_swap_ab = (!blockscale_offsets.has_value()) &&
147
- (topk_ids.numel() <= SWAP_AB_THRESHOLD);
147
+ bool may_swap_ab =
148
+ force_swap_ab.value_or((!blockscale_offsets.has_value()) &&
149
+ (topk_ids.numel() <= SWAP_AB_THRESHOLD));
150
151
launch_compute_problem_sizes(topk_ids, problem_sizes1, problem_sizes2,
152
atomic_buffer, num_experts, n, k, stream,
csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu
+5
-3
@@ -80,7 +80,8 @@ void get_cutlass_moe_mm_data_caller(
80
void get_cutlass_moe_mm_problem_sizes_caller(
81
const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
82
torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
83
- const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets);
83
+ const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
84
+ std::optional<bool> force_swap_ab = std::nullopt);
85
86
void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
87
torch::Tensor& problem_sizes1,
@@ -303,14 +304,15 @@ void get_cutlass_moe_mm_data(
304
void get_cutlass_moe_mm_problem_sizes(
305
const torch::Tensor& topk_ids, torch::Tensor& problem_sizes1,
306
torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
306
- const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {
307
+ const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets,
308
+ std::optional<bool> force_swap_ab = std::nullopt) {
309
int32_t version_num = get_sm_version_num();
310
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
311
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
312
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
313
get_cutlass_moe_mm_problem_sizes_caller(topk_ids, problem_sizes1,
314
problem_sizes2, num_experts, n, k,
313
- blockscale_offsets);
315
+ blockscale_offsets, force_swap_ab);
316
return;
317
#endif
318
TORCH_CHECK_NOT_IMPLEMENTED(
csrc/torch_bindings.cpp
+25
-1
@@ -350,6 +350,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
350
ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
351
// conditionally compiled so impl registration is in source file
352
353
+ // CUTLASS w4a8 grouped GEMM
354
+ ops.def(
355
+ "cutlass_w4a8_moe_mm("
356
+ " Tensor! out_tensors,"
357
+ " Tensor a_tensors,"
358
+ " Tensor b_tensors,"
359
+ " Tensor a_scales,"
360
+ " Tensor b_scales,"
361
+ " Tensor b_group_scales,"
362
+ " int b_group_size,"
363
+ " Tensor expert_offsets,"
364
+ " Tensor problem_sizes,"
365
+ " Tensor a_strides,"
366
+ " Tensor b_strides,"
367
+ " Tensor c_strides,"
368
+ " Tensor group_scale_strides,"
369
+ " str? maybe_schedule"
370
+ ") -> ()");
371
+ ops.def(
372
+ "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
373
+ "Tensor)");
374
+ // conditionally compiled so impl registration is in source file
375
+
376
#endif
377
378
// Dequantization for GGML.
@@ -466,7 +489,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
489
" Tensor! problem_sizes1, "
490
" Tensor! problem_sizes2, "
491
" int num_experts, int n, int k, "
469
- " Tensor? blockscale_offsets) -> ()");
492
+ " Tensor? blockscale_offsets, "
493
+ " bool? force_swap_ab) -> ()");
494
ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
495
&get_cutlass_moe_mm_problem_sizes);
496
tests/kernels/quantization/test_cutlass_w4a8.py
+41
-5
@@ -12,8 +12,11 @@ import torch
12
13
from vllm import _custom_ops as ops
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
+ convert_packed_uint4b8_to_signed_int4_inplace,
16
+ pack_cols,
17
pack_rows,
18
quantize_weights,
19
+ unpack_quantized_values_into_int32,
20
)
21
from vllm.platforms import current_platform
22
from vllm.scalar_type import ScalarType, scalar_types
@@ -167,8 +170,7 @@ def create_test_tensors(
170
171
# for the practical use case we need per-tok scales for fp8 activations
172
w_tok_s = torch.randn((m,), device="cuda", dtype=types.token_scale_type)
170
- # weights are already per-group quantized, use placeholder here
171
- w_ch_s = torch.ones((n,), device="cuda", dtype=types.channel_scale_type)
173
+ w_ch_s = torch.randn((n,), device="cuda", dtype=types.channel_scale_type)
174
175
return Tensors(
176
w_ref=w_ref,
@@ -211,7 +213,7 @@ def mm_test_helper(
213
print(output_ref)
214
215
torch.testing.assert_close(
214
- output, output_ref.to(output.dtype), rtol=1e-3, atol=1e-3
216
+ output, output_ref.to(output.dtype), rtol=1e-2, atol=1e-2
217
)
218
219
@@ -257,7 +259,7 @@ def test_w4a8_cuda_graph():
259
)
260
261
w_tok_s = torch.randn((m,), device="cuda", dtype=torch.float32)
260
- w_ch_s = torch.ones((n,), device="cuda", dtype=torch.float32)
262
+ w_ch_s = torch.randn((n,), device="cuda", dtype=torch.float32)
263
264
# Construct a trivial model with a single layer that calls the kernel
265
model = W4A8Layer(
@@ -287,4 +289,38 @@ def test_w4a8_cuda_graph():
289
output.zero_()
290
g.replay()
291
290
- torch.testing.assert_close(output, output_ref, rtol=1e-3, atol=1e-3)
292
+ torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=1e-2)
293
+
294
+
295
+@pytest.mark.skipif(
296
+ not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type."
297
+)
298
+@pytest.mark.parametrize("shape", MNK_SHAPES)
299
+def test_convert_packed_uint4b8_to_signed_int4_inplace(shape):
300
+ """
301
+ The W4A16 checkpoints encode the weights as int4b8 packed to int32.
302
+ The CUTLASS kernels expect signed int4 packed to int32.
303
+ This tests checks that the runtime int4b8 -> signed int4 conversion
304
+ matches the offline conversion step exactly.
305
+ """
306
+ _, N, K = shape
307
+ # random weights packed to int32
308
+ t = torch.randint(
309
+ low=torch.iinfo(torch.int32).min,
310
+ high=torch.iinfo(torch.int32).max + 1,
311
+ size=(N, K // 8),
312
+ dtype=torch.int32,
313
+ device="cuda",
314
+ )
315
+
316
+ # compute reference
317
+ unpacked = unpack_quantized_values_into_int32(
318
+ t.clone(), scalar_types.uint4b8, packed_dim=1
319
+ )
320
+ unpacked = unpacked - 8 # int4b8 -> signed int4
321
+ ref = pack_cols(unpacked & 0x0F, 4, *unpacked.shape)
322
+
323
+ out = convert_packed_uint4b8_to_signed_int4_inplace(t.clone())
324
+
325
+ assert torch.equal(ref, out)
326
+ assert not torch.equal(ref, t)
tests/kernels/quantization/test_cutlass_w4a8_moe.py
new
+340
@@ -0,0 +1,340 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer.
5
+"""
6
+
7
+import random
8
+from dataclasses import dataclass
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm import _custom_ops as ops
14
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
+ pack_rows,
16
+ quantize_weights,
17
+)
18
+from vllm.platforms import current_platform
19
+from vllm.scalar_type import ScalarType, scalar_types
20
+
21
+IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9
22
+
23
+
24
+def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
25
+ finfo = torch.finfo(torch.float8_e4m3fn)
26
+ return tensor.clamp(min=finfo.min, max=finfo.max).to(dtype=torch.float8_e4m3fn)
27
+
28
+
29
+def cutlass_quantize(
30
+ atype: torch.dtype,
31
+ w: torch.Tensor,
32
+ wtype: ScalarType,
33
+ stype: torch.dtype | None,
34
+ group_size: int | None,
35
+ zero_points: bool = False,
36
+):
37
+ """
38
+ Quantize weights into W4 and compute reference dequantized weights.
39
+
40
+ Encoding/reordering of weights and packing of scales is deferred
41
+ until after all experts are combined.
42
+ """
43
+ assert wtype.is_integer(), "TODO: support floating point weights"
44
+
45
+ w_ref, w_q, w_s, w_zp = quantize_weights(
46
+ w, wtype, group_size=group_size, zero_points=zero_points
47
+ )
48
+
49
+ # Since scales are later cast to fp8, recompute w_ref in atype here.
50
+ w_ref = (
51
+ w_q.to(torch.float32)
52
+ * w_s.to(atype).to(torch.float32).repeat_interleave(group_size, dim=0)
53
+ ).to(atype)
54
+
55
+ # Bit mask prevents sign extension of int4 when packing.
56
+ w_q = pack_rows(w_q & 0x0F, wtype.size_bits, *w_q.shape)
57
+ # Make weights row-major (N, K).
58
+ w_q = w_q.t().contiguous()
59
+
60
+ return w_ref, w_q, w_s.to(atype), w_zp
61
+
62
+
63
+def cutlass_preprocess(
64
+ w_q_experts: list[torch.Tensor], w_s_experts: list[torch.Tensor]
65
+):
66
+ """
67
+ Reorder/encode expert weights and pack scales.
68
+
69
+ Returns:
70
+ w_q_packed: Packed/encoded int4 weights for all experts.
71
+ w_s_packed: Packed fp8 scales for all experts.
72
+ packed_layout: Layout/stride metadata for grouped GEMM.
73
+ """
74
+ w_s_packed = ops.cutlass_pack_scale_fp8(torch.stack(w_s_experts))
75
+ w_q_packed, packed_layout = ops.cutlass_encode_and_reorder_int4b_grouped(
76
+ torch.stack(w_q_experts)
77
+ ) # expects dim 3
78
+ return w_q_packed, w_s_packed, packed_layout
79
+
80
+
81
+GROUP_SIZE = 128
82
+# (num_experts, N, K)
83
+TEST_SHAPES = [
84
+ (8, 512, 2048),
85
+ (8, 2048, 2048),
86
+ (64, 512, 1024),
87
+ (64, 2048, 2048),
88
+ (4, 2048, 768),
89
+ (8, 768, 2048),
90
+ (64, 1536, 2048),
91
+ (128, 8192, 4096), # test overflow int32
92
+]
93
+ALIGNMENT = 16 # torch._scaled_mm alignment for M, needed for reference check
94
+
95
+
96
+@dataclass
97
+class MoETestSetup:
98
+ num_experts: int
99
+ K: int
100
+ N: int
101
+ Ms: list[int]
102
+ M_full: int
103
+ a: torch.Tensor
104
+ a_ref: torch.Tensor
105
+ a_strides: torch.Tensor
106
+ out: torch.Tensor
107
+ c_strides: torch.Tensor
108
+ per_tok_scales: torch.Tensor
109
+ per_chan_scales: torch.Tensor
110
+ w_refs: list[torch.Tensor]
111
+ w_q_packed: torch.Tensor
112
+ w_s_packed: torch.Tensor
113
+ problem_sizes: torch.Tensor
114
+ expert_offsets: torch.Tensor
115
+ b_strides: torch.Tensor
116
+ group_scale_strides: torch.Tensor
117
+
118
+
119
+def make_moe_test_setup(
120
+ num_experts: int,
121
+ K: int,
122
+ N: int,
123
+ *,
124
+ alignment: int = ALIGNMENT,
125
+ max_blocks: int = 64,
126
+ device: str = "cuda",
127
+ random_zero: bool = False,
128
+) -> MoETestSetup:
129
+ """Create a full set of tensors for testing cutlass_w4a8_moe_mm."""
130
+
131
+ assert K % GROUP_SIZE == 0
132
+ # Token counts per expert (multiples of `alignment`).
133
+ Ms = [alignment * random.randint(1, max_blocks) for _ in range(num_experts)]
134
+
135
+ # set random experts to 0 tokens
136
+ if random_zero and num_experts > 1:
137
+ num_zero = max(1, num_experts // 8)
138
+ zero_indices = random.sample(range(num_experts), k=num_zero)
139
+ for idx in zero_indices:
140
+ Ms[idx] = 0
141
+
142
+ M_full = sum(Ms)
143
+ assert M_full > 0
144
+
145
+ # Activations.
146
+ a = to_fp8(torch.randn((M_full, K), device=device))
147
+ a_ref = a.to(torch.float32)
148
+ a_strides = torch.full((num_experts,), K, dtype=torch.int64, device=device)
149
+
150
+ # Output buffer.
151
+ out = torch.empty((M_full, N), dtype=torch.bfloat16, device=device)
152
+ c_strides = torch.full((num_experts,), N, dtype=torch.int64, device=device)
153
+
154
+ # Channel/token scales.
155
+ per_tok_scales = torch.randn((M_full, 1), dtype=torch.float32, device=device)
156
+ per_chan_scales = torch.randn(
157
+ (num_experts, N, 1), dtype=torch.float32, device=device
158
+ )
159
+
160
+ # Expert weights and scales.
161
+ wtype = scalar_types.int4
162
+ atype = stype = torch.float8_e4m3fn
163
+ w_refs, w_qs, w_ss = [], [], []
164
+ for _ in range(num_experts):
165
+ b = to_fp8(torch.randn((K, N), device=device))
166
+ w_ref, w_q, w_s, _ = cutlass_quantize(
167
+ atype, b.to(torch.float16), wtype, stype, GROUP_SIZE, zero_points=False
168
+ )
169
+ w_refs.append(w_ref)
170
+ w_qs.append(w_q)
171
+ w_ss.append(w_s)
172
+
173
+ w_q_packed, w_s_packed, packed_layout = cutlass_preprocess(w_qs, w_ss)
174
+
175
+ problem_sizes = torch.tensor(
176
+ [[N, M, K] for M in Ms], dtype=torch.int32, device=device
177
+ )
178
+
179
+ expert_offsets = torch.cat(
180
+ [
181
+ torch.tensor([0], dtype=torch.int64),
182
+ torch.cumsum(torch.tensor(Ms, dtype=torch.int64), dim=0)[:-1],
183
+ ]
184
+ ).to(device=device)
185
+
186
+ # B strides and group scale strides.
187
+ b_strides = packed_layout
188
+ group_scale_strides = torch.zeros(
189
+ (num_experts, 2), dtype=torch.int64, device=device
190
+ )
191
+ group_scale_strides[:, 0] = N
192
+
193
+ return MoETestSetup(
194
+ num_experts=num_experts,
195
+ K=K,
196
+ N=N,
197
+ Ms=Ms,
198
+ M_full=M_full,
199
+ a=a,
200
+ a_ref=a_ref,
201
+ a_strides=a_strides,
202
+ out=out,
203
+ c_strides=c_strides,
204
+ per_tok_scales=per_tok_scales,
205
+ per_chan_scales=per_chan_scales,
206
+ w_refs=w_refs,
207
+ w_q_packed=w_q_packed,
208
+ w_s_packed=w_s_packed,
209
+ problem_sizes=problem_sizes,
210
+ expert_offsets=expert_offsets,
211
+ b_strides=b_strides,
212
+ group_scale_strides=group_scale_strides,
213
+ )
214
+
215
+
216
+def compute_moe_reference_output(setup: MoETestSetup) -> torch.Tensor:
217
+ """Compute reference output using torch._scaled_mm per expert."""
218
+ out_ref = torch.empty_like(setup.out)
219
+
220
+ ends = torch.cumsum(torch.tensor(setup.Ms), 0).tolist()
221
+ starts = setup.expert_offsets.cpu().tolist()
222
+
223
+ for i in range(setup.num_experts):
224
+ start, end = starts[i], ends[i]
225
+ if start == end:
226
+ continue
227
+
228
+ out_ref_i = torch._scaled_mm(
229
+ setup.a_ref[start:end].to(torch.float8_e4m3fn),
230
+ setup.w_refs[i].to(torch.float8_e4m3fn).t().contiguous().t(),
231
+ setup.per_tok_scales[start:end], # (M, 1)
232
+ setup.per_chan_scales[i].reshape(1, -1), # (1, N)
233
+ out_dtype=torch.bfloat16,
234
+ use_fast_accum=True,
235
+ )
236
+ out_ref[start:end] = out_ref_i
237
+
238
+ return out_ref
239
+
240
+
241
+@pytest.mark.skipif(
242
+ not IS_SUPPORTED_BY_GPU,
243
+ reason="W4A8 Grouped GEMM is not supported on this GPU type.",
244
+)
245
+@pytest.mark.parametrize("shape", TEST_SHAPES)
246
+@pytest.mark.parametrize("random_zero", [True, False])
247
+def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero):
248
+ num_experts, N, K = shape
249
+ current_platform.seed_everything(42)
250
+ setup = make_moe_test_setup(
251
+ num_experts=num_experts, K=K, N=N, max_blocks=64, random_zero=random_zero
252
+ )
253
+
254
+ ops.cutlass_w4a8_moe_mm(
255
+ setup.out,
256
+ setup.a,
257
+ setup.w_q_packed,
258
+ setup.per_tok_scales,
259
+ setup.per_chan_scales,
260
+ setup.w_s_packed,
261
+ GROUP_SIZE,
262
+ setup.expert_offsets,
263
+ setup.problem_sizes,
264
+ setup.a_strides,
265
+ setup.b_strides,
266
+ setup.c_strides,
267
+ setup.group_scale_strides,
268
+ )
269
+ torch.cuda.synchronize()
270
+
271
+ out_ref = compute_moe_reference_output(setup)
272
+ torch.testing.assert_close(setup.out, out_ref, rtol=1e-2, atol=1e-2)
273
+
274
+
275
+class W4A8MoELayer(torch.nn.Module):
276
+ """
277
+ Minimal wrapper module to test cuda graphs
278
+ """
279
+
280
+ def __init__(self, setup: MoETestSetup):
281
+ super().__init__()
282
+ self.setup = setup
283
+
284
+ def forward(self, a: torch.Tensor) -> torch.Tensor:
285
+ s = self.setup
286
+ ops.cutlass_w4a8_moe_mm(
287
+ s.out,
288
+ a,
289
+ s.w_q_packed,
290
+ s.per_tok_scales,
291
+ s.per_chan_scales,
292
+ s.w_s_packed,
293
+ GROUP_SIZE,
294
+ s.expert_offsets,
295
+ s.problem_sizes,
296
+ s.a_strides,
297
+ s.b_strides,
298
+ s.c_strides,
299
+ s.group_scale_strides,
300
+ )
301
+ return s.out
302
+
303
+
304
+@pytest.mark.skipif(
305
+ not IS_SUPPORTED_BY_GPU,
306
+ reason="W4A8 Grouped GEMM is not supported on this GPU type.",
307
+)
308
+def test_cutlass_w4a8_moe_mm_cuda_graph():
309
+ current_platform.seed_everything(42)
310
+ # Fixed config for CUDA graph test (single parameter point).
311
+ num_experts = 8
312
+ K = 512
313
+ N = 2048
314
+
315
+ setup = make_moe_test_setup(
316
+ num_experts=num_experts,
317
+ K=K,
318
+ N=N,
319
+ max_blocks=32,
320
+ )
321
+
322
+ # Construct model that calls the grouped GEMM kernel.
323
+ model = W4A8MoELayer(setup)
324
+
325
+ # Build reference output once.
326
+ out_ref = compute_moe_reference_output(setup)
327
+
328
+ # Capture and run the model in a CUDA graph.
329
+ a_static = setup.a.clone() # static input tensor for graph replay
330
+
331
+ stream = torch.cuda.Stream()
332
+ with torch.cuda.stream(stream):
333
+ g = torch.cuda.CUDAGraph()
334
+ with torch.cuda.graph(g):
335
+ out_static = model(a_static)
336
+
337
+ out_static.zero_()
338
+ g.replay()
339
+
340
+ torch.testing.assert_close(out_static, out_ref, rtol=1e-2, atol=1e-2)
vllm/_custom_ops.py
+89
-1
@@ -695,6 +695,10 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
695
def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
696
return torch.empty_like(b, memory_format=torch.contiguous_format)
697
698
+ @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped")
699
+ def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor:
700
+ return torch.empty_like(b, memory_format=torch.contiguous_format)
701
+
702
703
if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
704
@@ -1058,6 +1062,7 @@ def get_cutlass_moe_mm_problem_sizes(
1062
n: int,
1063
k: int,
1064
blockscale_offsets: torch.Tensor | None = None,
1065
+ force_swap_ab: bool | None = None,
1066
):
1067
"""
1068
Compute only the per-expert problem sizes needed by the two grouped matrix
@@ -1067,9 +1072,20 @@ def get_cutlass_moe_mm_problem_sizes(
1072
- problem_sizes1, problem_sizes2: M×N×K sizes of each expert's
1073
multiplication for the two grouped MMs
1074
used in the fused MoE operation.
1075
+ Optional:
1076
+ - force_swap_ab: If set to True or False, explicitly enable or disable the
1077
+ A/B input swap optimization. If None (default), the swap
1078
+ is selected automatically based on tensor sizes.
1079
"""
1080
return torch.ops._C.get_cutlass_moe_mm_problem_sizes(
1072
- topk_ids, problem_sizes1, problem_sizes2, num_experts, n, k, blockscale_offsets
1081
+ topk_ids,
1082
+ problem_sizes1,
1083
+ problem_sizes2,
1084
+ num_experts,
1085
+ n,
1086
+ k,
1087
+ blockscale_offsets,
1088
+ force_swap_ab,
1089
)
1090
1091
@@ -1457,6 +1473,78 @@ def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor:
1473
return torch.ops._C.cutlass_encode_and_reorder_int4b(b)
1474
1475
1476
+def cutlass_w4a8_moe_mm(
1477
+ out_tensors: torch.Tensor,
1478
+ a_tensors: torch.Tensor,
1479
+ b_tensors: torch.Tensor,
1480
+ a_scales: torch.Tensor,
1481
+ b_scales: torch.Tensor,
1482
+ b_group_scales: torch.Tensor,
1483
+ b_group_size: int,
1484
+ expert_offsets: torch.Tensor,
1485
+ problem_sizes: torch.Tensor,
1486
+ a_strides: torch.Tensor,
1487
+ b_strides: torch.Tensor,
1488
+ c_strides: torch.Tensor,
1489
+ group_scale_strides: torch.Tensor,
1490
+ maybe_schedule: str | None = None,
1491
+):
1492
+ """
1493
+ Executes the CUTLASS-based fused-MoE grouped matrix multiplication for the
1494
+ W4A8 quantization scheme. Uses group-wise quantization (INT4 -> FP8)
1495
+ and both per-channel + per-token scaling in the epilogue.
1496
+
1497
+ Args:
1498
+ out_tensors:
1499
+ Output buffer for all experts (updated in-place).
1500
+ a_tensors:
1501
+ FP8 (E4M3FN) activations for all experts.
1502
+ b_tensors:
1503
+ INT4-packed weight matrix for all experts, packed to INT32
1504
+ a_scales:
1505
+ Per-token FP8 activation scales, applied in the epilogue.
1506
+ b_scales:
1507
+ Per-channel FP8 weight scales for each expert, applied in the epilogue.
1508
+ b_group_scales:
1509
+ FP8 scale values for group-wise INT4 weight blocks.
1510
+ b_group_size:
1511
+ Number of elements grouped under each entry of b_group_scales.
1512
+ expert_offsets:
1513
+ Cumulative token offsets
1514
+ problem_sizes:
1515
+ Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher.
1516
+ a/b/c/group_scale_strides:
1517
+ Strides describing the memory layout of the input tensors.
1518
+ maybe_schedule:
1519
+ Optional override to choose a specific kernel or epilogue schedule.
1520
+
1521
+ Returns:
1522
+ out_tensors updated in-place with the dequantized INT4xFP8 grouped GEMM result.
1523
+ """
1524
+ return torch.ops._C.cutlass_w4a8_moe_mm(
1525
+ out_tensors,
1526
+ a_tensors,
1527
+ b_tensors,
1528
+ a_scales,
1529
+ b_scales,
1530
+ b_group_scales,
1531
+ b_group_size,
1532
+ expert_offsets,
1533
+ problem_sizes,
1534
+ a_strides,
1535
+ b_strides,
1536
+ c_strides,
1537
+ group_scale_strides,
1538
+ maybe_schedule,
1539
+ )
1540
+
1541
+
1542
+def cutlass_encode_and_reorder_int4b_grouped(
1543
+ b_tensors: torch.Tensor,
1544
+) -> tuple[torch.Tensor, torch.Tensor]:
1545
+ return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors)
1546
+
1547
+
1548
if hasattr(torch.ops._C, "permute_cols"):
1549
1550
@register_fake("_C::permute_cols")
vllm/model_executor/layers/fused_moe/__init__.py
+4
@@ -63,8 +63,10 @@ if HAS_TRITON:
63
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
64
CutlassBatchedExpertsFp8,
65
CutlassExpertsFp8,
66
+ CutlassExpertsW4A8Fp8,
67
cutlass_moe_fp4,
68
cutlass_moe_fp8,
69
+ cutlass_moe_w4a8_fp8,
70
)
71
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
72
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
@@ -88,8 +90,10 @@ if HAS_TRITON:
90
"grouped_topk",
91
"cutlass_moe_fp8",
92
"cutlass_moe_fp4",
93
+ "cutlass_moe_w4a8_fp8",
94
"CutlassExpertsFp8",
95
"CutlassBatchedExpertsFp8",
96
+ "CutlassExpertsW4A8Fp8",
97
"TritonExperts",
98
"BatchedTritonExperts",
99
"DeepGemmExperts",
vllm/model_executor/layers/fused_moe/config.py
+29
@@ -143,6 +143,7 @@ class FusedMoEQuantDesc:
143
scale: Union[torch.Tensor, "PrecisionConfig", None] = None
144
145
# Quantization alphas or gscales, used for nvfp4 types.
146
+ # W4A8 FP8: used for per-channel scales
147
# TODO(bnell): put some of these in subclasses
148
alpha_or_gscale: torch.Tensor | None = None
149
@@ -442,7 +443,9 @@ class FusedMoEQuantConfig:
443
- a1_scale: Optional scale to be used for a1.
444
- a2_scale: Optional scale to be used for a2.
445
- g1_alphas: Optional global quantization scales for w1 (for nvfp4).
446
+ per-channel scales for w1 (for W4A8 FP8).
447
- g2_alphas: Optional global quantization scales for w2 (for nvfp4).
448
+ per-channel scales for w2 (for W4A8 FP8).
449
- a1_gscale: Optional global quantization scales for a1 (for nvfp4).
450
- a2_gscale: Optional global quantization scales for a2 (for nvfp4).
451
- w1_bias: Optional biases for w1 (GPT OSS Triton).
@@ -461,6 +464,7 @@ class FusedMoEQuantConfig:
464
"mxfp4",
465
"mxfp6_e3m2",
466
"mxfp6_e2m3",
467
+ "int4",
468
}
469
470
if weight_dtype is None:
@@ -671,6 +675,31 @@ def int8_w8a16_moe_quant_config(
675
)
676
677
678
+def int4_w4afp8_moe_quant_config(
679
+ w1_scale: torch.Tensor,
680
+ w2_scale: torch.Tensor,
681
+ g1_alphas: torch.Tensor,
682
+ g2_alphas: torch.Tensor,
683
+ per_act_token_quant: bool = False,
684
+ per_out_ch_quant: bool = False,
685
+ block_shape: list[int] | None = None,
686
+) -> FusedMoEQuantConfig:
687
+ """
688
+ Construct a quant config for fp8 activations and int4 weights.
689
+ """
690
+ return FusedMoEQuantConfig.make(
691
+ torch.float8_e4m3fn, # quant dtype for activations
692
+ w1_scale=w1_scale,
693
+ w2_scale=w2_scale,
694
+ g1_alphas=g1_alphas,
695
+ g2_alphas=g2_alphas,
696
+ per_act_token_quant=per_act_token_quant,
697
+ per_out_ch_quant=per_out_ch_quant,
698
+ block_shape=block_shape,
699
+ weight_dtype="int4", # weight dtype for weights
700
+ )
701
+
702
+
703
def biased_moe_quant_config(
704
w1_bias: torch.Tensor | None,
705
w2_bias: torch.Tensor | None,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+401
@@ -1052,3 +1052,404 @@ def run_cutlass_block_scaled_fused_experts(
1052
return (
1053
c2[c_map].view(m, topk, k) * topk_weights.view(m, topk, 1).to(out_dtype)
1054
).sum(dim=1)
1055
+
1056
+
1057
+# W4A8
1058
+def run_cutlass_moe_w4a8_fp8(
1059
+ output: torch.Tensor,
1060
+ hidden_states: torch.Tensor,
1061
+ w1: torch.Tensor,
1062
+ w2: torch.Tensor,
1063
+ topk_ids: torch.Tensor,
1064
+ activation_callable: Callable,
1065
+ global_num_experts: int,
1066
+ expert_map: torch.Tensor | None,
1067
+ w1_scale: torch.Tensor | None,
1068
+ w2_scale: torch.Tensor | None,
1069
+ a1q_scale: torch.Tensor | None,
1070
+ a2_scale: torch.Tensor | None,
1071
+ w1_chan_scale: torch.Tensor,
1072
+ w2_chan_scale: torch.Tensor,
1073
+ a_strides1: torch.Tensor,
1074
+ a_strides2: torch.Tensor,
1075
+ b_strides1: torch.Tensor,
1076
+ b_strides2: torch.Tensor,
1077
+ c_strides1: torch.Tensor,
1078
+ c_strides2: torch.Tensor,
1079
+ s_strides1: torch.Tensor,
1080
+ s_strides2: torch.Tensor,
1081
+ workspace13: torch.Tensor,
1082
+ workspace2: torch.Tensor,
1083
+ expert_num_tokens: torch.Tensor | None,
1084
+ out_dtype: torch.dtype,
1085
+ per_act_token: bool,
1086
+ per_out_ch: bool,
1087
+ use_batched_format: bool,
1088
+ topk_weights: torch.Tensor | None,
1089
+ group_size: int,
1090
+):
1091
+ a1q = hidden_states
1092
+ M = a1q.size(0)
1093
+ local_E = w1.size(0)
1094
+ device = a1q.device
1095
+ _, K, N_packed = w2.shape
1096
+ N = N_packed * 8 # logical N, pack 8 int4 into 1 int32
1097
+
1098
+ assert per_act_token, "W4A8 must use per-token scales"
1099
+ assert per_out_ch, "W4A8 must use per-channel scales"
1100
+ assert w1_scale is not None
1101
+ assert w2_scale is not None
1102
+ assert w1_scale.dtype == torch.float8_e4m3fn
1103
+ assert w2_scale.dtype == torch.float8_e4m3fn
1104
+ assert w1.dtype == torch.int32
1105
+ assert w2.dtype == torch.int32
1106
+ assert w1_chan_scale.dtype == torch.float32
1107
+ assert w2_chan_scale.dtype == torch.float32
1108
+ assert w1.size(0) == w2.size(0), "Weights expert number mismatch"
1109
+ assert a1q_scale is not None
1110
+ assert a2_scale is None
1111
+ assert out_dtype in [torch.bfloat16], f"Invalid output dtype: {out_dtype}"
1112
+ if expert_map is not None:
1113
+ assert expert_num_tokens is None
1114
+ assert not use_batched_format, "batched format not supported yet"
1115
+ assert group_size == 128, f"Only group size 128 supported but got {group_size=}"
1116
+
1117
+ assert global_num_experts != -1
1118
+ assert w1.size(2) * 8 == K, (
1119
+ f"w1 hidden size mismatch: got {w1.size(2) * 8}, expected {K=}"
1120
+ )
1121
+
1122
+ # Translate info from expert_map to topk_ids
1123
+ if expert_map is not None:
1124
+ local_topk_ids = torch.where(
1125
+ expert_map[topk_ids] != -1, expert_map[topk_ids], -1
1126
+ )
1127
+ else:
1128
+ local_topk_ids = topk_ids
1129
+
1130
+ topk = local_topk_ids.size(1)
1131
+ a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn), (M * topk, K))
1132
+ mm1_out = _resize_cache(workspace13, (M * topk, N * 2))
1133
+ act_out = _resize_cache(workspace2, (M * topk, N))
1134
+ # original workspace are based on input hidden_states dtype (bf16)
1135
+ quant_out = _resize_cache(
1136
+ workspace13.view(dtype=torch.float8_e4m3fn), (M * topk, N)
1137
+ )
1138
+ mm2_out = _resize_cache(workspace2, (M * topk, K))
1139
+
1140
+ problem_sizes1 = torch.empty(
1141
+ (global_num_experts, 3), dtype=torch.int32, device=device
1142
+ )
1143
+ problem_sizes2 = torch.empty(
1144
+ (global_num_experts, 3), dtype=torch.int32, device=device
1145
+ )
1146
+
1147
+ num_expert = global_num_experts if expert_map is None else expert_map.size(0)
1148
+ # permuted a1q reuses workspace2
1149
+ a1q, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
1150
+ a1q,
1151
+ a1q_scale,
1152
+ topk_ids,
1153
+ num_expert,
1154
+ local_E,
1155
+ expert_map,
1156
+ permuted_hidden_states=a1q_perm,
1157
+ )
1158
+ expert_offsets = expert_offsets[:-1]
1159
+
1160
+ # For RS gemm SwapAB is always enabled (swap logical M, N in the problem shape)
1161
+ ops.get_cutlass_moe_mm_problem_sizes(
1162
+ local_topk_ids,
1163
+ problem_sizes1,
1164
+ problem_sizes2,
1165
+ global_num_experts,
1166
+ N,
1167
+ K,
1168
+ force_swap_ab=True,
1169
+ )
1170
+
1171
+ ops.cutlass_w4a8_moe_mm(
1172
+ mm1_out,
1173
+ a1q,
1174
+ w1,
1175
+ a1q_scale,
1176
+ w1_chan_scale,
1177
+ w1_scale,
1178
+ group_size,
1179
+ expert_offsets,
1180
+ problem_sizes1,
1181
+ a_strides1,
1182
+ b_strides1,
1183
+ c_strides1,
1184
+ s_strides1,
1185
+ )
1186
+
1187
+ activation_callable(act_out, mm1_out)
1188
+
1189
+ a2q, a2q_scale = ops.scaled_fp8_quant(
1190
+ act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out
1191
+ )
1192
+
1193
+ if expert_map is not None:
1194
+ mm2_out.fill_(0)
1195
+
1196
+ ops.cutlass_w4a8_moe_mm(
1197
+ mm2_out,
1198
+ a2q,
1199
+ w2,
1200
+ a2q_scale,
1201
+ w2_chan_scale,
1202
+ w2_scale,
1203
+ group_size,
1204
+ expert_offsets,
1205
+ problem_sizes2,
1206
+ a_strides2,
1207
+ b_strides2,
1208
+ c_strides2,
1209
+ s_strides2,
1210
+ )
1211
+
1212
+ # for non-chunking mode the output is resized from workspace13
1213
+ # so we need to make sure mm2_out uses workspace2.
1214
+ moe_unpermute(
1215
+ out=output,
1216
+ permuted_hidden_states=mm2_out,
1217
+ topk_weights=topk_weights,
1218
+ inv_permuted_idx=inv_perm,
1219
+ )
1220
+
1221
+
1222
+class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):
1223
+ def __init__(
1224
+ self,
1225
+ out_dtype: torch.dtype | None,
1226
+ a_strides1: torch.Tensor,
1227
+ a_strides2: torch.Tensor,
1228
+ b_strides1: torch.Tensor,
1229
+ b_strides2: torch.Tensor,
1230
+ c_strides1: torch.Tensor,
1231
+ c_strides2: torch.Tensor,
1232
+ s_strides1: torch.Tensor,
1233
+ s_strides2: torch.Tensor,
1234
+ quant_config: FusedMoEQuantConfig,
1235
+ group_size: int,
1236
+ ):
1237
+ super().__init__(quant_config)
1238
+ self.out_dtype = out_dtype
1239
+ self.a_strides1 = a_strides1
1240
+ self.a_strides2 = a_strides2
1241
+ self.b_strides1 = b_strides1
1242
+ self.b_strides2 = b_strides2
1243
+ self.c_strides1 = c_strides1
1244
+ self.c_strides2 = c_strides2
1245
+ self.s_strides1 = s_strides1
1246
+ self.s_strides2 = s_strides2
1247
+ self.group_size = group_size
1248
+
1249
+ @property
1250
+ def activation_formats(
1251
+ self,
1252
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
1253
+ return (
1254
+ mk.FusedMoEActivationFormat.Standard,
1255
+ mk.FusedMoEActivationFormat.Standard,
1256
+ )
1257
+
1258
+ def supports_chunking(self) -> bool:
1259
+ return True
1260
+
1261
+ def supports_expert_map(self) -> bool:
1262
+ return True
1263
+
1264
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1265
+ # topk weights and reduction are fused in moe_unpermute cuda kernel
1266
+ return TopKWeightAndReduceNoOP()
1267
+
1268
+ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
1269
+ return self.out_dtype if self.out_dtype is not None else act_dtype
1270
+
1271
+ def workspace_shapes(
1272
+ self,
1273
+ M: int,
1274
+ N: int,
1275
+ K: int,
1276
+ topk: int,
1277
+ global_num_experts: int,
1278
+ local_num_experts: int,
1279
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1280
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1281
+ workspace1 = (M * topk, max(N, K))
1282
+ workspace2 = (M * topk, max(N // 2, K))
1283
+ output = (M, K)
1284
+ return (workspace1, workspace2, output)
1285
+
1286
+ def apply(
1287
+ self,
1288
+ output: torch.Tensor,
1289
+ hidden_states: torch.Tensor,
1290
+ w1: torch.Tensor,
1291
+ w2: torch.Tensor,
1292
+ topk_weights: torch.Tensor,
1293
+ topk_ids: torch.Tensor,
1294
+ activation: str,
1295
+ global_num_experts: int,
1296
+ expert_map: torch.Tensor | None,
1297
+ a1q_scale: torch.Tensor | None,
1298
+ a2_scale: torch.Tensor | None,
1299
+ workspace13: torch.Tensor | None,
1300
+ workspace2: torch.Tensor | None,
1301
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1302
+ apply_router_weight_on_input: bool,
1303
+ ):
1304
+ assert self.w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
1305
+ assert self.w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
1306
+
1307
+ expert_num_tokens = None
1308
+ activation_callable = lambda o, i: self.activation(activation, o, i)
1309
+
1310
+ use_batched_format = (
1311
+ self.activation_formats[0] == mk.FusedMoEActivationFormat.BatchedExperts
1312
+ )
1313
+ assert not use_batched_format, "batched format not supported"
1314
+
1315
+ in_dtype = hidden_states.dtype
1316
+
1317
+ run_cutlass_moe_w4a8_fp8(
1318
+ output,
1319
+ hidden_states,
1320
+ w1,
1321
+ w2,
1322
+ topk_ids,
1323
+ activation_callable,
1324
+ global_num_experts,
1325
+ expert_map,
1326
+ self.w1_scale,
1327
+ self.w2_scale,
1328
+ a1q_scale,
1329
+ a2_scale,
1330
+ self.g1_alphas, # per-channel scales
1331
+ self.g2_alphas, # per-channel scales
1332
+ self.a_strides1,
1333
+ self.a_strides2,
1334
+ self.b_strides1,
1335
+ self.b_strides2,
1336
+ self.c_strides1,
1337
+ self.c_strides2,
1338
+ self.s_strides1,
1339
+ self.s_strides2,
1340
+ workspace13,
1341
+ workspace2,
1342
+ expert_num_tokens,
1343
+ self.out_dtype if self.out_dtype is not None else in_dtype,
1344
+ self.per_act_token_quant,
1345
+ self.per_out_ch_quant,
1346
+ use_batched_format,
1347
+ topk_weights,
1348
+ self.group_size,
1349
+ )
1350
+
1351
+
1352
+def cutlass_moe_w4a8_fp8(
1353
+ a: torch.Tensor,
1354
+ w1_q: torch.Tensor,
1355
+ w2_q: torch.Tensor,
1356
+ topk_weights: torch.Tensor,
1357
+ topk_ids: torch.Tensor,
1358
+ a_strides1: torch.Tensor,
1359
+ a_strides2: torch.Tensor,
1360
+ b_strides1: torch.Tensor,
1361
+ b_strides2: torch.Tensor,
1362
+ c_strides1: torch.Tensor,
1363
+ c_strides2: torch.Tensor,
1364
+ s_strides1: torch.Tensor,
1365
+ s_strides2: torch.Tensor,
1366
+ quant_config: FusedMoEQuantConfig,
1367
+ activation: str = "silu",
1368
+ expert_map: torch.Tensor | None = None,
1369
+ apply_router_weight_on_input: bool = False,
1370
+ global_num_experts: int = -1,
1371
+ group_size: int = 128,
1372
+) -> torch.Tensor:
1373
+ """
1374
+ This function computes a w4a8-quantized Mixture of Experts (MoE) layer
1375
+ using two sets of quantized weights, w1_q and w2_q, and top-k gating
1376
+ mechanism. The matrix multiplications are implemented with CUTLASS
1377
+ mixed-dtype grouped gemm.
1378
+
1379
+ Parameters:
1380
+ - a (torch.Tensor): The input tensor to the MoE layer.
1381
+ Shape: [M, K]
1382
+ - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
1383
+ Shape: [num_experts, 2*N, K // packed_factor]
1384
+ - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
1385
+ Shape: [num_experts, K, N // packed_factor]
1386
+ - topk_weights (torch.Tensor): The weights of each token->expert mapping.
1387
+ - topk_ids (torch.Tensor): The token->expert mappings.
1388
+ - a_strides1 (torch.Tensor): The input strides for the first gemm.
1389
+ Shape: [num_experts]
1390
+ - a_strides2 (torch.Tensor): The input strides for the second gemm.
1391
+ Shape: [num_experts]
1392
+ - b_strides1 (torch.Tensor): The packed layout for the first gemm weights.
1393
+ Shape: [num_experts, 3]
1394
+ dtype: torch.int32
1395
+ - b_strides2 (torch.Tensor): The packed layout for the second gemm weights.
1396
+ Shape: [num_experts, 3]
1397
+ dtype: torch.int32
1398
+ - c_strides1 (torch.Tensor): The output strides for the first gemm.
1399
+ Shape: [num_experts]
1400
+ - c_strides2 (torch.Tensor): The output strides for the second gemm.
1401
+ Shape: [num_experts]
1402
+ - s_strides1 (torch.Tensor): strides for the group-wise scales for the first gemm.
1403
+ Shape: [num_experts, 2]
1404
+ dtype: torch.int64
1405
+ - s_strides2 (torch.Tensor): strides for the group-wise scales for the second gemm.
1406
+ Shape: [num_experts, 2]
1407
+ dtype: torch.int64
1408
+ - per_act_token (Optional[bool]): Whether the scale is per-token or
1409
+ per-tensor.
1410
+ - activation (str): The activation function to use.
1411
+ - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
1412
+ every Rank is responsible for a subset of experts. expert_map is a
1413
+ mapping from global expert-id to local expert-id. When expert_map[i]
1414
+ is -1, it means that this Rank is not responsible for global
1415
+ expert-id i.
1416
+ - apply_router_weight_on_input (bool): When true, the topk weights are
1417
+ applied directly on the inputs. This is only applicable when topk is 1.
1418
+ - global_num_experts (int): The total number of experts.
1419
+ - group_size (int): The number of weights per scale factor
1420
+
1421
+ Returns:
1422
+ - torch.Tensor: The bf16 output tensor after applying the MoE layer.
1423
+ """
1424
+ assert quant_config is not None
1425
+
1426
+ num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0)
1427
+
1428
+ fn = mk.FusedMoEModularKernel(
1429
+ MoEPrepareAndFinalizeNoEP(),
1430
+ CutlassExpertsW4A8Fp8(
1431
+ out_dtype=a.dtype,
1432
+ a_strides1=a_strides1,
1433
+ a_strides2=a_strides2,
1434
+ b_strides1=b_strides1,
1435
+ b_strides2=b_strides2,
1436
+ c_strides1=c_strides1,
1437
+ c_strides2=c_strides2,
1438
+ s_strides1=s_strides1,
1439
+ s_strides2=s_strides2,
1440
+ quant_config=quant_config,
1441
+ group_size=group_size,
1442
+ ),
1443
+ )
1444
+
1445
+ return fn(
1446
+ a,
1447
+ w1_q,
1448
+ w2_q,
1449
+ topk_weights,
1450
+ topk_ids,
1451
+ activation=activation,
1452
+ global_num_experts=num_experts,
1453
+ expert_map=expert_map,
1454
+ apply_router_weight_on_input=apply_router_weight_on_input,
1455
+ )
vllm/model_executor/layers/fused_moe/modular_kernel.py
+1
-1
@@ -367,7 +367,7 @@ class FusedMoEPrepareAndFinalize(ABC):
367
class FusedMoEPermuteExpertsUnpermute(ABC):
368
"""
369
An abstract base class for the [Permute-Experts-Unpermute] step described
370
- above.
370
+ above.
371
"""
372
373
def __init__(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-1
@@ -256,7 +256,7 @@ class CompressedTensorsConfig(QuantizationConfig):
256
if format is not None
257
else is_activation_quantization_format(quant_format)
258
)
259
- # TODO(czhu): w4a8fp8 is in packed-quantized format
259
+ # w4a8fp8 is in packed-quantized format
260
# but needs input activation quantization
261
input_activations = quant_config.get("input_activations")
262
if act_quant_format or input_activations:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+339
-1
@@ -33,6 +33,7 @@ from vllm.model_executor.layers.fused_moe.config import (
33
FusedMoEQuantConfig,
34
fp8_w8a8_moe_quant_config,
35
int4_w4a16_moe_quant_config,
36
+ int4_w4afp8_moe_quant_config,
37
int8_w8a8_moe_quant_config,
38
int8_w8a16_moe_quant_config,
39
nvfp4_moe_quant_config,
@@ -79,7 +80,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
80
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
81
prepare_moe_fp8_layer_for_marlin,
82
)
82
-from vllm.model_executor.layers.quantization.utils.quant_utils import swizzle_blockscale
83
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
84
+ convert_bf16_scales_to_fp8,
85
+ convert_packed_uint4b8_to_signed_int4_inplace,
86
+ swizzle_blockscale,
87
+)
88
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
89
all_close_1d,
90
normalize_e4m3fn_to_e4m3fnuz,
@@ -204,6 +209,11 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
209
return CompressedTensorsW8A8Int8MoEMethod(
210
weight_quant, input_quant, layer.moe_config
211
)
212
+ elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
213
+ logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod")
214
+ return CompressedTensorsW4A8Fp8MoEMethod(
215
+ weight_quant, input_quant, layer.moe_config
216
+ )
217
elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
218
return CompressedTensorsW4A8Int8MoEMethod(
219
weight_quant, input_quant, layer.moe_config
@@ -2428,3 +2438,331 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2438
apply_router_weight_on_input,
2439
int(_act_kind(activation)),
2440
)
2441
+
2442
+
2443
+class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
2444
+ def __init__(
2445
+ self,
2446
+ weight_quant: QuantizationArgs,
2447
+ input_quant: QuantizationArgs,
2448
+ moe: FusedMoEConfig,
2449
+ layer_name: str | None = None,
2450
+ ):
2451
+ super().__init__(moe)
2452
+ self.weight_quant = weight_quant
2453
+ self.input_quant = input_quant
2454
+
2455
+ self.group_size = self.weight_quant.group_size
2456
+ self.num_bits = self.weight_quant.num_bits
2457
+ self.packed_factor = 32 // self.num_bits
2458
+
2459
+ assert self.weight_quant.symmetric, (
2460
+ "Only symmetric quantization is supported for W4A8 MoE"
2461
+ )
2462
+ assert self.weight_quant.actorder != "group"
2463
+ assert self.group_size == 128, "Only group size 128 supported for W4A8 MoE"
2464
+
2465
+ self.disable_expert_map = False
2466
+ self.layer_name = layer_name
2467
+
2468
+ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
2469
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
2470
+ GroupShape,
2471
+ )
2472
+
2473
+ self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
2474
+
2475
+ def create_weights(
2476
+ self,
2477
+ layer: torch.nn.Module,
2478
+ num_experts: int,
2479
+ hidden_size: int,
2480
+ intermediate_size_per_partition: int,
2481
+ params_dtype: torch.dtype,
2482
+ **extra_weight_attrs,
2483
+ ):
2484
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
2485
+ layer.hidden_size = hidden_size
2486
+ layer.num_experts = num_experts
2487
+ layer.orig_dtype = params_dtype
2488
+ layer.weight_block_size = None
2489
+
2490
+ # requirement for CUTLASS reorder_tensor
2491
+ assert hidden_size % 256 == 0, f"{hidden_size=} must be divisible by 256"
2492
+ assert intermediate_size_per_partition % 256 == 0, (
2493
+ f"{intermediate_size_per_partition=} must be divisible by 256"
2494
+ )
2495
+ # storage type, pack 8xint4 into int32
2496
+ params_dtype = torch.int32
2497
+
2498
+ # WEIGHTS
2499
+ w13_weight_packed = torch.nn.Parameter(
2500
+ torch.empty(
2501
+ num_experts,
2502
+ 2 * intermediate_size_per_partition,
2503
+ hidden_size // self.packed_factor,
2504
+ dtype=params_dtype,
2505
+ ),
2506
+ requires_grad=False,
2507
+ )
2508
+ layer.register_parameter("w13_weight_packed", w13_weight_packed)
2509
+ set_weight_attrs(w13_weight_packed, extra_weight_attrs)
2510
+
2511
+ w2_weight_packed = torch.nn.Parameter(
2512
+ torch.empty(
2513
+ num_experts,
2514
+ hidden_size,
2515
+ intermediate_size_per_partition // self.packed_factor,
2516
+ dtype=params_dtype,
2517
+ ),
2518
+ requires_grad=False,
2519
+ )
2520
+ layer.register_parameter("w2_weight_packed", w2_weight_packed)
2521
+ set_weight_attrs(w2_weight_packed, extra_weight_attrs)
2522
+
2523
+ # SCALES
2524
+ # weight_scale refers to the group-wise scales
2525
+ # they are initially loaded as bf16, we will convert to fp8
2526
+ # after loading
2527
+ w13_weight_scale = torch.nn.Parameter(
2528
+ torch.ones(
2529
+ num_experts,
2530
+ 2 * intermediate_size_per_partition,
2531
+ hidden_size // self.group_size,
2532
+ dtype=layer.orig_dtype,
2533
+ ),
2534
+ requires_grad=False,
2535
+ )
2536
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
2537
+
2538
+ w2_weight_scale = torch.nn.Parameter(
2539
+ torch.ones(
2540
+ num_experts,
2541
+ hidden_size,
2542
+ intermediate_size_per_partition // self.group_size,
2543
+ dtype=layer.orig_dtype,
2544
+ ),
2545
+ requires_grad=False,
2546
+ )
2547
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
2548
+ # Add PER-GROUP quantization for FusedMoE.weight_loader.
2549
+ extra_weight_attrs.update(
2550
+ {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
2551
+ )
2552
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
2553
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
2554
+
2555
+ # weight shapes
2556
+ w2_weight_shape = torch.nn.Parameter(
2557
+ torch.empty(num_experts, 2), requires_grad=False
2558
+ )
2559
+ layer.register_parameter("w2_weight_shape", w2_weight_shape)
2560
+ set_weight_attrs(w2_weight_shape, extra_weight_attrs)
2561
+ w13_weight_shape = torch.nn.Parameter(
2562
+ torch.empty(num_experts, 2), requires_grad=False
2563
+ )
2564
+ layer.register_parameter("w13_weight_shape", w13_weight_shape)
2565
+ set_weight_attrs(w13_weight_shape, extra_weight_attrs)
2566
+
2567
+ # don't use input scales
2568
+ layer.w13_input_scale = None
2569
+ layer.w2_input_scale = None
2570
+
2571
+ def process_weights_after_loading(self, layer):
2572
+ device = layer.w13_weight_packed.device
2573
+
2574
+ # STRIDES
2575
+ # A, C
2576
+ self.a_strides1_c_strides2 = torch.full(
2577
+ (layer.local_num_experts,),
2578
+ layer.hidden_size,
2579
+ device=device,
2580
+ dtype=torch.int64,
2581
+ )
2582
+ self.a_strides2 = torch.full(
2583
+ (layer.local_num_experts,),
2584
+ layer.intermediate_size_per_partition,
2585
+ device=device,
2586
+ dtype=torch.int64,
2587
+ )
2588
+ self.c_strides1 = torch.full(
2589
+ (layer.local_num_experts,),
2590
+ 2 * layer.intermediate_size_per_partition,
2591
+ device=device,
2592
+ dtype=torch.int64,
2593
+ )
2594
+
2595
+ # S (group-wise scales)
2596
+ # sizeof(StrideS) = 16 bytes, so we need to use 2xint64 to encode it
2597
+ self.s_strides1 = torch.zeros(
2598
+ (layer.local_num_experts, 2), device=device, dtype=torch.int64
2599
+ )
2600
+ self.s_strides1[:, 0] = 2 * layer.intermediate_size_per_partition
2601
+
2602
+ self.s_strides2 = torch.zeros(
2603
+ (layer.local_num_experts, 2), device=device, dtype=torch.int64
2604
+ )
2605
+ self.s_strides2[:, 0] = layer.hidden_size
2606
+
2607
+ # encode and reorder weight tensors, and get the layout to pass to
2608
+ # the grouped gemm kernel. `b_strides1/2` specifies the entire layout
2609
+ convert_packed_uint4b8_to_signed_int4_inplace(layer.w13_weight_packed)
2610
+ w13_weight_shuffled, self.b_strides1 = (
2611
+ ops.cutlass_encode_and_reorder_int4b_grouped(layer.w13_weight_packed)
2612
+ )
2613
+ replace_parameter(layer, "w13_weight_packed", w13_weight_shuffled)
2614
+ convert_packed_uint4b8_to_signed_int4_inplace(layer.w2_weight_packed)
2615
+ w2_weight_shuffled, self.b_strides2 = (
2616
+ ops.cutlass_encode_and_reorder_int4b_grouped(layer.w2_weight_packed)
2617
+ )
2618
+ replace_parameter(layer, "w2_weight_packed", w2_weight_shuffled)
2619
+
2620
+ # convert bf16 scales to (fp8_scales, channel_scales)
2621
+ w13_weight_scale, w13_weight_chan_scale = convert_bf16_scales_to_fp8(
2622
+ self.quant_fp8, layer.w13_weight_scale
2623
+ )
2624
+ w2_weight_scale, w2_weight_chan_scale = convert_bf16_scales_to_fp8(
2625
+ self.quant_fp8, layer.w2_weight_scale
2626
+ )
2627
+
2628
+ # register channel scales
2629
+ layer.register_parameter(
2630
+ "w13_weight_chan_scale",
2631
+ torch.nn.Parameter(w13_weight_chan_scale, requires_grad=False),
2632
+ )
2633
+ layer.register_parameter(
2634
+ "w2_weight_chan_scale",
2635
+ torch.nn.Parameter(w2_weight_chan_scale, requires_grad=False),
2636
+ )
2637
+
2638
+ # The scales are stored as (E, N, K // 128) but the kernel expects
2639
+ # (E, K // 128, N) in row-major format, so we need to permute the last 2 dims
2640
+ # and make it contiguous
2641
+ w13_weight_scale_packed = ops.cutlass_pack_scale_fp8(
2642
+ w13_weight_scale.permute(0, 2, 1).contiguous()
2643
+ )
2644
+ replace_parameter(layer, "w13_weight_scale", w13_weight_scale_packed)
2645
+ w2_weight_scale_packed = ops.cutlass_pack_scale_fp8(
2646
+ w2_weight_scale.permute(0, 2, 1).contiguous()
2647
+ )
2648
+ replace_parameter(layer, "w2_weight_scale", w2_weight_scale_packed)
2649
+
2650
+ def maybe_make_prepare_finalize(
2651
+ self,
2652
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
2653
+ ) -> mk.FusedMoEPrepareAndFinalize | None:
2654
+ return super().maybe_make_prepare_finalize(routing_tables)
2655
+
2656
+ def get_fused_moe_quant_config(
2657
+ self, layer: torch.nn.Module
2658
+ ) -> FusedMoEQuantConfig | None:
2659
+ # Store quantization scales; both per-group and per-channel
2660
+ # Note we haven't specified the group size here because
2661
+ # the quant config logic assumes group-wise scaling
2662
+ # and channel-wise scaling are exclusive.
2663
+ return int4_w4afp8_moe_quant_config(
2664
+ w1_scale=layer.w13_weight_scale, # group scale
2665
+ w2_scale=layer.w2_weight_scale, # group scale
2666
+ g1_alphas=layer.w13_weight_chan_scale,
2667
+ g2_alphas=layer.w2_weight_chan_scale,
2668
+ per_act_token_quant=True, # always use dynamc per-token
2669
+ per_out_ch_quant=True, # always use per-channel
2670
+ )
2671
+
2672
+ def select_gemm_impl(
2673
+ self,
2674
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
2675
+ layer: torch.nn.Module,
2676
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
2677
+ assert self.moe_quant_config is not None
2678
+ assert (
2679
+ prepare_finalize.activation_format == FusedMoEActivationFormat.Standard
2680
+ ), "BatchedExperts not supported"
2681
+
2682
+ from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8
2683
+
2684
+ experts: FusedMoEPermuteExpertsUnpermute
2685
+
2686
+ logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__)
2687
+ experts = CutlassExpertsW4A8Fp8(
2688
+ out_dtype=self.moe.in_dtype,
2689
+ a_strides1=self.a_strides1_c_strides2,
2690
+ a_strides2=self.a_strides2,
2691
+ b_strides1=self.b_strides1,
2692
+ b_strides2=self.b_strides2,
2693
+ c_strides1=self.c_strides1,
2694
+ c_strides2=self.a_strides1_c_strides2,
2695
+ s_strides1=self.s_strides1,
2696
+ s_strides2=self.s_strides2,
2697
+ quant_config=self.moe_quant_config,
2698
+ group_size=self.group_size,
2699
+ )
2700
+
2701
+ num_dispatchers = prepare_finalize.num_dispatchers()
2702
+ self.disable_expert_map = (
2703
+ num_dispatchers > 1 or not experts.supports_expert_map()
2704
+ )
2705
+
2706
+ return experts
2707
+
2708
+ def apply(
2709
+ self,
2710
+ layer: torch.nn.Module,
2711
+ x: torch.Tensor,
2712
+ router_logits: torch.Tensor,
2713
+ top_k: int,
2714
+ renormalize: bool,
2715
+ use_grouped_topk: bool = False,
2716
+ topk_group: int | None = None,
2717
+ num_expert_group: int | None = None,
2718
+ global_num_experts: int = -1,
2719
+ expert_map: torch.Tensor | None = None,
2720
+ custom_routing_function: Callable | None = None,
2721
+ scoring_func: str = "softmax",
2722
+ routed_scaling_factor: float = 1.0,
2723
+ e_score_correction_bias: torch.Tensor | None = None,
2724
+ apply_router_weight_on_input: bool = False,
2725
+ activation: str = "silu",
2726
+ enable_eplb: bool = False,
2727
+ expert_load_view: torch.Tensor | None = None,
2728
+ logical_to_physical_map: torch.Tensor | None = None,
2729
+ logical_replica_count: torch.Tensor | None = None,
2730
+ ):
2731
+ if enable_eplb:
2732
+ raise NotImplementedError(
2733
+ "EPLB not supported for `CompressedTensorsW4A8Fp8MoEMethod` yet."
2734
+ )
2735
+ assert self.moe_quant_config is not None
2736
+ topk_weights, topk_ids, _ = layer.select_experts(
2737
+ hidden_states=x,
2738
+ router_logits=router_logits,
2739
+ )
2740
+
2741
+ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
2742
+ cutlass_moe_w4a8_fp8,
2743
+ )
2744
+
2745
+ return cutlass_moe_w4a8_fp8(
2746
+ x,
2747
+ layer.w13_weight_packed,
2748
+ layer.w2_weight_packed,
2749
+ topk_weights,
2750
+ topk_ids,
2751
+ quant_config=self.moe_quant_config,
2752
+ activation=activation,
2753
+ global_num_experts=global_num_experts,
2754
+ expert_map=None if self.disable_expert_map else expert_map,
2755
+ a_strides1=self.a_strides1_c_strides2,
2756
+ a_strides2=self.a_strides2,
2757
+ b_strides1=self.b_strides1,
2758
+ b_strides2=self.b_strides2,
2759
+ c_strides1=self.c_strides1,
2760
+ c_strides2=self.a_strides1_c_strides2,
2761
+ s_strides1=self.s_strides1,
2762
+ s_strides2=self.s_strides2,
2763
+ group_size=self.group_size,
2764
+ )
2765
+
2766
+ @property
2767
+ def supports_eplb(self) -> bool:
2768
+ return False
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py
+4
-11
@@ -128,14 +128,15 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
128
),
129
)
130
131
- # TODO(czhu): allocate the packed fp8 scales memory here?
132
- # the scales will be expanded by 8x via `cutlass_pack_scale_fp8`
131
+ # After loading, we will transform bf16 -> fp8 ->
132
+ # expand by 8x via `cutlass_pack_scale_fp8`
133
+ # and construct per-channel fp32 scales.
134
weight_scale_args = {
135
"weight_loader": weight_loader,
136
"data": torch.empty(
137
output_size_per_partition,
138
scales_and_zp_size,
138
- dtype=torch.float8_e4m3fn,
139
+ dtype=params_dtype,
140
),
141
}
142
@@ -152,17 +153,9 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
153
data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
154
)
155
155
- # per-channel scales
156
- weight_chan_scale = ChannelQuantScaleParameter(
157
- data=torch.empty((output_size_per_partition, 1), dtype=torch.float32),
158
- output_dim=0,
159
- weight_loader=weight_loader,
160
- )
161
-
156
layer.register_parameter("weight_packed", weight)
157
layer.register_parameter("weight_scale", weight_scale)
158
layer.register_parameter("weight_shape", weight_shape)
165
- layer.register_parameter("weight_chan_scale", weight_chan_scale)
159
160
self.kernel = kernel_type(
161
mp_linear_kernel_config,
vllm/model_executor/layers/quantization/kernels/mixed_precision/cutlass.py
+15
-4
@@ -6,7 +6,11 @@ import torch
6
7
from vllm import _custom_ops as ops
8
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
9
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
9
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
10
+ GroupShape,
11
+ convert_bf16_scales_to_fp8,
12
+ convert_packed_uint4b8_to_signed_int4_inplace,
13
+)
14
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
15
from vllm.platforms import current_platform
16
from vllm.scalar_type import scalar_types
@@ -48,7 +52,6 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
52
"CUTLASS W4A8, only supported int4",
53
)
54
51
- # TODO(czhu): support -1 (column-wise)
55
if c.group_size != 128:
56
return False, "Only group_size 128 is supported"
57
@@ -71,9 +74,9 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
74
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
75
# `weight_scale` is: {input_dim = 0, output_dim = 1}
76
def process_weights_after_loading(self, layer: torch.nn.Module):
74
- # TODO(czhu): optimize speed/mem usage
77
def transform_w_q(x):
78
assert isinstance(x, BasevLLMParameter)
79
+ convert_packed_uint4b8_to_signed_int4_inplace(x.data)
80
permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
81
x.data = ops.cutlass_encode_and_reorder_int4b(x.data.t().contiguous().t())
82
return x
@@ -85,10 +88,18 @@ class CutlassW4A8LinearKernel(MPLinearKernel):
88
x.data = ops.cutlass_pack_scale_fp8(x.data)
89
return x
90
91
+ w_s = getattr(layer, self.w_s_name)
92
+ fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data)
93
+ w_s.data = fp8_scales
94
+
95
+ # register per-channel scales
96
+ layer.register_parameter(
97
+ "weight_chan_scale", torch.nn.Parameter(chan_scales, requires_grad=False)
98
+ )
99
+
100
# Encode/reorder weights and pack scales
101
self._transform_param(layer, self.w_q_name, transform_w_q)
102
self._transform_param(layer, self.w_s_name, transform_w_s)
91
- self._transform_param(layer, "weight_chan_scale", lambda x: x)
103
104
def apply_weights(
105
self,
vllm/model_executor/layers/quantization/utils/quant_utils.py
+49
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""This file is used for /tests and /benchmarks"""
4
5
-from collections.abc import Mapping
5
+from collections.abc import Callable, Mapping
6
from dataclasses import dataclass
7
from types import MappingProxyType
8
from typing import ClassVar, NamedTuple
@@ -691,3 +691,51 @@ def cutlass_fp4_supported() -> bool:
691
capability_tuple = current_platform.get_device_capability()
692
capability = -1 if capability_tuple is None else capability_tuple.to_int()
693
return cutlass_scaled_mm_supports_fp4(capability)
694
+
695
+
696
+def convert_bf16_scales_to_fp8(
697
+ quant_fp8: Callable, scales: torch.Tensor
698
+) -> tuple[torch.Tensor, torch.Tensor]:
699
+ """
700
+ Convert a BF16 scale tensor into the pair of (fp8_scales, channel_scales)
701
+ expected by W4A8 GEMM kernels.
702
+ """
703
+ assert scales.is_contiguous(), (
704
+ f"scale tensor must be contiguous, got {scales.stride()=}"
705
+ )
706
+ assert scales.is_cuda, "scales must be on gpu"
707
+
708
+ orig_shape = scales.shape
709
+ k_groups = orig_shape[-1]
710
+ flat_scales = scales.view(-1, k_groups)
711
+
712
+ fp8_scales, chan_scales = quant_fp8(flat_scales)
713
+ fp8_scales = (fp8_scales.float() / 8.0).to(torch.float8_e4m3fn)
714
+ chan_scales *= 8.0
715
+
716
+ # restore original shape
717
+ fp8_scales = fp8_scales.view(orig_shape)
718
+ chan_scales = chan_scales.view(orig_shape[:-1], -1)
719
+
720
+ return fp8_scales, chan_scales
721
+
722
+
723
+def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tensor:
724
+ """
725
+ Convert int4b8 (packed to int32) to signed int4
726
+ """
727
+ assert t.is_cuda, "tensor must be on gpu"
728
+ assert t.dtype == torch.int32, f"expected int32 packed weights but got {t.dtype}"
729
+
730
+ # loop through the 8 4-bit nibbles in each int32 entry
731
+ for i in range(8):
732
+ shift = 4 * i
733
+ # extract the i-th 4-bit nibble
734
+ nib = (t >> shift) & 0xF
735
+ # clear the original nibble by masking out
736
+ t &= ~(0xF << shift)
737
+ # convert int4b8 [0..15] to signed int4 [-8..7] by subtracting 8
738
+ # and update in-place
739
+ t |= ((nib - 8) & 0xF) << shift
740
+
741
+ return t