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 4, 15:00 UTC
Cutlass W4A16 (Machete) Tests (#35450)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com>
ojhaanshika committed
Apr 26, 2026 at 22:15 UTC
592ae6805cb87c9a44c29bd9c6eef9d04d91e39b
2 files changed
+187
.buildkite/test_areas/kernels.yaml
+2
@@ -95,11 +95,13 @@ steps:
95
- tests/kernels/moe/test_deepgemm.py
96
- tests/kernels/moe/test_batched_deepgemm.py
97
- tests/kernels/attention/test_deepgemm_attention.py
98
+ - tests/quantization/test_cutlass_w4a16.py
99
commands:
100
- pytest -v -s kernels/quantization/test_block_fp8.py
101
- pytest -v -s kernels/moe/test_deepgemm.py
102
- pytest -v -s kernels/moe/test_batched_deepgemm.py
103
- pytest -v -s kernels/attention/test_deepgemm_attention.py
104
+ - pytest -v -s quantization/test_cutlass_w4a16.py
105
106
- label: Kernels (B200)
107
timeout_in_minutes: 30
tests/quantization/test_cutlass_w4a16.py
new
+185
@@ -0,0 +1,185 @@
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):
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
+ has_g_idx=False,
68
+ )
69
+ kernel = choose_mp_linear_kernel(config)
70
+ assert kernel is MacheteLinearKernel, (
71
+ f"Expected MacheteLinearKernel, got {kernel.__name__}"
72
+ )
73
+
74
+
75
+@pytest.mark.parametrize(
76
+ "full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
77
+ [
78
+ ((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
79
+ (
80
+ (4096, 4096),
81
+ (4096, 4096),
82
+ scalar_types.float6_e3m2f,
83
+ 128,
84
+ False,
85
+ "Quant type",
86
+ ),
87
+ ((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
88
+ ],
89
+ ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
90
+)
91
+def test_machete_rejects_invalid_config(
92
+ full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
93
+):
94
+ """Verify Machete rejects unsupported configurations."""
95
+ config = MPLinearLayerConfig(
96
+ full_weight_shape=full_shape,
97
+ partition_weight_shape=part_shape,
98
+ act_type=torch.float16,
99
+ weight_type=weight_type,
100
+ group_size=group_size,
101
+ zero_points=False,
102
+ has_g_idx=has_g_idx,
103
+ )
104
+ can_impl, reason = MacheteLinearKernel.can_implement(config)
105
+ assert not can_impl
106
+ assert expected_reason in reason
107
+
108
+
109
+def test_kernel_selection_with_disabled_machete(monkeypatch):
110
+ """Verify kernel selection falls back when Machete is disabled."""
111
+ monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
112
+
113
+ config = MPLinearLayerConfig(
114
+ full_weight_shape=(4096, 4096),
115
+ partition_weight_shape=(4096, 4096),
116
+ act_type=torch.float16,
117
+ weight_type=scalar_types.uint4b8,
118
+ group_size=128,
119
+ zero_points=False,
120
+ has_g_idx=False,
121
+ )
122
+ kernel = choose_mp_linear_kernel(config)
123
+ assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
124
+
125
+
126
+@pytest.mark.parametrize(
127
+ "model_name",
128
+ [
129
+ "nm-testing/tinyllama-oneshot-w4a16-channel-v2",
130
+ "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
131
+ ],
132
+)
133
+def test_w4a16_machete_e2e(vllm_runner, model_name):
134
+ """Load a W4A16 model, verify Machete kernel is used, and generate."""
135
+ with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
136
+
137
+ def check_model(model):
138
+ layer = model.model.layers[0]
139
+ qkv_proj = layer.self_attn.qkv_proj
140
+
141
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
142
+ assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
143
+ assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
144
+ f"Expected MacheteLinearKernel on Hopper, "
145
+ f"got {type(qkv_proj.scheme.kernel).__name__}"
146
+ )
147
+
148
+ assert hasattr(qkv_proj, "weight_packed")
149
+ assert hasattr(qkv_proj, "weight_scale")
150
+ assert qkv_proj.weight_packed.dtype == torch.int32
151
+
152
+ llm.apply_model(check_model)
153
+
154
+ output = llm.generate_greedy("Hello my name is", max_tokens=10)
155
+ assert output
156
+ assert len(output[0][1]) > 0
157
+
158
+
159
+def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
160
+ """Verify Machete works with bf16 activations and is deterministic."""
161
+ model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
162
+ prompt = "The capital of France is"
163
+
164
+ with vllm_runner(
165
+ model_name,
166
+ enforce_eager=True,
167
+ dtype="bfloat16",
168
+ gpu_memory_utilization=0.5,
169
+ ) as llm:
170
+
171
+ def check_kernel_type(model):
172
+ layer = model.model.layers[0]
173
+ scheme = layer.self_attn.qkv_proj.scheme
174
+ assert isinstance(scheme.kernel, MacheteLinearKernel), (
175
+ f"Expected MacheteLinearKernel with bf16, "
176
+ f"got {type(scheme.kernel).__name__}"
177
+ )
178
+
179
+ llm.apply_model(check_kernel_type)
180
+
181
+ out1 = llm.generate_greedy(prompt, max_tokens=10)
182
+ out2 = llm.generate_greedy(prompt, max_tokens=10)
183
+ assert out1[0][1] == out2[0][1], (
184
+ f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
185
+ )