@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d468322dc1
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 5, 07:15 UTC

[Kernel][Hardware][AMD] Add TritonW4A16LinearKernel for ROCm (#37352)

Signed-off-by: jatseng-ai <jatseng@amd.com> Signed-off-by: jatseng-ai <janet.tseng@amd.com> Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Matthias Gehre <matthias.gehre@amd.com>

jatseng-ai committed Apr 10, 2026 at 03:25 UTC d468322dc1c89377f65488973d2f1927e9bbe5cb
6 files changed +831
tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py new
+35
@@ -0,0 +1,35 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""End-to-end smoke test for CT W4A16 models on ROCm.
5 +
6 +This validates that a real compressed-tensors W4A16 model can run inference
7 +end-to-end (which will exercise the Triton W4A16 kernel when selected).
8 +
9 +Run `pytest tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py`.
10 +"""
11 +
12 +import pytest
13 +
14 +from vllm.platforms import current_platform
15 +
16 +
17 +@pytest.mark.parametrize(
18 + "model_path",
19 + [
20 + # Listed in tests/weight_loading/models.txt
21 + "nm-testing/tinyllama-oneshot-w4a16-group128-v2",
22 + ],
23 +)
24 +@pytest.mark.parametrize("max_tokens", [32])
25 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm")
26 +def test_rocm_compressed_tensors_w4a16_e2e(
27 + vllm_runner, example_prompts, model_path, max_tokens
28 +):
29 + # Use fp16 activations for maximum compatibility.
30 + # gpu_memory_utilization lowered to work on shared nodes.
31 + with vllm_runner(
32 + model_path, dtype="float16", gpu_memory_utilization=0.3
33 + ) as vllm_model:
34 + # If the W4A16 kernel is broken, this will typically throw.
35 + vllm_model.generate_greedy(example_prompts, max_tokens=max_tokens)
tests/kernels/quantization/test_triton_w4a16.py new
+304
@@ -0,0 +1,304 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Tests for the ROCm Triton W4A16 GEMM kernel.
5 +
6 +Run `pytest tests/kernels/quantization/test_triton_w4a16.py`.
7 +"""
8 +
9 +import importlib
10 +
11 +import pytest
12 +import torch
13 +
14 +from vllm.platforms import current_platform
15 +from vllm.utils.torch_utils import set_random_seed
16 +
17 +# This test module is ROCm/Triton specific. Avoid import-time failures on
18 +# non-ROCm or environments without Triton by skipping early.
19 +if not current_platform.is_rocm():
20 + pytest.skip("ROCm only", allow_module_level=True)
21 +
22 +pytest.importorskip("triton")
23 +
24 +device = "cuda"
25 +
26 +triton_w4a16_module = importlib.import_module(
27 + "vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16"
28 +)
29 +triton_w4a16_gemm = triton_w4a16_module.triton_w4a16_gemm
30 +TritonW4A16LinearKernel = triton_w4a16_module.TritonW4A16LinearKernel
31 +
32 +
33 +def _pack_int4_along_n(w_int4_kn: torch.Tensor) -> torch.Tensor:
34 + """Pack int4 values along N: [K, N] -> [K, N//8] int32."""
35 + assert w_int4_kn.dtype == torch.int32
36 + K, N = w_int4_kn.shape
37 + assert N % 8 == 0
38 + shifts = torch.arange(8, device=w_int4_kn.device, dtype=torch.int32) * 4
39 + return torch.sum(
40 + (w_int4_kn.view(K, N // 8, 8) & 0xF) << shifts,
41 + dim=2,
42 + dtype=torch.int32,
43 + ).contiguous()
44 +
45 +
46 +def _unpack_int4_along_n(w_packed_kn8: torch.Tensor) -> torch.Tensor:
47 + """Unpack int4 values along N: [K, N//8] -> [K, N] int32."""
48 + assert w_packed_kn8.dtype == torch.int32
49 + K, N8 = w_packed_kn8.shape
50 + shifts = torch.arange(8, device=w_packed_kn8.device, dtype=torch.int32) * 4
51 + nibbles = (w_packed_kn8.unsqueeze(-1) >> shifts) & 0xF
52 + return nibbles.reshape(K, N8 * 8)
53 +
54 +
55 +def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor:
56 + """Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8]."""
57 + assert w_int4_kn.dtype == torch.int32
58 + K, N = w_int4_kn.shape
59 + assert K % 8 == 0
60 + out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device)
61 + for i in range(8):
62 + out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4)
63 + return out.contiguous()
64 +
65 +
66 +def _w4a16_reference(
67 + a_mk: torch.Tensor,
68 + b_packed_kn8: torch.Tensor,
69 + scales_gn: torch.Tensor,
70 + *,
71 + group_size: int,
72 + qzeros_gn8: torch.Tensor | None,
73 + zp_bias: int,
74 +) -> torch.Tensor:
75 + """Reference implementation for W4A16.
76 +
77 + a_mk: [M,K] fp16/bf16
78 + b_packed_kn8: [K, N//8] int32, N-packed int4 weights
79 + scales_gn: [K//G, N] fp16/bf16
80 + qzeros_gn8: [K//G, N//8] int32, N-packed int4 zeros, or None
81 + """
82 + assert a_mk.dtype in (torch.float16, torch.bfloat16)
83 + assert b_packed_kn8.dtype == torch.int32
84 + assert scales_gn.dtype == a_mk.dtype
85 +
86 + M, K = a_mk.shape
87 + N = b_packed_kn8.shape[1] * 8
88 + assert b_packed_kn8.shape[0] == K
89 +
90 + assert group_size > 0 and K % group_size == 0
91 + G = group_size
92 + num_groups = K // G
93 + assert scales_gn.shape == (num_groups, N)
94 +
95 + w_int4 = _unpack_int4_along_n(b_packed_kn8) # [K,N]
96 + if qzeros_gn8 is None:
97 + z_full = torch.full((K, N), zp_bias, dtype=torch.int32, device=a_mk.device)
98 + else:
99 + assert qzeros_gn8.shape == (num_groups, N // 8)
100 + z_gn = _unpack_int4_along_n(qzeros_gn8) # [G,N] in groups
101 + z_full = z_gn.repeat_interleave(G, dim=0) # [K,N]
102 +
103 + s_full = scales_gn.repeat_interleave(G, dim=0).to(torch.float32) # [K,N]
104 + w_fp = (w_int4 - z_full).to(torch.float32) * s_full # [K,N]
105 +
106 + out = a_mk.to(torch.float32) @ w_fp # [M,N]
107 + return out.to(a_mk.dtype)
108 +
109 +
110 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
111 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
112 +@pytest.mark.parametrize(
113 + "M,K,N,G,has_zp",
114 + [
115 + (1, 256, 256, 32, False),
116 + (17, 256, 512, 32, False),
117 + (32, 512, 256, 64, False),
118 + (33, 512, 512, 128, False),
119 + (64, 1024, 256, 256, False),
120 + (128, 256, 1024, 32, True),
121 + (64, 512, 512, 64, True),
122 + ],
123 +)
124 +def test_triton_w4a16_gemm_matches_reference(dtype, M, K, N, G, has_zp):
125 + if not torch.cuda.is_available():
126 + pytest.skip("CUDA/HIP device not available")
127 + if N % 8 != 0 or K % G != 0:
128 + pytest.skip("Invalid test shape")
129 +
130 + set_random_seed(0)
131 +
132 + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
133 + w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
134 + b_packed = _pack_int4_along_n(w_int4)
135 +
136 + scales = (0.05 * torch.rand((K // G, N), device=device, dtype=torch.float32)).to(
137 + dtype
138 + )
139 +
140 + qzeros = None
141 + if has_zp:
142 + zeros_int4 = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
143 + qzeros = _pack_int4_along_n(zeros_int4)
144 +
145 + out = triton_w4a16_gemm(
146 + a=a,
147 + b_q=b_packed,
148 + scales=scales,
149 + qzeros=qzeros,
150 + group_size=G,
151 + zp_bias=8,
152 + )
153 + ref = _w4a16_reference(
154 + a,
155 + b_packed,
156 + scales,
157 + group_size=G,
158 + qzeros_gn8=qzeros,
159 + zp_bias=8,
160 + )
161 +
162 + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)
163 +
164 +
165 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
166 +def test_triton_w4a16_gemm_requires_contiguous_inputs():
167 + if not torch.cuda.is_available():
168 + pytest.skip("CUDA/HIP device not available")
169 +
170 + set_random_seed(0)
171 + M, K, N, G = 32, 256, 256, 32
172 + a = torch.randn((K, M), device=device, dtype=torch.float16).t() # non-contiguous
173 + w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
174 + b_packed = _pack_int4_along_n(w_int4)
175 + scales = torch.rand((K // G, N), device=device, dtype=torch.float16)
176 +
177 + with pytest.raises(AssertionError):
178 + triton_w4a16_gemm(
179 + a=a,
180 + b_q=b_packed,
181 + scales=scales,
182 + qzeros=None,
183 + group_size=G,
184 + zp_bias=8,
185 + )
186 +
187 +
188 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
189 +def test_triton_w4a16_process_weights_after_loading_repacks_layout():
190 + if not torch.cuda.is_available():
191 + pytest.skip("CUDA/HIP device not available")
192 +
193 + from vllm.config import VllmConfig, set_current_vllm_config
194 + from vllm.distributed import (
195 + ensure_model_parallel_initialized,
196 + init_distributed_environment,
197 + )
198 + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
199 + MPLinearLayerConfig,
200 + )
201 + from vllm.model_executor.parameter import (
202 + GroupQuantScaleParameter,
203 + PackedColumnParameter,
204 + PackedvLLMParameter,
205 + )
206 + from vllm.scalar_type import scalar_types
207 +
208 + with set_current_vllm_config(VllmConfig()):
209 + init_distributed_environment(
210 + world_size=1,
211 + rank=0,
212 + distributed_init_method="tcp://127.0.0.1:0",
213 + local_rank=0,
214 + )
215 + ensure_model_parallel_initialized(1, 1)
216 +
217 + set_random_seed(0)
218 +
219 + # Small-but-nontrivial shapes.
220 + K, N = 256, 256
221 + G = 32
222 + assert K % 8 == 0 and N % 8 == 0 and K % G == 0
223 +
224 + # Build a canonical int4 weight grid then pack into the CT checkpoint layout.
225 + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
226 + w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn) # [N, K//8]
227 +
228 + # Scales in CT checkpoint layout for WNA16: [N, K//G]
229 + scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
230 +
231 + # Asymmetric case: zero points in CT checkpoint layout [N//8, K//G] (N-packed)
232 + zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
233 + zeros_packed_gn8 = _pack_int4_along_n(zeros_int4_gn) # [K//G, N//8]
234 + zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G]
235 +
236 + config = MPLinearLayerConfig(
237 + full_weight_shape=(K, N),
238 + partition_weight_shape=(K, N),
239 + weight_type=scalar_types.uint4, # asymmetric
240 + act_type=torch.float16,
241 + group_size=G,
242 + zero_points=True,
243 + has_g_idx=False,
244 + )
245 + kernel = TritonW4A16LinearKernel(
246 + config,
247 + w_q_param_name="weight_packed",
248 + w_s_param_name="weight_scale",
249 + w_zp_param_name="weight_zero_point",
250 + w_gidx_param_name=None,
251 + )
252 +
253 + # Build dummy layer with vLLM parameter wrappers.
254 + weight_loader = lambda *args, **kwargs: None
255 +
256 + class DummyLayer(torch.nn.Module):
257 + pass
258 +
259 + layer = DummyLayer()
260 + layer.register_parameter(
261 + "weight_packed",
262 + PackedvLLMParameter(
263 + data=w_ckpt_nk8,
264 + weight_loader=weight_loader,
265 + input_dim=1,
266 + output_dim=0,
267 + packed_factor=8,
268 + packed_dim=1,
269 + ),
270 + )
271 + layer.register_parameter(
272 + "weight_scale",
273 + GroupQuantScaleParameter(
274 + data=scales_ckpt_nkg,
275 + weight_loader=weight_loader,
276 + input_dim=1,
277 + output_dim=0,
278 + ),
279 + )
280 + layer.register_parameter(
281 + "weight_zero_point",
282 + PackedColumnParameter(
283 + data=zeros_ckpt_n8kg,
284 + weight_loader=weight_loader,
285 + output_dim=0,
286 + packed_factor=8,
287 + packed_dim=0,
288 + ),
289 + )
290 +
291 + kernel.process_weights_after_loading(layer)
292 +
293 + # Expected transformed layouts.
294 + expected_w_kn8 = _pack_int4_along_n(w_int4_kn) # [K, N//8]
295 + expected_scales_gn = scales_ckpt_nkg.t().contiguous() # [K//G, N]
296 + expected_zeros_gn8 = zeros_ckpt_n8kg.t().contiguous() # [K//G, N//8]
297 +
298 + assert tuple(layer.weight_packed.shape) == (K, N // 8)
299 + assert tuple(layer.weight_scale.shape) == (K // G, N)
300 + assert tuple(layer.weight_zero_point.shape) == (K // G, N // 8)
301 +
302 + torch.testing.assert_close(layer.weight_packed, expected_w_kn8)
303 + torch.testing.assert_close(layer.weight_scale, expected_scales_gn)
304 + torch.testing.assert_close(layer.weight_zero_point, expected_zeros_gn8)
tests/kernels/quantization/test_w4a16_kernel_selection.py new
+53
@@ -0,0 +1,53 @@
1 +#!/usr/bin/env python3
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +"""Tests for W4A16 kernel selection logic (ROCm).
5 +
6 +Run `pytest tests/kernels/quantization/test_w4a16_kernel_selection.py`.
7 +"""
8 +
9 +import pytest
10 +import torch
11 +
12 +from vllm.model_executor.kernels.linear import (
13 + MPLinearLayerConfig,
14 + choose_mp_linear_kernel,
15 +)
16 +from vllm.platforms import current_platform
17 +from vllm.scalar_type import scalar_types
18 +
19 +
20 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
21 +def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4b8():
22 + # int4 weights, 16-bit activations (CT W4A16 typical config).
23 + K, N = 1024, 256
24 + config = MPLinearLayerConfig(
25 + full_weight_shape=(K, N),
26 + partition_weight_shape=(K, N),
27 + weight_type=scalar_types.uint4b8, # symmetric int4 (bias=8)
28 + act_type=torch.float16,
29 + group_size=128,
30 + zero_points=False,
31 + has_g_idx=False,
32 + )
33 +
34 + kernel_type = choose_mp_linear_kernel(config)
35 + assert kernel_type.__name__ == "TritonW4A16LinearKernel"
36 +
37 +
38 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
39 +def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4_asymmetric():
40 + # Asymmetric int4 weights should also be supported (explicit zero points).
41 + K, N = 512, 512
42 + config = MPLinearLayerConfig(
43 + full_weight_shape=(K, N),
44 + partition_weight_shape=(K, N),
45 + weight_type=scalar_types.uint4, # asymmetric int4 (explicit zeros)
46 + act_type=torch.bfloat16,
47 + group_size=64,
48 + zero_points=True,
49 + has_g_idx=False,
50 + )
51 +
52 + kernel_type = choose_mp_linear_kernel(config)
53 + assert kernel_type.__name__ == "TritonW4A16LinearKernel"
vllm/model_executor/kernels/linear/__init__.py
+5
@@ -51,6 +51,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import (
51 from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
52 MarlinLinearKernel,
53 )
54 +from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
55 + TritonW4A16LinearKernel,
56 +)
57 from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
58 XPUW4A8IntLinearKernel,
59 XPUwNa16LinearKernel,
@@ -203,6 +206,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
206 ExllamaLinearKernel,
207 ],
208 PlatformEnum.ROCM: [
209 + TritonW4A16LinearKernel,
210 ConchLinearKernel,
211 ExllamaLinearKernel,
212 ],
@@ -667,6 +671,7 @@ __all__ = [
671 "ExllamaLinearKernel",
672 "MacheteLinearKernel",
673 "MarlinLinearKernel",
674 + "TritonW4A16LinearKernel",
675 "XPUW4A8IntLinearKernel",
676 "XPUwNa16LinearKernel",
677 "CutlassNvFp4LinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+4
@@ -29,6 +29,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
29 MPLinearKernel,
30 MPLinearLayerConfig,
31 )
32 +from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
33 + TritonW4A16LinearKernel,
34 +)
35 from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
36 XPUW4A8IntLinearKernel,
37 XPUwNa16LinearKernel,
@@ -45,6 +48,7 @@ __all__ = [
48 "ExllamaLinearKernel",
49 "MacheteLinearKernel",
50 "MarlinLinearKernel",
51 + "TritonW4A16LinearKernel",
52 "XPUW4A8IntLinearKernel",
53 "XPUwNa16LinearKernel",
54 ]
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py new
+430
@@ -0,0 +1,430 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Triton-based W4A16 GEMM kernel for ROCm MI300.
5 +
6 +Implements fused int4-weight dequantization + fp16 GEMM in a single kernel,
7 +using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]).
8 +Plugs into the MPLinearKernel selection system and is preferred over
9 +MarlinLinearKernel/ExllamaLinearKernel on ROCm.
10 +
11 +Weight layout expected by this kernel (post-process_weights_after_loading):
12 + qweight: [K, N//8] int32 — rows=K (input), cols=N//8 (N is packed)
13 + scales: [K//G, N] fp16/bf16
14 + qzeros: [K//G, N//8] int32 (optional; None for symmetric uint4b8)
15 +
16 +Checkpoint layout from compressed_tensors_wNa16 create_weights:
17 + weight_packed: [N, K//8] int32 (output_dim=0, input_dim=1, packed_dim=1)
18 + weight_scale: [N, K//G] fp16 (output_dim=0, input_dim=1)
19 + weight_zero_point: [N//8, K//G] int32 (output_dim=0, packed_dim=0)
20 +"""
21 +
22 +import torch
23 +
24 +from vllm.model_executor.layers.quantization.utils import replace_parameter
25 +from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
26 +from vllm.platforms import current_platform
27 +from vllm.scalar_type import scalar_types
28 +from vllm.triton_utils import tl, triton
29 +
30 +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
31 +
32 +TRITON_W4A16_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128, 256]
33 +TRITON_W4A16_SUPPORTED_QUANT_TYPES = [
34 + scalar_types.uint4b8, # symmetric GPTQ (bias=8)
35 + scalar_types.uint4, # asymmetric with explicit zeros
36 +]
37 +
38 +
39 +@triton.jit
40 +def triton_w4a16_gemm_kernel(
41 + # Pointers
42 + a_ptr, # [M, K] fp16/bf16 activations
43 + b_ptr, # [K, N//8] int32 packed 4-bit weights (N is the packed dim)
44 + scales_ptr, # [K//G, N] fp16/bf16 scales
45 + zeros_ptr, # [K//G, N//8] int32 packed zeros (unused when HAS_ZP=False)
46 + c_ptr, # [M, N] fp16/bf16 output
47 + # Dimensions
48 + M,
49 + N,
50 + K,
51 + # Strides
52 + stride_am,
53 + stride_ak,
54 + stride_bk,
55 + stride_bn, # stride in b along the packed N//8 dim
56 + stride_cm,
57 + stride_cn,
58 + # Quantization parameters
59 + group_size,
60 + # Whether explicit zero points are provided
61 + HAS_ZP: tl.constexpr,
62 + # Zero bias used when HAS_ZP is False (e.g. 8 for uint4b8)
63 + ZP_BIAS: tl.constexpr,
64 + # Block sizes (tuned for MI300 wavefront=64)
65 + BLOCK_M: tl.constexpr,
66 + BLOCK_N: tl.constexpr,
67 + BLOCK_K: tl.constexpr,
68 +):
69 + """
70 + Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]
71 +
72 + B is stored as [K, N//8] int32 using GPTQ sequential packing:
73 + each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28].
74 +
75 + Dequant: w_fp = (w_int4 - zero) * scale
76 + HAS_ZP=True: zero is loaded from zeros_ptr and unpacked
77 + HAS_ZP=False: zero = ZP_BIAS constant (e.g. 8 for uint4b8 symmetric)
78 + """
79 + pid_m = tl.program_id(0)
80 + pid_n = tl.program_id(1)
81 +
82 + # Row/col offsets for this tile
83 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
84 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
85 +
86 + # b/zeros are stored with N packed: N//8 int32 columns per K row
87 + offs_bn = pid_n * (BLOCK_N // 8) + tl.arange(0, BLOCK_N // 8)
88 +
89 + # GPTQ sequential shifts tiled across BLOCK_N:
90 + # [0,4,8,...,28] repeating for every group of 8 N-values.
91 + # Build 1D shifts_1d of length BLOCK_N: column j gets shift (j % 8) * 4.
92 + shifts_row = tl.arange(0, 8) * 4 # [8]
93 + shifts_1d_2d = tl.broadcast_to(shifts_row[None, :], (BLOCK_N // 8, 8))
94 + shifts_1d = tl.reshape(shifts_1d_2d, (BLOCK_N,)) # [BLOCK_N]
95 + # Broadcast to [BLOCK_K, BLOCK_N] for weight unpacking
96 + shifts = tl.broadcast_to(shifts_1d[None, :], (BLOCK_K, BLOCK_N))
97 +
98 + # Scales column offsets: full N-width (one scale per output neuron)
99 + offs_sn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
100 +
101 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
102 +
103 + for k_start in range(0, tl.cdiv(K, BLOCK_K)):
104 + offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
105 + mask_k = offs_k < K
106 +
107 + # ---- Load activations A: [BLOCK_M, BLOCK_K] ----
108 + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
109 + mask_a = (offs_m[:, None] < M) & mask_k[None, :]
110 + a = tl.load(a_ptrs, mask=mask_a, other=0.0)
111 +
112 + # ---- Load packed weights B: [BLOCK_K, BLOCK_N//8] int32 ----
113 + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
114 + mask_b = mask_k[:, None] & (offs_bn[None, :] < N // 8)
115 + b_packed = tl.load(b_ptrs, mask=mask_b, other=0)
116 +
117 + # ---- Unpack int4 weights → [BLOCK_K, BLOCK_N] ----
118 + # tl.interleave(x, x) doubles the last dim by interleaving.
119 + # Starting from [BLOCK_K, BLOCK_N//8], three interleaves give
120 + # [BLOCK_K, BLOCK_N], where each int32 is replicated 8 times.
121 + b = tl.interleave(b_packed, b_packed)
122 + b = tl.interleave(b, b)
123 + b = tl.interleave(b, b)
124 + # Extract the correct 4-bit nibble for each output column
125 + b = (b >> shifts) & 0xF
126 +
127 + # ---- Compute scale/zero group row index ----
128 + g_idx = (k_start * BLOCK_K) // group_size
129 +
130 + # ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ----
131 + scale_offset = g_idx * N + offs_sn
132 + scale_mask = offs_sn < N
133 + scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0)
134 + scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N))
135 +
136 + # ---- Load / compute zeros ----
137 + if HAS_ZP:
138 + # Load packed zeros row: [BLOCK_N//8] int32
139 + zero_offset = g_idx * (N // 8) + offs_bn
140 + zero_mask = offs_bn < N // 8
141 + z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0)
142 + # Unpack to [BLOCK_N] using same interleave+shift pattern
143 + z = tl.interleave(z_packed, z_packed)
144 + z = tl.interleave(z, z)
145 + z = tl.interleave(z, z)
146 + z = (z >> shifts_1d) & 0xF
147 + z = tl.broadcast_to(z[None, :], (BLOCK_K, BLOCK_N))
148 + else:
149 + z = tl.full((BLOCK_K, BLOCK_N), ZP_BIAS, dtype=tl.int32)
150 +
151 + # ---- Dequantize: (w - zero) * scale ----
152 + b_fp = (b - z).to(a.dtype) * scales
153 +
154 + # ---- Accumulate ----
155 + accumulator += tl.dot(a, b_fp, out_dtype=tl.float32)
156 +
157 + # ---- Store output C: [BLOCK_M, BLOCK_N] ----
158 + c = accumulator.to(c_ptr.type.element_ty)
159 + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
160 + mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N)
161 + tl.store(c_ptrs, c, mask=mask_c)
162 +
163 +
164 +def triton_w4a16_gemm(
165 + a: torch.Tensor, # [M, K] fp16/bf16
166 + b_q: torch.Tensor, # [K, N//8] int32
167 + scales: torch.Tensor, # [K//G, N] fp16/bf16
168 + qzeros: torch.Tensor | None, # [K//G, N//8] int32, or None
169 + group_size: int,
170 + zp_bias: int = 8, # bias for uint4b8 when qzeros is None
171 +) -> torch.Tensor:
172 + """
173 + Fused W4A16 GEMM using GPTQ-packed int4 weights.
174 +
175 + Args:
176 + a: Activation matrix [M, K], float16 or bfloat16.
177 + b_q: Packed weight matrix [K, N//8], int32 (GPTQ sequential).
178 + scales: Per-group scales [K//G, N], same dtype as a.
179 + qzeros: Per-group packed zero points [K//G, N//8] int32, or None
180 + for symmetric quantization (uses zp_bias instead).
181 + group_size: Quantization group size (resolved from -1 to K by caller).
182 + zp_bias: Constant zero used when qzeros is None (default 8 for uint4b8).
183 +
184 + Returns:
185 + Output matrix [M, N], same dtype as a.
186 + """
187 + assert a.is_contiguous(), "Activation matrix must be contiguous"
188 + assert b_q.is_contiguous(), "Weight matrix must be contiguous"
189 + assert scales.is_contiguous(), "Scales must be contiguous"
190 +
191 + M, K = a.shape
192 + N = b_q.shape[1] * 8
193 +
194 + assert b_q.shape == (K, N // 8), (
195 + f"b_q shape mismatch: {b_q.shape} vs ({K}, {N // 8})"
196 + )
197 + assert scales.shape == (K // group_size, N), (
198 + f"scales shape mismatch: {scales.shape} vs ({K // group_size}, {N})"
199 + )
200 + if qzeros is not None:
201 + assert qzeros.shape == (K // group_size, N // 8), (
202 + f"qzeros shape mismatch: {qzeros.shape}"
203 + )
204 +
205 + c = torch.empty((M, N), dtype=a.dtype, device=a.device)
206 +
207 + has_zp = qzeros is not None
208 + # Provide a dummy pointer when HAS_ZP=False (Triton requires a valid ptr)
209 + zeros_ptr = qzeros if has_zp else b_q
210 +
211 + if current_platform.is_rocm():
212 + from vllm.platforms.rocm import on_gfx1x
213 +
214 + if on_gfx1x():
215 + # Tuned for RDNA 3.5 (gfx1151, 40 CUs, 32-wide wavefronts).
216 + if M <= 32:
217 + BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 64
218 + elif M <= 64:
219 + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
220 + else:
221 + BLOCK_M, BLOCK_N, BLOCK_K = 128, 32, 64
222 + else:
223 + # Tuned for MI300 (gfx942, 304 CUs, 64-wide wavefronts).
224 + if M <= 32:
225 + BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
226 + elif M <= 64:
227 + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
228 + else:
229 + BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
230 + else:
231 + if M <= 32:
232 + BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
233 + elif M <= 64:
234 + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
235 + else:
236 + BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
237 +
238 + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
239 +
240 + triton_w4a16_gemm_kernel[grid](
241 + a,
242 + b_q,
243 + scales,
244 + zeros_ptr,
245 + c,
246 + M,
247 + N,
248 + K,
249 + a.stride(0),
250 + a.stride(1),
251 + b_q.stride(0),
252 + b_q.stride(1),
253 + c.stride(0),
254 + c.stride(1),
255 + group_size=group_size,
256 + HAS_ZP=has_zp,
257 + ZP_BIAS=zp_bias,
258 + BLOCK_M=BLOCK_M,
259 + BLOCK_N=BLOCK_N,
260 + BLOCK_K=BLOCK_K,
261 + )
262 + return c
263 +
264 +
265 +class TritonW4A16LinearKernel(MPLinearKernel):
266 + """
267 + Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).
268 +
269 + Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric)
270 + with grouped quantization. Weight tensors are transposed from the
271 + compressed-tensors checkpoint layout to the kernel's [K, N//8] layout.
272 + """
273 +
274 + SUPPORTED_QUANT_TYPES = TRITON_W4A16_SUPPORTED_QUANT_TYPES
275 +
276 + @classmethod
277 + def get_min_capability(cls) -> int:
278 + # Triton handles capability checks itself
279 + return 0
280 +
281 + @classmethod
282 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
283 + if not current_platform.is_rocm():
284 + return False, "TritonW4A16LinearKernel only targets ROCm"
285 +
286 + if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
287 + return (
288 + False,
289 + f"Quant type {c.weight_type} not supported; "
290 + f"supported: {cls.SUPPORTED_QUANT_TYPES}",
291 + )
292 +
293 + if c.act_type not in (torch.float16, torch.bfloat16):
294 + return False, "Only float16/bfloat16 activations are supported"
295 +
296 + N = c.partition_weight_shape[1]
297 + if N % 8 != 0:
298 + return (
299 + False,
300 + f"Output features ({N}) must be divisible by 8 "
301 + "(8 int4 values packed per int32)",
302 + )
303 +
304 + if c.has_g_idx:
305 + return (
306 + False,
307 + "Activation reordering (g_idx) is not supported by "
308 + "TritonW4A16LinearKernel",
309 + )
310 +
311 + gs = c.group_size
312 + if (
313 + gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES
314 + and gs != c.full_weight_shape[0]
315 + ):
316 + return (
317 + False,
318 + f"Group size {gs} not supported; "
319 + f"supported: {TRITON_W4A16_SUPPORTED_GROUP_SIZES} "
320 + f"or full K ({c.full_weight_shape[0]})",
321 + )
322 +
323 + K = c.partition_weight_shape[0]
324 + eff_gs = gs if gs != -1 else K
325 + if K % eff_gs != 0:
326 + return (False, f"Input features {K} not divisible by group size {eff_gs}")
327 +
328 + return True, None
329 +
330 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
331 + """
332 + Convert compressed-tensors checkpoint layout to kernel layout.
333 +
334 + Checkpoint (from compressed_tensors_wNa16.create_weights):
335 + weight_packed: [N, K//8] int32 input_dim=1, output_dim=0, packed_dim=1
336 + weight_scale: [N, K//G] fp16 input_dim=1, output_dim=0
337 + weight_zero_point: [N//8, K//G] int32 output_dim=0, packed_dim=0
338 +
339 + Kernel needs:
340 + qweight: [K, N//8] int32 (transpose weight_packed)
341 + scales: [K//G, N] fp16 (transpose weight_scale)
342 + qzeros: [K//G, N//8] int32 (transpose weight_zero_point)
343 + """
344 +
345 + # ---- Transform qweight: [N, K//8] → [K//8, N] → back to [K, N//8] ----
346 + # permute_param_layout_(x, input_dim=0, output_dim=1) rearranges so that
347 + # the input(K) dimension is at physical dim 0 and output(N) at dim 1.
348 + # Checkpoint has input_dim=1, output_dim=0, packed_dim=1 (K is packed).
349 + # After permute we get [K//8, N] (K packed at dim 0, N at dim 1).
350 + # The kernel wants [K, N//8] (K at dim 0, N packed at dim 1), so we
351 + # then transpose: [K//8, N].T = [N, K//8] — that's not right.
352 + #
353 + # Actually we need to change WHAT is packed:
354 + # Original packing: K packed into K//8 (8 K-values per int32)
355 + # Kernel packing: N packed into N//8 (8 N-values per int32)
356 + # These require a full repack, not just a transpose.
357 + #
358 + # Simple approach: unpack → transpose the full [N, K] → repack as [K, N//8].
359 + # This is done CPU-side at load time (one-time cost).
360 + def repack_w_q(x: BasevLLMParameter) -> BasevLLMParameter:
361 + # x.data is [N, K//8] int32, K packed (GPTQ checkpoint format)
362 + # Step 1: bring to [N, K//8] with output(N) at dim 0
363 + permute_param_layout_(x, input_dim=1, output_dim=0, packed_dim=1)
364 + w = x.data # [N, K//8] int32
365 +
366 + N_dim, K8 = w.shape
367 + K_dim = K8 * 8
368 + # Step 2: unpack to [N, K] int32 (vectorized)
369 + shifts = torch.arange(8, device=w.device, dtype=torch.int32) * 4
370 + w_unpacked = ((w.unsqueeze(-1) >> shifts) & 0xF).reshape(N_dim, K_dim)
371 + # Step 3: transpose to [K, N] int32
372 + w_KN = w_unpacked.t().contiguous()
373 + # Step 4: repack N into N//8 int32 values → [K, N//8] (vectorized)
374 + N8 = N_dim // 8
375 + w_repacked = torch.sum(
376 + (w_KN.view(K_dim, N8, 8) & 0xF) << shifts,
377 + dim=2,
378 + dtype=torch.int32,
379 + )
380 + x.data = w_repacked.contiguous()
381 + return x
382 +
383 + def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter:
384 + # x.data is [N, K//G] fp16, bring to [K//G, N]
385 + permute_param_layout_(x, input_dim=1, output_dim=0)
386 + x.data = x.data.t().contiguous()
387 + return x
388 +
389 + self._transform_param(layer, self.w_q_name, repack_w_q)
390 + self._transform_param(layer, self.w_s_name, repack_w_s)
391 +
392 + if self.w_zp_name is not None:
393 + zp = getattr(layer, self.w_zp_name, None)
394 + if zp is not None:
395 + # Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1)
396 + # Kernel needs: [K//G, N//8] — just transpose
397 + replace_parameter(
398 + layer,
399 + self.w_zp_name,
400 + torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False),
401 + )
402 +
403 + def apply_weights(
404 + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
405 + ) -> torch.Tensor:
406 + c = self.config
407 + w_q, w_s, w_zp, _ = self._get_weight_params(layer)
408 +
409 + x_2d = x.reshape(-1, x.shape[-1]).contiguous()
410 + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
411 +
412 + K = c.partition_weight_shape[0]
413 + group_size = c.group_size if c.group_size != -1 else K
414 +
415 + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor
416 + zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0
417 +
418 + output = triton_w4a16_gemm(
419 + a=x_2d,
420 + b_q=w_q,
421 + scales=w_s,
422 + qzeros=w_zp,
423 + group_size=group_size,
424 + zp_bias=zp_bias,
425 + )
426 +
427 + if bias is not None:
428 + output.add_(bias)
429 +
430 + return output.reshape(out_shape)