[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>

tests/kernels/quantization/test_triton_w4a16.py

@@ -10,6 +10,8 @@ import importlib
 
 import pytest
 import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.proxy_tensor import make_fx
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
@@ -107,6 +109,37 @@ def _w4a16_reference(
     return out.to(a_mk.dtype)
 
 
+def test_triton_w4a16_gemm_fake_impl():
+    with FakeTensorMode():
+        a = torch.empty((4, 256), device=device, dtype=torch.float16)
+        b_q = torch.empty((256, 64), device=device, dtype=torch.int32)
+        scales = torch.empty((2, 512), device=device, dtype=torch.float16)
+
+        out = triton_w4a16_gemm(
+            a=a,
+            b_q=b_q,
+            scales=scales,
+            qzeros=None,
+            group_size=128,
+        )
+
+        assert out.shape == (4, 512)
+        assert out.dtype == a.dtype
+        assert out.device == a.device
+
+        graph = make_fx(
+            lambda a, b_q, scales: triton_w4a16_gemm(
+                a, b_q, scales, qzeros=None, group_size=128
+            )
+        )(a, b_q, scales)
+
+        assert any(
+            node.op == "call_function"
+            and node.target == torch.ops.vllm.triton_w4a16_gemm.default
+            for node in graph.graph.nodes
+        )
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize(

vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py

@@ -26,6 +26,7 @@ from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layou
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
 
@@ -161,7 +162,7 @@ def triton_w4a16_gemm_kernel(
     tl.store(c_ptrs, c, mask=mask_c)
 
 
-def triton_w4a16_gemm(
+def _triton_w4a16_gemm_impl(
     a: torch.Tensor,  # [M, K] fp16/bf16
     b_q: torch.Tensor,  # [K, N//8] int32
     scales: torch.Tensor,  # [K//G, N] fp16/bf16
@@ -270,6 +271,36 @@ def triton_w4a16_gemm(
     return c
 
 
+def _triton_w4a16_gemm_fake(
+    a: torch.Tensor,
+    b_q: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor | None,
+    group_size: int,
+    zp_bias: int = 8,
+) -> torch.Tensor:
+    return torch.empty((a.size(0), b_q.size(1) * 8), dtype=a.dtype, device=a.device)
+
+
+direct_register_custom_op(
+    op_name="triton_w4a16_gemm",
+    op_func=_triton_w4a16_gemm_impl,
+    mutates_args=[],
+    fake_impl=_triton_w4a16_gemm_fake,
+)
+
+
+def triton_w4a16_gemm(
+    a: torch.Tensor,
+    b_q: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor | None,
+    group_size: int,
+    zp_bias: int = 8,
+) -> torch.Tensor:
+    return torch.ops.vllm.triton_w4a16_gemm(a, b_q, scales, qzeros, group_size, zp_bias)
+
+
 class TritonW4A16LinearKernel(MPLinearKernel):
     """
     Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).