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