@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e76e233540
Fork 0

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 14, 05:30 UTC

[kernel] Support W4A8 on Hopper (#23198)

Signed-off-by: czhu-cohere <conway.zhu@cohere.com>

czhu-cohere committed Aug 24, 2025 at 02:18 UTC e76e23354033f167b778f5e49fd384e301681d65
12 files changed +1128 -7
CMakeLists.txt
+27
@@ -750,6 +750,33 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
750 "found in CUDA target architectures")
751 endif()
752 endif()
753 +
754 + # Only build W4A8 kernels if we are building for something compatible with sm90a
755 + cuda_archs_loose_intersection(W4A8_ARCHS "9.0a" "${CUDA_ARCHS}")
756 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND W4A8_ARCHS)
757 + set(SRCS
758 + "csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu")
759 +
760 + set_gencode_flags_for_srcs(
761 + SRCS "${SRCS}"
762 + CUDA_ARCHS "${W4A8_ARCHS}")
763 +
764 + list(APPEND VLLM_EXT_SRC "${SRCS}")
765 +
766 + message(STATUS "Building W4A8 kernels for archs: ${W4A8_ARCHS}")
767 + else()
768 + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0
769 + AND W4A8_ARCHS)
770 + message(STATUS "Not building W4A8 kernels as CUDA Compiler version is "
771 + "not >= 12.0, we recommend upgrading to CUDA 12.0 or "
772 + "later if you intend on running w4a16 quantized models on "
773 + "Hopper.")
774 + else()
775 + message(STATUS "Not building W4A8 kernels as no compatible archs "
776 + "found in CUDA target architectures")
777 + endif()
778 + endif()
779 +
780 # if CUDA endif
781 endif()
782
benchmarks/kernels/benchmark_machete.py
+33
@@ -284,6 +284,25 @@ def machete_create_bench_fn(
284 )
285
286
287 +def cutlass_w4a8_create_bench_fn(
288 + bt: BenchmarkTensors, out_type=torch.dtype, schedule=None
289 +) -> Callable:
290 + w_q = bt.w_q.t().contiguous().t() # make col major
291 + w_q = ops.cutlass_encode_and_reorder_int4b(w_q)
292 + # expects fp8 scales
293 + w_s = ops.cutlass_pack_scale_fp8(bt.w_g_s.to(torch.float8_e4m3fn))
294 +
295 + return lambda: ops.cutlass_w4a8_mm(
296 + a=bt.a,
297 + b_q=w_q,
298 + b_group_scales=w_s,
299 + b_group_size=bt.group_size,
300 + b_channel_scales=bt.w_ch_s,
301 + a_token_scales=bt.w_tok_s,
302 + maybe_schedule=schedule,
303 + )
304 +
305 +
306 # impl
307
308 # bench
@@ -385,6 +404,20 @@ def bench(
404 )
405 )
406
407 + # cutlass w4a8
408 + if types.act_type == torch.float8_e4m3fn and group_size == 128:
409 + timers.append(
410 + bench_fns(
411 + label,
412 + sub_label,
413 + f"cutlass w4a8 ({name_type_string})",
414 + [
415 + cutlass_w4a8_create_bench_fn(bt, out_type=types.output_type)
416 + for bt in benchmark_tensors
417 + ],
418 + )
419 + )
420 +
421 if sweep_schedules:
422 global _SWEEP_SCHEDULES_RESULTS
423
benchmarks/kernels/weight_shapes.py
+6
@@ -95,4 +95,10 @@ WEIGHT_SHAPES = {
95 ([2048, 2816], 1),
96 ([1408, 2048], 0),
97 ],
98 + "CohereLabs/c4ai-command-a-03-2025": [
99 + ([12288, 14336], 1),
100 + ([12288, 12288], 0),
101 + ([12288, 73728], 1),
102 + ([36864, 12288], 0),
103 + ],
104 }
csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu new
+418
@@ -0,0 +1,418 @@
1 +//
2 +// Based off of:
3 +// https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu
4 +//
5 +
6 +#include <ATen/cuda/CUDAContext.h>
7 +#include <c10/cuda/CUDAGuard.h>
8 +#include <torch/all.h>
9 +#include "cutlass_extensions/torch_utils.hpp"
10 +
11 +#include "core/registration.h"
12 +
13 +#include "cutlass/cutlass.h"
14 +
15 +#include "cute/tensor.hpp"
16 +#include "cutlass/gemm/collective/collective_builder.hpp"
17 +#include "cutlass/epilogue/collective/collective_builder.hpp"
18 +#include "cutlass/gemm/device/gemm_universal_adapter.h"
19 +#include "cutlass/gemm/kernel/gemm_universal.hpp"
20 +
21 +#include "cutlass/util/packed_stride.hpp"
22 +#include "cutlass/util/mixed_dtype_utils.hpp"
23 +
24 +#include "cutlass_extensions/common.hpp"
25 +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
26 +
27 +namespace vllm::cutlass_w4a8 {
28 +
29 +using namespace cute;
30 +
31 +// -------------------------------------------------------------------------------------
32 +// Static configuration shared across all instantiations
33 +// -------------------------------------------------------------------------------------
34 +using MmaType = cutlass::float_e4m3_t; // A/scale element type
35 +using QuantType = cutlass::int4b_t; // B element type (packed int4)
36 +
37 +static int constexpr TileShapeK = 128 * 8 / sizeof_bits<MmaType>::value;
38 +static int constexpr ScalePackSize = 8; // pack 8 scale elements together
39 +static int constexpr PackFactor = 8; // 8 4-bit packed into int32
40 +
41 +// A matrix configuration
42 +using ElementA = MmaType; // Element type for A matrix operand
43 +using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand
44 +using LayoutA_Transpose =
45 + typename cutlass::layout::LayoutTranspose<LayoutA>::type;
46 +constexpr int AlignmentA =
47 + 128 / cutlass::sizeof_bits<
48 + ElementA>::value; // Memory access granularity/alignment of A
49 + // matrix in units of elements (up to 16 bytes)
50 +using StrideA = cutlass::detail::TagToStrideA_t<LayoutA>;
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 +using LayoutB_Transpose =
57 + typename cutlass::layout::LayoutTranspose<LayoutB>::type;
58 +constexpr int AlignmentB =
59 + 128 / cutlass::sizeof_bits<
60 + ElementB>::value; // Memory access granularity/alignment of B
61 + // matrix in units of elements (up to 16 bytes)
62 +using StrideB = cutlass::detail::TagToStrideB_t<LayoutB>;
63 +
64 +// Define the CuTe layout for reordered quantized tensor B
65 +// LayoutAtomQuant places values that will be read by the same thread in
66 +// contiguous locations in global memory. It specifies the reordering within a
67 +// single warp's fragment
68 +using LayoutAtomQuant =
69 + decltype(cutlass::compute_memory_reordering_atom<MmaType>());
70 +using LayoutB_Reordered = decltype(cute::tile_to_shape(
71 + LayoutAtomQuant{}, Layout<Shape<int, int, int>, StrideB>{}));
72 +
73 +// Group-wise scales
74 +using ElementScale = MmaType;
75 +using LayoutScale = cutlass::layout::RowMajor;
76 +
77 +// Per-tok, per-chan scales
78 +using ElementSChannel = float;
79 +
80 +// C/D matrix configuration
81 +using ElementC =
82 + cutlass::bfloat16_t; // Element type for C and D matrix operands
83 +using LayoutC =
84 + cutlass::layout::RowMajor; // Layout type for C and D matrix operands
85 +constexpr int AlignmentC =
86 + 128 / cutlass::sizeof_bits<
87 + ElementC>::value; // Memory access granularity/alignment of C
88 + // matrix in units of elements (up to 16 bytes)
89 +
90 +using ElementD = ElementC;
91 +using LayoutD = LayoutC;
92 +constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
93 +
94 +// Core kernel configurations
95 +using ElementAccumulator = float; // Element type for internal accumulation
96 +using ElementCompute = float; // Element type for epilogue computation
97 +using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that
98 + // supports the intended feature
99 +using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag
100 +using KernelSchedule =
101 + cutlass::gemm::KernelTmaWarpSpecializedCooperative; // Kernel to launch
102 + // based on the default
103 + // setting in the
104 + // Collective Builder
105 +using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
106 +using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
107 +
108 +// ----------------------------------------------------------------------------
109 +// Kernel template — Tile/Cluster shapes
110 +// ----------------------------------------------------------------------------
111 +template <class TileShape_MN, class ClusterShape_MNK>
112 +struct W4A8GemmKernel {
113 + using TileShape =
114 + decltype(cute::append(TileShape_MN{}, cute::Int<TileShapeK>{}));
115 + using ClusterShape = ClusterShape_MNK;
116 +
117 + // Epilogue per-tok, per-chan scales
118 + using ChTokScalesEpilogue =
119 + typename vllm::c3x::ScaledEpilogue<ElementAccumulator, ElementD,
120 + TileShape>;
121 + using EVTCompute = typename ChTokScalesEpilogue::EVTCompute;
122 + using CollectiveEpilogue =
123 + typename cutlass::epilogue::collective::CollectiveBuilder<
124 + ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType,
125 + ElementAccumulator, ElementSChannel,
126 + // Transpose layout of D here since we use explicit swap + transpose
127 + // the void type for C tells the builder to allocate 0 smem for the C
128 + // matrix. We can enable this if beta == 0 by changing ElementC to
129 + // void below.
130 + ElementC, typename cutlass::layout::LayoutTranspose<LayoutC>::type,
131 + AlignmentC, ElementD,
132 + typename cutlass::layout::LayoutTranspose<LayoutD>::type, AlignmentD,
133 + EpilogueSchedule, // This is the only epi supporting the required
134 + // swap + transpose.
135 + EVTCompute>::CollectiveOp;
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, ScalePackSize>>,
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 + Shape<int, int, int, int>, // Indicates ProblemShape
152 + CollectiveMainloopShuffled, CollectiveEpilogue>;
153 + using GemmShuffled =
154 + cutlass::gemm::device::GemmUniversalAdapter<GemmKernelShuffled>;
155 +
156 + using StrideC = typename GemmKernelShuffled::StrideC;
157 + using StrideD = typename GemmKernelShuffled::StrideD;
158 + using StrideS = typename CollectiveMainloopShuffled::StrideScale;
159 +
160 + static torch::Tensor mm(torch::Tensor const& A,
161 + torch::Tensor const& B, // already packed
162 + torch::Tensor const& group_scales, // already packed
163 + int64_t group_size,
164 + torch::Tensor const& channel_scales,
165 + torch::Tensor const& token_scales,
166 + std::optional<at::ScalarType> const& maybe_out_type) {
167 + // TODO: param validation
168 + int m = A.size(0);
169 + int k = A.size(1);
170 + int n = B.size(1);
171 +
172 + // Allocate output
173 + const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
174 + auto device = A.device();
175 + auto stream = at::cuda::getCurrentCUDAStream(device.index());
176 + torch::Tensor D =
177 + torch::empty({m, n}, torch::TensorOptions()
178 + .dtype(equivalent_scalar_type_v<ElementD>)
179 + .device(device));
180 + // prepare arg pointers
181 + auto A_ptr = static_cast<MmaType const*>(A.const_data_ptr());
182 + auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
183 + auto D_ptr = static_cast<ElementD*>(D.data_ptr());
184 + // can we avoid harcode the 8 here
185 + auto S_ptr =
186 + static_cast<cutlass::Array<ElementScale, ScalePackSize> const*>(
187 + group_scales.const_data_ptr());
188 +
189 + // runtime layout for B
190 + auto shape_B = cute::make_shape(n, k, 1);
191 + LayoutB_Reordered layout_B_reordered =
192 + cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
193 +
194 + // strides
195 + int const scale_k = cutlass::ceil_div(k, group_size);
196 + StrideA stride_A =
197 + cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
198 + // Reverse stride here due to swap and transpose
199 + StrideD stride_D =
200 + cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(n, m, 1));
201 + StrideS stride_S = cutlass::make_cute_packed_stride(
202 + StrideS{}, cute::make_shape(n, scale_k, 1));
203 +
204 + // Create a structure of gemm kernel arguments suitable for invoking an
205 + // instance of Gemm auto arguments =
206 + // args_from_options<GemmShuffled>(options);
207 + /// Populates a Gemm::Arguments structure from the given arguments
208 + /// Swap the A and B tensors, as well as problem shapes here.
209 + using Args = typename GemmShuffled::Arguments;
210 + using MainloopArguments = typename GemmKernelShuffled::MainloopArguments;
211 + using EpilogueArguments = typename GemmKernelShuffled::EpilogueArguments;
212 +
213 + MainloopArguments mainloop_arguments{
214 + B_ptr, layout_B_reordered, A_ptr, stride_A,
215 + S_ptr, stride_S, group_size};
216 +
217 + EpilogueArguments epilogue_arguments{
218 + ChTokScalesEpilogue::prepare_args(channel_scales, token_scales),
219 + nullptr,
220 + {}, // no C
221 + D_ptr,
222 + stride_D};
223 +
224 + Args arguments{cutlass::gemm::GemmUniversalMode::kGemm,
225 + {n, m, k, 1}, // shape
226 + mainloop_arguments,
227 + epilogue_arguments};
228 +
229 + // Workspace
230 + size_t workspace_size = GemmShuffled::get_workspace_size(arguments);
231 + torch::Tensor workspace =
232 + torch::empty(workspace_size,
233 + torch::TensorOptions().dtype(torch::kU8).device(device));
234 +
235 + // Run GEMM
236 + GemmShuffled gemm;
237 + CUTLASS_CHECK(gemm.can_implement(arguments));
238 + CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
239 + CUTLASS_CHECK(gemm.run(stream));
240 +
241 + return D;
242 + }
243 +};
244 +
245 +// ----------------------------------------------------------------------------
246 +// Kernel instantiations and dispatch logic
247 +// ----------------------------------------------------------------------------
248 +using Kernel_256x128_1x1x1 =
249 + W4A8GemmKernel<Shape<_256, _128>, Shape<_1, _1, _1>>;
250 +using Kernel_256x64_1x1x1 = W4A8GemmKernel<Shape<_256, _64>, Shape<_1, _1, _1>>;
251 +using Kernel_256x32_1x1x1 = W4A8GemmKernel<Shape<_256, _32>, Shape<_1, _1, _1>>;
252 +using Kernel_256x16_1x1x1 = W4A8GemmKernel<Shape<_256, _16>, Shape<_1, _1, _1>>;
253 +using Kernel_128x256_2x1x1 =
254 + W4A8GemmKernel<Shape<_128, _256>, Shape<_2, _1, _1>>;
255 +using Kernel_128x256_1x1x1 =
256 + W4A8GemmKernel<Shape<_128, _256>, Shape<_1, _1, _1>>;
257 +using Kernel_128x128_1x1x1 =
258 + W4A8GemmKernel<Shape<_128, _128>, Shape<_1, _1, _1>>;
259 +using Kernel_128x64_1x1x1 = W4A8GemmKernel<Shape<_128, _64>, Shape<_1, _1, _1>>;
260 +using Kernel_128x32_1x1x1 = W4A8GemmKernel<Shape<_128, _32>, Shape<_1, _1, _1>>;
261 +using Kernel_128x16_1x1x1 = W4A8GemmKernel<Shape<_128, _16>, Shape<_1, _1, _1>>;
262 +
263 +torch::Tensor mm_dispatch(torch::Tensor const& A,
264 + torch::Tensor const& B, // already packed
265 + torch::Tensor const& group_scales, // already packed
266 + int64_t group_size,
267 + torch::Tensor const& channel_scales,
268 + torch::Tensor const& token_scales,
269 + std::optional<at::ScalarType> const& maybe_out_type,
270 + const std::string& schedule) {
271 + if (schedule == "256x128_1x1x1") {
272 + return Kernel_256x128_1x1x1::mm(A, B, group_scales, group_size,
273 + channel_scales, token_scales,
274 + maybe_out_type);
275 + } else if (schedule == "256x64_1x1x1") {
276 + return Kernel_256x64_1x1x1::mm(A, B, group_scales, group_size,
277 + channel_scales, token_scales,
278 + maybe_out_type);
279 + } else if (schedule == "256x32_1x1x1") {
280 + return Kernel_256x32_1x1x1::mm(A, B, group_scales, group_size,
281 + channel_scales, token_scales,
282 + maybe_out_type);
283 + } else if (schedule == "256x16_1x1x1") {
284 + return Kernel_256x16_1x1x1::mm(A, B, group_scales, group_size,
285 + channel_scales, token_scales,
286 + maybe_out_type);
287 + } else if (schedule == "128x256_2x1x1") {
288 + return Kernel_128x256_2x1x1::mm(A, B, group_scales, group_size,
289 + channel_scales, token_scales,
290 + maybe_out_type);
291 + } else if (schedule == "128x256_1x1x1") {
292 + return Kernel_128x256_1x1x1::mm(A, B, group_scales, group_size,
293 + channel_scales, token_scales,
294 + maybe_out_type);
295 + } else if (schedule == "128x128_1x1x1") {
296 + return Kernel_128x128_1x1x1::mm(A, B, group_scales, group_size,
297 + channel_scales, token_scales,
298 + maybe_out_type);
299 + } else if (schedule == "128x64_1x1x1") {
300 + return Kernel_128x64_1x1x1::mm(A, B, group_scales, group_size,
301 + channel_scales, token_scales,
302 + maybe_out_type);
303 + } else if (schedule == "128x32_1x1x1") {
304 + return Kernel_128x32_1x1x1::mm(A, B, group_scales, group_size,
305 + channel_scales, token_scales,
306 + maybe_out_type);
307 + } else if (schedule == "128x16_1x1x1") {
308 + return Kernel_128x16_1x1x1::mm(A, B, group_scales, group_size,
309 + channel_scales, token_scales,
310 + maybe_out_type);
311 + }
312 + TORCH_CHECK(false, "Unknown W4A8 schedule: ", schedule);
313 + return {};
314 +}
315 +
316 +torch::Tensor mm(torch::Tensor const& A,
317 + torch::Tensor const& B, // already packed
318 + torch::Tensor const& group_scales, // already packed
319 + int64_t group_size, torch::Tensor const& channel_scales,
320 + torch::Tensor const& token_scales,
321 + std::optional<at::ScalarType> const& maybe_out_type,
322 + std::optional<std::string> maybe_schedule) {
323 + // requested a specific schedule
324 + if (maybe_schedule) {
325 + return mm_dispatch(A, B, group_scales, group_size, channel_scales,
326 + token_scales, maybe_out_type, *maybe_schedule);
327 + }
328 + std::string schedule;
329 + int M = A.size(0);
330 + int K = A.size(1);
331 + int N = B.size(1);
332 + // heuristic
333 + if (M <= 16) {
334 + schedule = (K == 16384 && N == 18432) ? "256x16_1x1x1" : "128x16_1x1x1";
335 + } else if (M <= 32) {
336 + schedule = (K == 16384 && N == 18432) ? "256x32_1x1x1" : "128x32_1x1x1";
337 + } else if (M <= 64) {
338 + if (K == 16384 && N == 18432)
339 + schedule = "256x64_1x1x1";
340 + else if (N <= 8192 && K <= 8192)
341 + schedule = "128x32_1x1x1";
342 + else
343 + schedule = "128x64_1x1x1";
344 + } else if (M <= 128) {
345 + if (K == 16384 && N == 18432)
346 + schedule = "256x128_1x1x1";
347 + else if (N <= 8192)
348 + schedule = "128x64_1x1x1";
349 + else
350 + schedule = "128x128_1x1x1";
351 + } else if (M <= 256) {
352 + if (N <= 4096)
353 + schedule = "128x64_1x1x1";
354 + else if (N <= 8192)
355 + schedule = "128x128_1x1x1";
356 + else
357 + schedule = "128x256_1x1x1";
358 + } else if (M <= 512 && N <= 4096) {
359 + schedule = "128x128_1x1x1";
360 + } else if (M <= 1024) {
361 + schedule = "128x256_1x1x1";
362 + } else {
363 + schedule = "128x256_2x1x1";
364 + }
365 + return mm_dispatch(A, B, group_scales, group_size, channel_scales,
366 + token_scales, maybe_out_type, schedule);
367 +}
368 +
369 +// ----------------------------------------------------------------------------
370 +// Pre-processing utils
371 +// ----------------------------------------------------------------------------
372 +torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
373 + TORCH_CHECK(scales.dtype() == torch::kFloat8_e4m3fn);
374 + TORCH_CHECK(scales.is_contiguous());
375 + TORCH_CHECK(scales.is_cuda());
376 +
377 + auto packed_scales = torch::empty(
378 + {scales.numel() * ScalePackSize},
379 + torch::TensorOptions().dtype(scales.dtype()).device(scales.device()));
380 + auto scales_ptr = static_cast<MmaType const*>(scales.const_data_ptr());
381 + auto packed_scales_ptr =
382 + static_cast<cutlass::Array<ElementScale, ScalePackSize>*>(
383 + packed_scales.data_ptr());
384 +
385 + cutlass::pack_scale_fp8(scales_ptr, packed_scales_ptr, scales.numel());
386 +
387 + return packed_scales;
388 +}
389 +
390 +torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
391 + TORCH_CHECK(B.dtype() == torch::kInt32);
392 + TORCH_CHECK(B.dim() == 2);
393 +
394 + torch::Tensor B_packed = torch::empty_like(B);
395 +
396 + int k = B.size(0) * PackFactor; // logical k
397 + int n = B.size(1);
398 +
399 + auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
400 + auto B_packed_ptr = static_cast<QuantType*>(B_packed.data_ptr());
401 + auto shape_B = cute::make_shape(n, k, 1);
402 + auto layout_B = make_layout(shape_B, LayoutRight{}); // row major
403 + LayoutB_Reordered layout_B_reordered =
404 + cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
405 +
406 + cutlass::unified_encode_int4b(B_ptr, B_packed_ptr, n * k);
407 + cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
408 +
409 + return B_packed;
410 +}
411 +
412 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
413 + m.impl("cutlass_w4a8_mm", &mm);
414 + m.impl("cutlass_pack_scale_fp8", &pack_scale_fp8);
415 + m.impl("cutlass_encode_and_reorder_int4b", &encode_and_reorder_int4b);
416 +}
417 +
418 +} // namespace vllm::cutlass_w4a8
\ No newline at end of file
csrc/torch_bindings.cpp
+20
@@ -309,6 +309,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
309 "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
310 "SymInt size_n, int num_bits) -> Tensor");
311 // conditionally compiled so impl registrations are in source file
312 +
313 + // CUTLASS w4a8 GEMM
314 + ops.def(
315 + "cutlass_w4a8_mm("
316 + " Tensor A,"
317 + " Tensor B,"
318 + " Tensor group_scales,"
319 + " int group_size,"
320 + " Tensor channel_scales,"
321 + " Tensor token_scales,"
322 + " ScalarType? out_type,"
323 + " str? maybe_schedule"
324 + ") -> Tensor",
325 + {stride_tag});
326 + // pack scales
327 + ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
328 + // encode and reorder weight matrix
329 + ops.def("cutlass_encode_and_reorder_int4b(Tensor B) -> Tensor");
330 + // conditionally compiled so impl registration is in source file
331 +
332 #endif
333
334 // Dequantization for GGML.
tests/kernels/quantization/test_cutlass_w4a8.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the CUTLASS W4A8 kernel.
4 +
5 +Run `pytest tests/kernels/test_cutlass_w4a8.py`.
6 +"""
7 +
8 +from dataclasses import dataclass
9 +from typing import Optional
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm import _custom_ops as ops
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + pack_rows, quantize_weights)
17 +from vllm.platforms import current_platform
18 +from vllm.scalar_type import ScalarType, scalar_types
19 +
20 +# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
21 +# unit tests to a common utility function. Currently the use of
22 +# `is_quant_method_supported` conflates kernels with quantization methods
23 +# an assumption which is breaking down as quantizations methods can have
24 +# have kernels and some kernels support multiple quantization methods.
25 +IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9
26 +
27 +MNK_SHAPES = [(1, 128, 128), (1, 512, 1024), (1, 4096, 4096), (1, 8192, 28672),
28 + (13, 8192, 4096), (26, 4096, 8192), (64, 4096, 4096),
29 + (64, 8192, 28672), (257, 128, 4096), (257, 4096, 4096),
30 + (1024, 4096, 8192), (1024, 8192, 4096)]
31 +
32 +# TODO(czhu): get supported schedules from fn
33 +SCHEDULES = [
34 + '128x16_1x1x1', '256x16_1x1x1', '128x32_1x1x1', '256x32_1x1x1',
35 + '128x64_1x1x1', '256x64_1x1x1', '128x128_1x1x1', '256x128_1x1x1',
36 + '128x256_1x1x1', '128x256_2x1x1'
37 +]
38 +
39 +
40 +@dataclass
41 +class TypeConfig:
42 + act_type: torch.dtype
43 + weight_type: ScalarType
44 + output_type: Optional[torch.dtype]
45 + group_scale_type: Optional[torch.dtype]
46 + channel_scale_type: Optional[torch.dtype]
47 + token_scale_type: Optional[torch.dtype]
48 +
49 +
50 +@dataclass
51 +class Tensors:
52 + w_ref: torch.Tensor
53 + a_ref: torch.Tensor
54 + a: torch.Tensor
55 + w_q: torch.Tensor
56 + w_g_s: torch.Tensor
57 + w_ch_s: torch.Tensor
58 + w_tok_s: torch.Tensor
59 +
60 +
61 +# (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints,
62 +# Ch Scales Type, Tok Scales Type)
63 +TestTypeTuple = tuple[list[torch.dtype], ScalarType, Optional[torch.dtype],
64 + Optional[torch.dtype], bool]
65 +TEST_TYPES = [
66 + *(
67 + TypeConfig(act_type=torch.float8_e4m3fn,
68 + weight_type=w_type,
69 + output_type=o_type,
70 + group_scale_type=torch.float8_e4m3fn,
71 + channel_scale_type=torch.float32,
72 + token_scale_type=torch.float32)
73 + for w_type in [scalar_types.int4]
74 + # TODO(czhu): fp16 out type
75 + for o_type in [torch.bfloat16]),
76 +]
77 +
78 +# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
79 +# unit tests to a common utility function. Currently the use of
80 +# `is_quant_method_supported` conflates kernels with quantization methods
81 +# an assumption which is breaking down as quantizations methods can have
82 +# have kernels and some kernels support multiple quantization methods.
83 +IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90)
84 +
85 +
86 +# For testing quantized linear kernels
87 +def to_fp8(tensor: torch.Tensor):
88 + finfo = torch.finfo(torch.float8_e4m3fn)
89 + return tensor.clamp(min=finfo.min,
90 + max=finfo.max).to(dtype=torch.float8_e4m3fn)
91 +
92 +
93 +def cutlass_quantize_and_pack(atype: torch.dtype,
94 + w: torch.Tensor,
95 + wtype: ScalarType,
96 + stype: Optional[torch.dtype],
97 + group_size: Optional[int],
98 + zero_points: bool = False):
99 + assert wtype.is_integer(), "TODO: support floating point weights"
100 +
101 + w_ref, w_q, w_s, w_zp = quantize_weights(w,
102 + wtype,
103 + group_size=group_size,
104 + zero_points=zero_points)
105 +
106 + # since scales are cast to fp8, we need to compute w_ref this way
107 + w_ref = ((w_q).to(torch.float32) * w_s.to(atype).to(
108 + torch.float32).repeat_interleave(group_size, dim=0)).to(atype)
109 +
110 + # bit mask prevents sign extending int4 when packing
111 + w_q = pack_rows(w_q & 0x0F, wtype.size_bits, *w_q.shape)
112 + w_q = w_q.t().contiguous().t() # convert to col major
113 +
114 + w_q_packed = ops.cutlass_encode_and_reorder_int4b(w_q)
115 + w_s_packed = ops.cutlass_pack_scale_fp8(w_s.to(atype))
116 +
117 + return w_ref, w_q_packed, w_s_packed, w_zp
118 +
119 +
120 +def create_test_tensors(shape: tuple[int, int, int], types: TypeConfig,
121 + group_size: Optional[int]) -> Tensors:
122 + m, n, k = shape
123 +
124 + print("create_test_tensors, shape:", shape, "types:", types, "group_size:",
125 + group_size)
126 +
127 + a = to_fp8(torch.randn((m, k), device="cuda"))
128 + w = to_fp8(torch.randn((k, n), device="cuda"))
129 +
130 + if types.group_scale_type is not None:
131 + w = w.to(types.group_scale_type)
132 + if w.dtype.itemsize == 1:
133 + w = w.to(torch.float16)
134 +
135 + w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack(
136 + a.dtype, w, types.weight_type, types.group_scale_type, group_size,
137 + False)
138 +
139 + a_ref = a.to(torch.float32)
140 + w_ref = w_ref.to(torch.float32)
141 +
142 + # for the practical use case we need per-tok scales for fp8 activations
143 + w_tok_s = torch.randn((m, ), device='cuda', dtype=types.token_scale_type)
144 + # weights are already per-group quantized, use placeholder here
145 + w_ch_s = torch.ones((n, ), device='cuda', dtype=types.channel_scale_type)
146 +
147 + return Tensors(w_ref=w_ref,
148 + a_ref=a_ref,
149 + a=a,
150 + w_q=w_q_packed,
151 + w_g_s=w_s,
152 + w_ch_s=w_ch_s,
153 + w_tok_s=w_tok_s)
154 +
155 +
156 +def mm_test_helper(types: TypeConfig,
157 + tensors: Tensors,
158 + group_size: Optional[int] = None,
159 + schedule: Optional[str] = None):
160 + # CUTLASS upstream uses fp8 with fastaccum as reference
161 + # https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu#L406
162 + output_ref = torch._scaled_mm(
163 + tensors.a_ref.to(types.act_type),
164 + tensors.w_ref.to(types.act_type).t().contiguous().t(), # col major
165 + tensors.w_tok_s.unsqueeze(1),
166 + tensors.w_ch_s.unsqueeze(0),
167 + out_dtype=types.output_type,
168 + use_fast_accum=True)
169 +
170 + output = ops.cutlass_w4a8_mm(
171 + a=tensors.a,
172 + b_q=tensors.w_q,
173 + b_group_scales=tensors.w_g_s,
174 + b_group_size=group_size,
175 + b_channel_scales=tensors.w_ch_s,
176 + a_token_scales=tensors.w_tok_s,
177 + )
178 +
179 + print(output)
180 + print(output_ref)
181 +
182 + torch.testing.assert_close(output,
183 + output_ref.to(output.dtype),
184 + rtol=1e-3,
185 + atol=1e-3)
186 +
187 +
188 +@pytest.mark.skipif(not IS_SUPPORTED_BY_GPU,
189 + reason="CUTLASS W4A8 is not supported on this GPU type.")
190 +@pytest.mark.parametrize("shape",
191 + MNK_SHAPES,
192 + ids=lambda x: "x".join(str(v) for v in x))
193 +@pytest.mark.parametrize("types", TEST_TYPES)
194 +@pytest.mark.parametrize("schedule", SCHEDULES)
195 +def test_cutlass_w4a8(shape, types: TypeConfig, schedule):
196 + group_sizes = [128]
197 + for group_size in group_sizes:
198 + tensors = create_test_tensors(shape, types, group_size)
199 + mm_test_helper(types, tensors, group_size, schedule)
200 +
201 +
202 +# Test to make sure cuda graphs work
203 +class W4A8Layer(torch.nn.Module):
204 +
205 + def __init__(self, **kwargs):
206 + super().__init__()
207 + self.kwargs = kwargs
208 +
209 + def forward(self, a):
210 + return ops.cutlass_w4a8_mm(a=a, **self.kwargs)
211 +
212 +
213 +@pytest.mark.skipif(not IS_SUPPORTED_BY_GPU,
214 + reason="CUTLASS W4A8 is not supported on this GPU type.")
215 +def test_w4a8_cuda_graph():
216 + m, n, k = 512, 4096, 4096
217 +
218 + a = to_fp8(torch.randn((m, k), device="cuda"))
219 + b = to_fp8(torch.randn((k, n), device="cuda"))
220 +
221 + wtype = scalar_types.int4
222 + stype = torch.float8_e4m3fn
223 + group_size = 128
224 + zero_points = False
225 +
226 + w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack(
227 + a.dtype, b.to(torch.float16), wtype, stype, group_size, zero_points)
228 +
229 + w_tok_s = torch.randn((m, ), device='cuda', dtype=torch.float32)
230 + w_ch_s = torch.ones((n, ), device='cuda', dtype=torch.float32)
231 +
232 + # Construct a trivial model with a single layer that calls the kernel
233 + model = W4A8Layer(
234 + b_q=w_q_packed,
235 + b_group_scales=w_s,
236 + b_group_size=group_size,
237 + b_channel_scales=w_ch_s,
238 + a_token_scales=w_tok_s,
239 + )
240 +
241 + output_ref = torch._scaled_mm(
242 + a,
243 + w_ref.to(a.dtype).t().contiguous().t(), # col major
244 + w_tok_s.unsqueeze(1),
245 + w_ch_s.unsqueeze(0),
246 + out_dtype=torch.bfloat16,
247 + use_fast_accum=True)
248 +
249 + # Run the model with a cuda graph
250 + stream = torch.cuda.Stream()
251 + with torch.cuda.stream(stream):
252 + g = torch.cuda.CUDAGraph()
253 + with torch.cuda.graph(g):
254 + output = model(a)
255 +
256 + output.zero_()
257 + g.replay()
258 +
259 + torch.testing.assert_close(output, output_ref, rtol=1e-3, atol=1e-3)
vllm/_custom_ops.py
+48
@@ -474,6 +474,30 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
474 return torch.empty_like(b_q_weight,
475 memory_format=torch.contiguous_format)
476
477 + @register_fake("_C::cutlass_w4a8_mm")
478 + def cutlass_w4a8_mm_fake(
479 + a: torch.Tensor,
480 + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
481 + b_q: torch.Tensor,
482 + b_group_scales: torch.Tensor,
483 + b_group_size: int,
484 + b_channel_scales: torch.Tensor,
485 + a_token_scales: torch.Tensor,
486 + out_type: Optional[torch.dtype] = None,
487 + maybe_schedule: Optional[str] = None) -> torch.Tensor:
488 + m = a.size(0)
489 + n = b_q.size(1)
490 + out_dtype = out_type if out_type is not None else torch.bfloat16
491 + return torch.empty((m, n), device=a.device, dtype=out_dtype)
492 +
493 + @register_fake("_C::cutlass_pack_scale_fp8")
494 + def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor:
495 + return torch.empty_like(scales, memory_format=torch.contiguous_format)
496 +
497 + @register_fake("_C::cutlass_encode_and_reorder_int4b")
498 + def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
499 + return torch.empty_like(b, memory_format=torch.contiguous_format)
500 +
501
502 if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
503
@@ -1032,6 +1056,30 @@ def machete_prepack_B(
1056 group_scales_type)
1057
1058
1059 +# CUTLASS W4A8
1060 +def cutlass_w4a8_mm(
1061 + a: torch.Tensor,
1062 + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
1063 + b_q: torch.Tensor,
1064 + b_group_scales: torch.Tensor,
1065 + b_group_size: int,
1066 + b_channel_scales: torch.Tensor,
1067 + a_token_scales: torch.Tensor,
1068 + out_type: Optional[torch.dtype] = None,
1069 + maybe_schedule: Optional[str] = None) -> torch.Tensor:
1070 + return torch.ops._C.cutlass_w4a8_mm(a, b_q, b_group_scales, b_group_size,
1071 + b_channel_scales, a_token_scales,
1072 + out_type, maybe_schedule)
1073 +
1074 +
1075 +def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor:
1076 + return torch.ops._C.cutlass_pack_scale_fp8(scales)
1077 +
1078 +
1079 +def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor:
1080 + return torch.ops._C.cutlass_encode_and_reorder_int4b(b)
1081 +
1082 +
1083 if hasattr(torch.ops._C, "permute_cols"):
1084
1085 @register_fake("_C::permute_cols")
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+37 -6
@@ -26,10 +26,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
26 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
27 W4A16SPARSE24_SUPPORTED_BITS, WNA16_SUPPORTED_BITS, CompressedTensors24,
28 CompressedTensorsScheme, CompressedTensorsW4A4Fp4,
29 - CompressedTensorsW4A8Int, CompressedTensorsW4A16Fp4,
30 - CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
31 - CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
32 - CompressedTensorsWNA16)
29 + CompressedTensorsW4A8Fp8, CompressedTensorsW4A8Int,
30 + CompressedTensorsW4A16Fp4, CompressedTensorsW4A16Sparse24,
31 + CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8,
32 + CompressedTensorsW8A16Fp8, CompressedTensorsWNA16)
33 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
34 find_matched_target, is_activation_quantization_format,
35 should_ignore_layer)
@@ -200,8 +200,10 @@ class CompressedTensorsConfig(QuantizationConfig):
200 format
201 ) if format is not None else is_activation_quantization_format(
202 quant_format)
203 - if act_quant_format:
204 - input_activations = quant_config.get("input_activations")
203 + # TODO(czhu): w4a8fp8 is in packed-quantized format
204 + # but needs input activation quantization
205 + input_activations = quant_config.get("input_activations")
206 + if act_quant_format or input_activations:
207 # The only case where we have activation quant supported
208 # but no input_activations provided in the config
209 # should be w8a16fp8 w8a16fp8 can also run for cases where
@@ -352,6 +354,28 @@ class CompressedTensorsConfig(QuantizationConfig):
354 input_quant.strategy == QuantizationStrategy.TENSOR)
355 return is_symmetric_activation and is_per_tensor_activation
356
357 + def _is_fp8_w4a8(self, weight_quant: BaseModel,
358 + input_quant: BaseModel) -> bool:
359 + if not weight_quant or not input_quant:
360 + return False
361 + is_weight_4_bits = weight_quant.num_bits == 4
362 + is_activation_8_bits = input_quant.num_bits == 8
363 + weight_strategy = (
364 + weight_quant.strategy == QuantizationStrategy.GROUP.value)
365 + is_token = (weight_strategy and input_quant.strategy
366 + == QuantizationStrategy.TOKEN.value)
367 + is_dynamic = not weight_quant.dynamic and input_quant.dynamic
368 + is_symmetric = weight_quant.symmetric and input_quant.symmetric
369 + # Only per-group symmetric weight (4bit)
370 + # + per-tok symmetric activation (8bit) quantization supported.
371 + return (is_weight_4_bits and is_activation_8_bits and is_token
372 + and is_symmetric and is_dynamic)
373 +
374 + def _is_fp8_w4a8_sm90(self, weight_quant: BaseModel,
375 + input_quant: BaseModel) -> bool:
376 + return (self._check_scheme_supported(90, error=False, match_exact=True)
377 + and self._is_fp8_w4a8(weight_quant, input_quant))
378 +
379 def _is_fp8_w8a8_sm90(self, weight_quant: BaseModel,
380 input_quant: BaseModel) -> bool:
381 return (self._check_scheme_supported(90, error=False, match_exact=True)
@@ -405,6 +429,13 @@ class CompressedTensorsConfig(QuantizationConfig):
429 if self._is_fp4a16_nvfp4(weight_quant, input_quant):
430 return CompressedTensorsW4A16Fp4()
431
432 + if self._is_fp8_w4a8_sm90(weight_quant, input_quant):
433 + return CompressedTensorsW4A8Fp8(num_bits=weight_quant.num_bits,
434 + strategy=weight_quant.strategy,
435 + symmetric=weight_quant.symmetric,
436 + group_size=weight_quant.group_size,
437 + actorder=weight_quant.actorder)
438 +
439 if self._is_wNa16_group_channel(weight_quant, input_quant):
440 if (self.quant_format == CompressionFormat.marlin_24.value
441 and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS):
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+3 -1
@@ -3,6 +3,7 @@
3
4 from .compressed_tensors_scheme import CompressedTensorsScheme
5 from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
6 +from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
7 from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
8 from .compressed_tensors_w4a16_24 import (W4A16SPARSE24_SUPPORTED_BITS,
9 CompressedTensorsW4A16Sparse24)
@@ -21,5 +22,6 @@ __all__ = [
22 "CompressedTensorsW8A8Int8", "CompressedTensorsW8A8Fp8",
23 "WNA16_SUPPORTED_BITS", "W4A16SPARSE24_SUPPORTED_BITS",
24 "CompressedTensors24", "CompressedTensorsW4A16Fp4",
24 - "CompressedTensorsW4A4Fp4", "CompressedTensorsW4A8Int"
25 + "CompressedTensorsW4A4Fp4", "CompressedTensorsW4A8Int",
26 + "CompressedTensorsW4A8Fp8"
27 ]
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py new
+160
@@ -0,0 +1,160 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Callable, Optional
5 +
6 +import torch
7 +from compressed_tensors.quantization import ActivationOrdering
8 +
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11 + CompressedTensorsScheme)
12 +from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
13 + MPLinearLayerConfig, choose_mp_linear_kernel)
14 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
15 + marlin_repeat_scales_on_all_ranks)
16 +# yapf conflicts with isort for this block
17 +# yapf: disable
18 +from vllm.model_executor.parameter import (BasevLLMParameter,
19 + ChannelQuantScaleParameter,
20 + GroupQuantScaleParameter,
21 + PackedvLLMParameter)
22 +# yapf: enable
23 +from vllm.scalar_type import scalar_types
24 +
25 +logger = init_logger(__name__)
26 +
27 +__all__ = ["CompressedTensorsW4A8Fp8"]
28 +W4A8_SUPPORTED_TYPES_MAP = {
29 + 4: scalar_types.int4,
30 +}
31 +W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys())
32 +
33 +
34 +class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):
35 + _kernel_backends_being_used: set[str] = set()
36 +
37 + def __init__(self,
38 + strategy: str,
39 + num_bits: int,
40 + group_size: Optional[int] = None,
41 + symmetric: Optional[bool] = True,
42 + actorder: Optional[ActivationOrdering] = None):
43 +
44 + self.pack_factor = 32 // num_bits
45 + self.strategy = strategy
46 + self.symmetric = symmetric
47 + self.group_size = -1 if group_size is None else group_size
48 + self.has_g_idx = actorder == ActivationOrdering.GROUP
49 +
50 + if self.group_size != 128 or self.strategy != "group":
51 + raise ValueError("W4A8 kernels require group quantization " \
52 + "with group size 128")
53 +
54 + if num_bits not in W4A8_SUPPORTED_TYPES_MAP:
55 + raise ValueError(
56 + f"Unsupported num_bits = {num_bits}. "
57 + f"Supported num_bits = {W4A8_SUPPORTED_TYPES_MAP.keys()}")
58 +
59 + self.quant_type = W4A8_SUPPORTED_TYPES_MAP[num_bits]
60 +
61 + @classmethod
62 + def get_min_capability(cls) -> int:
63 + # hopper
64 + return 90
65 +
66 + def create_weights(self, layer: torch.nn.Module, output_size: int,
67 + input_size: int, output_partition_sizes: list[int],
68 + input_size_per_partition: int,
69 + params_dtype: torch.dtype, weight_loader: Callable,
70 + **kwargs):
71 +
72 + output_size_per_partition = sum(output_partition_sizes)
73 +
74 + mp_linear_kernel_config = MPLinearLayerConfig(
75 + full_weight_shape=(input_size, output_size),
76 + partition_weight_shape=\
77 + (input_size_per_partition, output_size_per_partition),
78 + weight_type=self.quant_type,
79 + act_type=torch.float8_e4m3fn, # always use fp8(e4m3)
80 + group_size=self.group_size,
81 + zero_points=not self.symmetric,
82 + has_g_idx=self.has_g_idx
83 + )
84 +
85 + kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
86 +
87 + if kernel_type.__name__ not in self._kernel_backends_being_used:
88 + logger.info("Using %s for CompressedTensorsW4A8Fp8",
89 + kernel_type.__name__)
90 + self._kernel_backends_being_used.add(kernel_type.__name__)
91 +
92 + # If group_size is -1, we are in channelwise case.
93 + group_size = self.group_size if self.group_size != -1 else input_size
94 + row_parallel = (input_size != input_size_per_partition)
95 + partition_scales = not marlin_repeat_scales_on_all_ranks(
96 + self.has_g_idx, self.group_size, row_parallel)
97 +
98 + scales_and_zp_size = input_size // group_size
99 +
100 + if partition_scales:
101 + assert input_size_per_partition % group_size == 0
102 + scales_and_zp_size = input_size_per_partition // group_size
103 +
104 + weight = PackedvLLMParameter(input_dim=1,
105 + output_dim=0,
106 + weight_loader=weight_loader,
107 + packed_factor=self.pack_factor,
108 + packed_dim=1,
109 + data=torch.empty(
110 + output_size_per_partition,
111 + input_size_per_partition //
112 + self.pack_factor,
113 + dtype=torch.int32,
114 + ))
115 +
116 + # TODO(czhu): allocate the packed fp8 scales memory here?
117 + # the scales will be expanded by 8x via `cutlass_pack_scale_fp8`
118 + weight_scale_args = {
119 + "weight_loader":
120 + weight_loader,
121 + "data":
122 + torch.empty(
123 + output_size_per_partition,
124 + scales_and_zp_size,
125 + dtype=params_dtype,
126 + )
127 + }
128 +
129 + if not partition_scales:
130 + weight_scale = ChannelQuantScaleParameter(output_dim=0,
131 + **weight_scale_args)
132 + else:
133 + weight_scale = GroupQuantScaleParameter(output_dim=0,
134 + input_dim=1,
135 + **weight_scale_args)
136 +
137 + # A 2D array defining the original shape of the weights
138 + # before packing
139 + weight_shape = BasevLLMParameter(data=torch.empty(2,
140 + dtype=torch.int64),
141 + weight_loader=weight_loader)
142 +
143 + layer.register_parameter("weight_packed", weight)
144 + layer.register_parameter("weight_scale", weight_scale)
145 + layer.register_parameter("weight_shape", weight_shape)
146 +
147 + self.kernel = kernel_type(mp_linear_kernel_config,
148 + w_q_param_name="weight_packed",
149 + w_s_param_name="weight_scale",
150 + w_zp_param_name="weight_zero_point",
151 + w_gidx_param_name="weight_g_idx")
152 +
153 + # Checkpoints are serialized in compressed-tensors format, which is
154 + # different from the format the kernel may want. Handle repacking here.
155 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
156 + self.kernel.process_weights_after_loading(layer)
157 +
158 + def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
159 + bias: Optional[torch.Tensor]) -> torch.Tensor:
160 + return self.kernel.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py
+3
@@ -10,6 +10,8 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision.bitblas imp
10 BitBLASLinearKernel)
11 from vllm.model_executor.layers.quantization.kernels.mixed_precision.conch import ( # noqa: E501
12 ConchLinearKernel)
13 +from vllm.model_executor.layers.quantization.kernels.mixed_precision.cutlass import ( # noqa: E501
14 + CutlassW4A8LinearKernel)
15 from vllm.model_executor.layers.quantization.kernels.mixed_precision.dynamic_4bit import ( # noqa: E501
16 Dynamic4bitLinearKernel)
17 from vllm.model_executor.layers.quantization.kernels.mixed_precision.exllama import ( # noqa: E501
@@ -24,6 +26,7 @@ from vllm.platforms import current_platform
26
27 # in priority/performance order (when available)
28 _POSSIBLE_KERNELS: list[type[MPLinearKernel]] = [
29 + CutlassW4A8LinearKernel,
30 MacheteLinearKernel,
31 AllSparkLinearKernel,
32 MarlinLinearKernel,
vllm/model_executor/layers/quantization/kernels/mixed_precision/cutlass.py new
+114
@@ -0,0 +1,114 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Optional
5 +
6 +import torch
7 +
8 +from vllm import _custom_ops as ops
9 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
10 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 + GroupShape)
12 +from vllm.model_executor.parameter import (BasevLLMParameter,
13 + permute_param_layout_)
14 +from vllm.platforms import current_platform
15 +from vllm.scalar_type import scalar_types
16 +
17 +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
18 +
19 +
20 +class CutlassW4A8LinearKernel(MPLinearKernel):
21 +
22 + def __init__(self, *args, **kwargs):
23 + super().__init__(*args, **kwargs)
24 + # dynamic per-tok fp8 activation quantization
25 + self.quant_fp8 = QuantFP8(static=False,
26 + group_shape=GroupShape.PER_TOKEN)
27 +
28 + @classmethod
29 + def get_min_capability(cls) -> int:
30 + return 90
31 +
32 + @classmethod
33 + def can_implement(cls,
34 + c: MPLinearLayerConfig) -> tuple[bool, Optional[str]]:
35 + if not current_platform.is_cuda():
36 + return False, "CUTLASS only supported on CUDA"
37 +
38 + if not current_platform.is_device_capability(90):
39 + return False, "CUTLASS W4A8 requires compute capability of 90 "\
40 + "(Hopper)"
41 +
42 + if c.act_type != torch.float8_e4m3fn:
43 + return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations"
44 +
45 + if c.has_g_idx:
46 + return False, "Act reordering not supported by CUTLASS W4A8"
47 +
48 + if c.zero_points:
49 + return False, "Zero points not supported by CUTLASS W4A8"
50 +
51 + if c.weight_type != scalar_types.int4:
52 + return False, f"Quant type ({c.weight_type}) not supported by "\
53 + "CUTLASS W4A8, only supported int4"
54 +
55 + # TODO(czhu): support -1 (column-wise)
56 + if c.group_size != 128:
57 + return False, "Only group_size 128 is supported"
58 +
59 + in_features, out_features = c.partition_weight_shape
60 + if in_features % 128 or out_features % 128:
61 + return False, "K and N must be divisible by 128, got "\
62 + f"{c.partition_weight_shape}"
63 + return True, None
64 +
65 + # note assumes that
66 + # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
67 + # `weight_scale` is: {input_dim = 0, output_dim = 1}
68 + def process_weights_after_loading(self, layer: torch.nn.Module):
69 + c = self.config
70 +
71 + # TODO(czhu): optimize speed/mem usage
72 + def transform_w_q(x):
73 + assert isinstance(x, BasevLLMParameter)
74 + permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
75 + x.data = ops.cutlass_encode_and_reorder_int4b(
76 + x.data.t().contiguous().t())
77 + return x
78 +
79 + def transform_w_s(x):
80 + assert isinstance(x, BasevLLMParameter)
81 + permute_param_layout_(x, input_dim=0, output_dim=1)
82 + x.data = x.data.contiguous().to(torch.float8_e4m3fn)
83 + x.data = ops.cutlass_pack_scale_fp8(x.data)
84 + return x
85 +
86 + # Encode/reorder weights and pack scales
87 + self._transform_param(layer, self.w_q_name, transform_w_q)
88 + self._transform_param(layer, self.w_s_name, transform_w_s)
89 +
90 + # TODO(czhu): support loading channel scales
91 + self.w_ch_s = torch.ones((c.partition_weight_shape[1], ),
92 + dtype=torch.float32,
93 + device='cuda')
94 +
95 + def apply_weights(self,
96 + layer: torch.nn.Module,
97 + x: torch.Tensor,
98 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:
99 + assert bias is None, "bias not supported by CUTLASS W4A8"
100 + c = self.config
101 + w_q, w_s, _, _ = self._get_weight_params(layer)
102 +
103 + x_2d = x.reshape(-1, x.shape[-1])
104 + out_shape = x.shape[:-1] + (c.partition_weight_shape[1], )
105 +
106 + x_2d, act_scales = self.quant_fp8(x_2d)
107 + output = ops.cutlass_w4a8_mm(a=x_2d,
108 + b_q=w_q,
109 + b_group_scales=w_s,
110 + b_group_size=c.group_size,
111 + a_token_scales=act_scales,
112 + b_channel_scales=self.w_ch_s)
113 +
114 + return output.reshape(out_shape)