@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 180 lines 5.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for Cutlass W4A16 (Machete) kernel on Hopper.
4
5 Verifies that W4A16 quantized models loaded through vllm select the
6 MacheteLinearKernel on sm_90 GPUs, that weights are correctly repacked,
7 and that inference produces valid output.
8
9 Run `pytest tests/quantization/test_cutlass_w4a16.py`.
10 """
11
12 import pytest
13 import torch
14
15 from vllm.platforms import current_platform
16
17 if not current_platform.has_device_capability(90) or current_platform.is_rocm():
18 pytest.skip(
19 "Machete W4A16 requires Hopper (sm_90).",
20 allow_module_level=True,
21 )
22
23 from vllm.model_executor.kernels.linear import (
24 MPLinearLayerConfig,
25 choose_mp_linear_kernel,
26 )
27 from vllm.model_executor.kernels.linear.mixed_precision import (
28 MacheteLinearKernel,
29 )
30 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
31 CompressedTensorsLinearMethod,
32 CompressedTensorsWNA16,
33 )
34 from vllm.scalar_type import scalar_types
35
36
37 @pytest.fixture(scope="function", autouse=True)
38 def enable_pickle(monkeypatch):
39 """`LLM.apply_model` requires pickling a function."""
40 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
41
42
43 @pytest.mark.parametrize(
44 "act_type,weight_type,group_size,zero_points",
45 [
46 (torch.float16, scalar_types.uint4b8, 128, False),
47 (torch.bfloat16, scalar_types.uint4b8, 128, False),
48 (torch.float16, scalar_types.uint4, 128, True),
49 (torch.float16, scalar_types.uint4b8, -1, False),
50 ],
51 ids=[
52 "fp16-gptq-g128",
53 "bf16-gptq-g128",
54 "fp16-awq-g128",
55 "fp16-channelwise",
56 ],
57 )
58 def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points):
59 """Verify choose_mp_linear_kernel picks MacheteLinearKernel."""
60 config = MPLinearLayerConfig(
61 full_weight_shape=(4096, 4096),
62 partition_weight_shape=(4096, 4096),
63 act_type=act_type,
64 weight_type=weight_type,
65 group_size=group_size,
66 zero_points=zero_points,
67 )
68 kernel = choose_mp_linear_kernel(config)
69 assert kernel is MacheteLinearKernel, (
70 f"Expected MacheteLinearKernel, got {kernel.__name__}"
71 )
72
73
74 @pytest.mark.parametrize(
75 "full_shape,part_shape,weight_type,group_size,expected_reason",
76 [
77 (
78 (4096, 4096),
79 (4096, 4096),
80 scalar_types.float6_e3m2f,
81 128,
82 "Quant type",
83 ),
84 ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, "Group size"),
85 ],
86 ids=["unsupported-quant-type", "unsupported-group-size"],
87 )
88 def test_machete_rejects_invalid_config(
89 full_shape, part_shape, weight_type, group_size, expected_reason
90 ):
91 """Verify Machete rejects unsupported configurations."""
92 config = MPLinearLayerConfig(
93 full_weight_shape=full_shape,
94 partition_weight_shape=part_shape,
95 act_type=torch.float16,
96 weight_type=weight_type,
97 group_size=group_size,
98 zero_points=False,
99 )
100 can_impl, reason = MacheteLinearKernel.can_implement(config)
101 assert not can_impl
102 assert reason is not None
103 assert expected_reason in reason
104
105
106 def test_kernel_selection_with_disabled_machete(monkeypatch):
107 """Verify kernel selection falls back when Machete is disabled."""
108 monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
109
110 config = MPLinearLayerConfig(
111 full_weight_shape=(4096, 4096),
112 partition_weight_shape=(4096, 4096),
113 act_type=torch.float16,
114 weight_type=scalar_types.uint4b8,
115 group_size=128,
116 zero_points=False,
117 )
118 kernel = choose_mp_linear_kernel(config)
119 assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
120
121
122 @pytest.mark.parametrize(
123 "model_name",
124 [
125 "nm-testing/tinyllama-oneshot-w4a16-channel-v2",
126 ],
127 )
128 def test_w4a16_machete_e2e(vllm_runner, model_name):
129 """Load a W4A16 model, verify Machete kernel is used, and generate."""
130 with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
131
132 def check_model(model):
133 layer = model.model.layers[0]
134 qkv_proj = layer.self_attn.qkv_proj
135
136 assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
137 assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
138 assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
139 f"Expected MacheteLinearKernel on Hopper, "
140 f"got {type(qkv_proj.scheme.kernel).__name__}"
141 )
142
143 assert hasattr(qkv_proj, "weight_packed")
144 assert hasattr(qkv_proj, "weight_scale")
145 assert qkv_proj.weight_packed.dtype == torch.int32
146
147 llm.apply_model(check_model)
148
149 output = llm.generate_greedy("Hello my name is", max_tokens=10)
150 assert output
151 assert len(output[0][1]) > 0
152
153
154 def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
155 """Verify Machete works with bf16 activations and is deterministic."""
156 model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
157 prompt = "The capital of France is"
158
159 with vllm_runner(
160 model_name,
161 enforce_eager=True,
162 dtype="bfloat16",
163 gpu_memory_utilization=0.5,
164 ) as llm:
165
166 def check_kernel_type(model):
167 layer = model.model.layers[0]
168 scheme = layer.self_attn.qkv_proj.scheme
169 assert isinstance(scheme.kernel, MacheteLinearKernel), (
170 f"Expected MacheteLinearKernel with bf16, "
171 f"got {type(scheme.kernel).__name__}"
172 )
173
174 llm.apply_model(check_kernel_type)
175
176 out1 = llm.generate_greedy(prompt, max_tokens=10)
177 out2 = llm.generate_greedy(prompt, max_tokens=10)
178 assert out1[0][1] == out2[0][1], (
179 f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
180 )