[Performance] Register Triton W4A16 GEMM as a custom op (#51453)
Signed-off-by: giuseppegrossi <ggrossi@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Giuseppe Grossi committed
Sep 1, 2026 at 23:55 UTC
ee3c00bbf47e0ef7e975705cc980b06ee5576bb0
2 files changed
+65
-1
tests/kernels/quantization/test_triton_w4a16.py
+33
@@ -10,6 +10,8 @@ import importlib
10
11
import pytest
12
import torch
13
+from torch._subclasses.fake_tensor import FakeTensorMode
14
+from torch.fx.experimental.proxy_tensor import make_fx
15
16
from vllm.platforms import current_platform
17
from vllm.utils.torch_utils import set_random_seed
@@ -107,6 +109,37 @@ def _w4a16_reference(
109
return out.to(a_mk.dtype)
110
111
112
+def test_triton_w4a16_gemm_fake_impl():
113
+ with FakeTensorMode():
114
+ a = torch.empty((4, 256), device=device, dtype=torch.float16)
115
+ b_q = torch.empty((256, 64), device=device, dtype=torch.int32)
116
+ scales = torch.empty((2, 512), device=device, dtype=torch.float16)
117
+
118
+ out = triton_w4a16_gemm(
119
+ a=a,
120
+ b_q=b_q,
121
+ scales=scales,
122
+ qzeros=None,
123
+ group_size=128,
124
+ )
125
+
126
+ assert out.shape == (4, 512)
127
+ assert out.dtype == a.dtype
128
+ assert out.device == a.device
129
+
130
+ graph = make_fx(
131
+ lambda a, b_q, scales: triton_w4a16_gemm(
132
+ a, b_q, scales, qzeros=None, group_size=128
133
+ )
134
+ )(a, b_q, scales)
135
+
136
+ assert any(
137
+ node.op == "call_function"
138
+ and node.target == torch.ops.vllm.triton_w4a16_gemm.default
139
+ for node in graph.graph.nodes
140
+ )
141
+
142
+
143
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
144
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
145
@pytest.mark.parametrize(
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+32
-1
@@ -26,6 +26,7 @@ from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layou
26
from vllm.platforms import current_platform
27
from vllm.scalar_type import scalar_types
28
from vllm.triton_utils import tl, triton
29
+from vllm.utils.torch_utils import direct_register_custom_op
30
31
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
32
@@ -161,7 +162,7 @@ def triton_w4a16_gemm_kernel(
162
tl.store(c_ptrs, c, mask=mask_c)
163
164
164
-def triton_w4a16_gemm(
165
+def _triton_w4a16_gemm_impl(
166
a: torch.Tensor, # [M, K] fp16/bf16
167
b_q: torch.Tensor, # [K, N//8] int32
168
scales: torch.Tensor, # [K//G, N] fp16/bf16
@@ -270,6 +271,36 @@ def triton_w4a16_gemm(
271
return c
272
273
274
+def _triton_w4a16_gemm_fake(
275
+ a: torch.Tensor,
276
+ b_q: torch.Tensor,
277
+ scales: torch.Tensor,
278
+ qzeros: torch.Tensor | None,
279
+ group_size: int,
280
+ zp_bias: int = 8,
281
+) -> torch.Tensor:
282
+ return torch.empty((a.size(0), b_q.size(1) * 8), dtype=a.dtype, device=a.device)
283
+
284
+
285
+direct_register_custom_op(
286
+ op_name="triton_w4a16_gemm",
287
+ op_func=_triton_w4a16_gemm_impl,
288
+ mutates_args=[],
289
+ fake_impl=_triton_w4a16_gemm_fake,
290
+)
291
+
292
+
293
+def triton_w4a16_gemm(
294
+ a: torch.Tensor,
295
+ b_q: torch.Tensor,
296
+ scales: torch.Tensor,
297
+ qzeros: torch.Tensor | None,
298
+ group_size: int,
299
+ zp_bias: int = 8,
300
+) -> torch.Tensor:
301
+ return torch.ops.vllm.triton_w4a16_gemm(a, b_q, scales, qzeros, group_size, zp_bias)
302
+
303
+
304
class TritonW4A16LinearKernel(MPLinearKernel):
305
"""
306
Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).