| 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/quantization/test_cutlass_w4a8.py`. |
| 6 | """ |
| 7 | |
| 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 | 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 |
| 23 | |
| 24 | if not current_platform.is_cuda(): |
| 25 | pytest.skip("These tests use CUTLASS which requires CUDA", allow_module_level=True) |
| 26 | |
| 27 | # TODO: in future PR refactor this and `is_quant_method_supported` in the kernel |
| 28 | # unit tests to a common utility function. Currently the use of |
| 29 | # `is_quant_method_supported` conflates kernels with quantization methods |
| 30 | # an assumption which is breaking down as quantizations methods can have |
| 31 | # have kernels and some kernels support multiple quantization methods. |
| 32 | capability = current_platform.get_device_capability() |
| 33 | IS_SUPPORTED_BY_GPU = capability is not None and capability[0] >= 9 |
| 34 | |
| 35 | MNK_SHAPES = [ |
| 36 | (1, 128, 128), |
| 37 | (1, 512, 1024), |
| 38 | (1, 4096, 4096), |
| 39 | (1, 8192, 28672), |
| 40 | (13, 8192, 4096), |
| 41 | (26, 4096, 8192), |
| 42 | (64, 4096, 4096), |
| 43 | (64, 8192, 28672), |
| 44 | (257, 128, 4096), |
| 45 | (257, 4096, 4096), |
| 46 | (1024, 4096, 8192), |
| 47 | (1024, 8192, 4096), |
| 48 | ] |
| 49 | |
| 50 | # TODO(czhu): get supported schedules from fn |
| 51 | SCHEDULES = [ |
| 52 | "128x16_1x1x1", |
| 53 | "256x16_1x1x1", |
| 54 | "128x32_1x1x1", |
| 55 | "256x32_1x1x1", |
| 56 | "128x64_1x1x1", |
| 57 | "256x64_1x1x1", |
| 58 | "128x128_1x1x1", |
| 59 | "256x128_1x1x1", |
| 60 | "128x256_1x1x1", |
| 61 | "128x256_2x1x1", |
| 62 | ] |
| 63 | |
| 64 | |
| 65 | @dataclass |
| 66 | class TypeConfig: |
| 67 | act_type: torch.dtype |
| 68 | weight_type: ScalarType |
| 69 | output_type: torch.dtype | None |
| 70 | group_scale_type: torch.dtype | None |
| 71 | channel_scale_type: torch.dtype | None |
| 72 | token_scale_type: torch.dtype | None |
| 73 | |
| 74 | |
| 75 | @dataclass |
| 76 | class Tensors: |
| 77 | w_ref: torch.Tensor |
| 78 | a_ref: torch.Tensor |
| 79 | a: torch.Tensor |
| 80 | w_q: torch.Tensor |
| 81 | w_g_s: torch.Tensor |
| 82 | w_ch_s: torch.Tensor |
| 83 | w_tok_s: torch.Tensor |
| 84 | |
| 85 | |
| 86 | # (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints, |
| 87 | # Ch Scales Type, Tok Scales Type) |
| 88 | TestTypeTuple = tuple[ |
| 89 | list[torch.dtype], ScalarType, torch.dtype | None, torch.dtype | None, bool |
| 90 | ] |
| 91 | TEST_TYPES = [ |
| 92 | *( |
| 93 | TypeConfig( |
| 94 | act_type=torch.float8_e4m3fn, |
| 95 | weight_type=w_type, |
| 96 | output_type=o_type, |
| 97 | group_scale_type=torch.float8_e4m3fn, |
| 98 | channel_scale_type=torch.float32, |
| 99 | token_scale_type=torch.float32, |
| 100 | ) |
| 101 | for w_type in [scalar_types.int4] |
| 102 | # TODO(czhu): fp16 out type |
| 103 | for o_type in [torch.bfloat16] |
| 104 | ), |
| 105 | ] |
| 106 | |
| 107 | # TODO: in future PR refactor this and `is_quant_method_supported` in the kernel |
| 108 | # unit tests to a common utility function. Currently the use of |
| 109 | # `is_quant_method_supported` conflates kernels with quantization methods |
| 110 | # an assumption which is breaking down as quantizations methods can have |
| 111 | # have kernels and some kernels support multiple quantization methods. |
| 112 | IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90) |
| 113 | |
| 114 | |
| 115 | # For testing quantized linear kernels |
| 116 | def to_fp8(tensor: torch.Tensor): |
| 117 | finfo = torch.finfo(torch.float8_e4m3fn) |
| 118 | return tensor.clamp(min=finfo.min, max=finfo.max).to(dtype=torch.float8_e4m3fn) |
| 119 | |
| 120 | |
| 121 | def cutlass_quantize_and_pack( |
| 122 | atype: torch.dtype, |
| 123 | w: torch.Tensor, |
| 124 | wtype: ScalarType, |
| 125 | stype: torch.dtype | None, |
| 126 | group_size: int | None, |
| 127 | zero_points: bool = False, |
| 128 | ): |
| 129 | assert wtype.is_integer(), "TODO: support floating point weights" |
| 130 | |
| 131 | w_ref, w_q, w_s, w_zp = quantize_weights( |
| 132 | w, wtype, group_size=group_size, zero_points=zero_points |
| 133 | ) |
| 134 | |
| 135 | # since scales are cast to fp8, we need to compute w_ref this way |
| 136 | w_ref = ( |
| 137 | (w_q).to(torch.float32) |
| 138 | * w_s.to(atype).to(torch.float32).repeat_interleave(group_size, dim=0) |
| 139 | ).to(atype) |
| 140 | |
| 141 | # bit mask prevents sign extending int4 when packing |
| 142 | w_q = pack_rows(w_q & 0x0F, wtype.size_bits, *w_q.shape) |
| 143 | w_q = w_q.t().contiguous().t() # convert to col major |
| 144 | |
| 145 | w_q_packed = ops.cutlass_encode_and_reorder_int4b(w_q) |
| 146 | w_s_packed = ops.cutlass_pack_scale_fp8(w_s.to(atype)) |
| 147 | |
| 148 | return w_ref, w_q_packed, w_s_packed, w_zp |
| 149 | |
| 150 | |
| 151 | def create_test_tensors( |
| 152 | shape: tuple[int, int, int], types: TypeConfig, group_size: int | None |
| 153 | ) -> Tensors: |
| 154 | m, n, k = shape |
| 155 | |
| 156 | print( |
| 157 | "create_test_tensors, shape:", shape, "types:", types, "group_size:", group_size |
| 158 | ) |
| 159 | |
| 160 | a = to_fp8(torch.randn((m, k), device="cuda")) |
| 161 | w = to_fp8(torch.randn((k, n), device="cuda")) |
| 162 | |
| 163 | if types.group_scale_type is not None: |
| 164 | w = w.to(types.group_scale_type) |
| 165 | if w.dtype.itemsize == 1: |
| 166 | w = w.to(torch.float16) |
| 167 | |
| 168 | w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack( |
| 169 | a.dtype, w, types.weight_type, types.group_scale_type, group_size, False |
| 170 | ) |
| 171 | |
| 172 | a_ref = a.to(torch.float32) |
| 173 | w_ref = w_ref.to(torch.float32) |
| 174 | |
| 175 | # for the practical use case we need per-tok scales for fp8 activations |
| 176 | w_tok_s = torch.randn((m,), device="cuda", dtype=types.token_scale_type) |
| 177 | w_ch_s = torch.randn((n,), device="cuda", dtype=types.channel_scale_type) |
| 178 | |
| 179 | return Tensors( |
| 180 | w_ref=w_ref, |
| 181 | a_ref=a_ref, |
| 182 | a=a, |
| 183 | w_q=w_q_packed, |
| 184 | w_g_s=w_s, |
| 185 | w_ch_s=w_ch_s, |
| 186 | w_tok_s=w_tok_s, |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | def mm_test_helper( |
| 191 | types: TypeConfig, |
| 192 | tensors: Tensors, |
| 193 | group_size: int | None = None, |
| 194 | schedule: str | None = None, |
| 195 | ): |
| 196 | # CUTLASS upstream uses fp8 with fastaccum as reference |
| 197 | # https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu#L406 |
| 198 | output_ref = torch._scaled_mm( |
| 199 | tensors.a_ref.to(types.act_type), |
| 200 | tensors.w_ref.to(types.act_type).t().contiguous().t(), # col major |
| 201 | tensors.w_tok_s.unsqueeze(1), |
| 202 | tensors.w_ch_s.unsqueeze(0), |
| 203 | out_dtype=types.output_type, |
| 204 | use_fast_accum=True, |
| 205 | ) |
| 206 | |
| 207 | output = ops.cutlass_w4a8_mm( |
| 208 | a=tensors.a, |
| 209 | b_q=tensors.w_q, |
| 210 | b_group_scales=tensors.w_g_s, |
| 211 | b_group_size=group_size, |
| 212 | b_channel_scales=tensors.w_ch_s, |
| 213 | a_token_scales=tensors.w_tok_s, |
| 214 | ) |
| 215 | |
| 216 | print(output) |
| 217 | print(output_ref) |
| 218 | |
| 219 | torch.testing.assert_close( |
| 220 | output, output_ref.to(output.dtype), rtol=1e-2, atol=1e-2 |
| 221 | ) |
| 222 | |
| 223 | |
| 224 | @pytest.mark.skipif( |
| 225 | not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type." |
| 226 | ) |
| 227 | @pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x)) |
| 228 | @pytest.mark.parametrize("types", TEST_TYPES) |
| 229 | @pytest.mark.parametrize("schedule", SCHEDULES) |
| 230 | def test_cutlass_w4a8(shape, types: TypeConfig, schedule): |
| 231 | group_sizes = [128] |
| 232 | for group_size in group_sizes: |
| 233 | tensors = create_test_tensors(shape, types, group_size) |
| 234 | mm_test_helper(types, tensors, group_size, schedule) |
| 235 | |
| 236 | |
| 237 | # Test to make sure cuda graphs work |
| 238 | class W4A8Layer(torch.nn.Module): |
| 239 | def __init__(self, **kwargs): |
| 240 | super().__init__() |
| 241 | self.kwargs = kwargs |
| 242 | |
| 243 | def forward(self, a): |
| 244 | return ops.cutlass_w4a8_mm(a=a, **self.kwargs) |
| 245 | |
| 246 | |
| 247 | @pytest.mark.skipif( |
| 248 | not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type." |
| 249 | ) |
| 250 | def test_w4a8_cuda_graph(): |
| 251 | m, n, k = 512, 4096, 4096 |
| 252 | |
| 253 | a = to_fp8(torch.randn((m, k), device="cuda")) |
| 254 | b = to_fp8(torch.randn((k, n), device="cuda")) |
| 255 | |
| 256 | wtype = scalar_types.int4 |
| 257 | stype = torch.float8_e4m3fn |
| 258 | group_size = 128 |
| 259 | zero_points = False |
| 260 | |
| 261 | w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack( |
| 262 | a.dtype, b.to(torch.float16), wtype, stype, group_size, zero_points |
| 263 | ) |
| 264 | |
| 265 | w_tok_s = torch.randn((m,), device="cuda", dtype=torch.float32) |
| 266 | w_ch_s = torch.randn((n,), device="cuda", dtype=torch.float32) |
| 267 | |
| 268 | # Construct a trivial model with a single layer that calls the kernel |
| 269 | model = W4A8Layer( |
| 270 | b_q=w_q_packed, |
| 271 | b_group_scales=w_s, |
| 272 | b_group_size=group_size, |
| 273 | b_channel_scales=w_ch_s, |
| 274 | a_token_scales=w_tok_s, |
| 275 | ) |
| 276 | |
| 277 | output_ref = torch._scaled_mm( |
| 278 | a, |
| 279 | w_ref.to(a.dtype).t().contiguous().t(), # col major |
| 280 | w_tok_s.unsqueeze(1), |
| 281 | w_ch_s.unsqueeze(0), |
| 282 | out_dtype=torch.bfloat16, |
| 283 | use_fast_accum=True, |
| 284 | ) |
| 285 | |
| 286 | # Run the model with a cuda graph |
| 287 | stream = torch.cuda.Stream() |
| 288 | with torch.cuda.stream(stream): |
| 289 | g = torch.cuda.CUDAGraph() |
| 290 | with torch.cuda.graph(g): |
| 291 | output = model(a) |
| 292 | |
| 293 | output.zero_() |
| 294 | g.replay() |
| 295 | |
| 296 | torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=1e-2) |
| 297 | |
| 298 | |
| 299 | @pytest.mark.skipif( |
| 300 | not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type." |
| 301 | ) |
| 302 | @pytest.mark.parametrize("shape", MNK_SHAPES) |
| 303 | def test_convert_packed_uint4b8_to_signed_int4_inplace(shape): |
| 304 | """The W4A16 checkpoints encode the weights as int4b8 packed to int32. |
| 305 | The CUTLASS kernels expect signed int4 packed to int32. |
| 306 | This tests checks that the runtime int4b8 -> signed int4 conversion |
| 307 | matches the offline conversion step exactly. |
| 308 | """ |
| 309 | _, N, K = shape |
| 310 | # random weights packed to int32 |
| 311 | t = torch.randint( |
| 312 | low=torch.iinfo(torch.int32).min, |
| 313 | high=torch.iinfo(torch.int32).max + 1, |
| 314 | size=(N, K // 8), |
| 315 | dtype=torch.int32, |
| 316 | device="cuda", |
| 317 | ) |
| 318 | |
| 319 | # compute reference |
| 320 | unpacked = unpack_quantized_values_into_int32( |
| 321 | t.clone(), scalar_types.uint4b8, packed_dim=1 |
| 322 | ) |
| 323 | unpacked = unpacked - 8 # int4b8 -> signed int4 |
| 324 | ref = pack_cols(unpacked & 0x0F, 4, *unpacked.shape) |
| 325 | |
| 326 | out = convert_packed_uint4b8_to_signed_int4_inplace(t.clone()) |
| 327 | |
| 328 | assert torch.equal(ref, out) |
| 329 | assert not torch.equal(ref, t) |