[Kernel][Helion][1/N] Add Helion kernel for per_token_group_fp8_quant (vllm-project#36902)

Signed-off-by: Sean Chen <seachen@redhat.com>
Co-authored-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

.buildkite/test-amd.yaml

@@ -398,7 +398,7 @@ steps:
   - tests/kernels/helion/
   - vllm/platforms/rocm.py
   commands:
-  - pip install helion==1.0.0
+  - pip install helion==1.1.0
   - pytest -v -s kernels/helion/
 
 - label: Kernels Mamba Test # TBD

.buildkite/test_areas/kernels.yaml

@@ -237,7 +237,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==1.0.0
+    - pip install helion==1.1.0
     - pytest -v -s kernels/helion/
 
 

setup.py

@@ -1229,7 +1229,7 @@ def add_vllm_package_data(filename: str) -> None:
         # NOTE: When updating helion version, also update CI files:
         #   - .buildkite/test_areas/kernels.yaml
         #   - .buildkite/test-amd.yaml
-        "helion": ["helion==1.0.0"],
+        "helion": ["helion==1.1.0"],
         # Optional deps for gRPC server (vllm serve --grpc)
         "grpc": ["smg-grpc-servicer[vllm] >= 0.5.2"],
         # Optional deps for OpenTelemetry tracing

tests/kernels/helion/test_per_token_group_fp8_quant.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the per_token_group_fp8_quant helion kernel
+
+Run `pytest tests/kernels/helion/test_per_token_group_fp8_quant.py`.
+"""
+
+from typing import Any
+
+import pytest
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from tests.kernels.helion.utils import skip_if_platform_unsupported
+from tests.kernels.quant_utils import FP8_DTYPE
+from vllm.kernels.helion.case_key import CaseKey
+from vllm.kernels.helion.config_manager import ConfigManager
+from vllm.kernels.helion.ops.per_token_group_fp8_quant import (
+    _pick_cache,
+    baseline,
+    per_token_group_fp8_quant,
+    pick_config,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+)
+from vllm.utils.import_utils import has_helion
+
+if not has_helion():
+    pytest.skip(
+        "Helion is not installed. Install with: pip install vllm[helion]",
+        allow_module_level=True,
+    )
+
+
+def _generate_fake_input(
+    num_tokens: int, hidden_size: int, group_size: int
+) -> tuple[Any, ...]:
+    with FakeTensorMode():
+        input = torch.randn(
+            (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
+        )
+        output_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
+        output_s = torch.empty(
+            (num_tokens, hidden_size // group_size),
+            device=input.device,
+            dtype=torch.float32,
+        )
+        use_ue8m0 = False
+        column_major = False
+        fp8_min, fp8_max = get_fp8_min_max()
+        eps = 1e-10
+        args = (
+            input,
+            output_q,
+            output_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            use_ue8m0,
+            column_major,
+        )
+        return args
+
+
+@pytest.fixture(autouse=True)
+def reset_config_manager_singleton():
+    ConfigManager.reset_instance()
+    ConfigManager()
+    yield
+    ConfigManager.reset_instance()
+
+
+class TestPerTokenGroupFp8QuantConfigPicker:
+    def setup_method(self):
+        _pick_cache.clear()
+
+    def test_config_picker_exact_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+        ]
+
+        args = _generate_fake_input(16, 4096, 128)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
+        )
+
+    def test_config_picker_closest_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
+        ]
+
+        args = _generate_fake_input(20, 3000, 70)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
+        )
+
+    def test_config_picker_no_configs(self):
+        config_keys: list[dict] = []
+
+        args = _generate_fake_input(16, 4096, 128)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key is None
+
+    def test_config_picker_fallback_to_largest(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
+        ]
+
+        args = _generate_fake_input(64, 8192, 256)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey(
+            {"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
+        )
+
+
+class TestPerTokenGroupFp8QuantCorrectness:
+    @pytest.mark.parametrize(
+        "shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512), (2048, 5120)]
+    )
+    @pytest.mark.parametrize("column_major", [False, True])
+    @pytest.mark.parametrize("tma_aligned", [False, True])
+    @pytest.mark.parametrize("scale_ue8m0", [False, True])
+    @pytest.mark.parametrize("group_size", [64, 128])
+    def test_per_token_group_fp8_quant(
+        self,
+        shape,
+        column_major: bool,
+        tma_aligned: bool,
+        scale_ue8m0: bool,
+        group_size: int,
+    ):
+        skip_if_platform_unsupported("per_token_group_fp8_quant")
+
+        torch.manual_seed(42)
+        num_tokens, hidden_size = shape
+        fp8_min, fp8_max = get_fp8_min_max()
+        eps = 1e-10
+        input = (
+            torch.randn((num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16)
+            * 8
+        )
+        ref_q = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
+        ops_q = ref_q.clone()
+
+        groups_per_row = hidden_size // group_size
+        if column_major:
+            if tma_aligned:
+                tma_alignment = 4
+                tma_aligned_m = (
+                    (num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
+                )
+                shape = (num_tokens, groups_per_row)
+                stride = (1, tma_aligned_m)
+                ref_s = torch.empty_strided(
+                    shape, stride, device=input.device, dtype=torch.float32
+                )
+            else:
+                ref_s = torch.empty(
+                    (groups_per_row, num_tokens),
+                    device=input.device,
+                    dtype=torch.float32,
+                ).transpose(0, 1)
+        else:
+            ref_s = torch.empty(
+                (num_tokens, groups_per_row), device=input.device, dtype=torch.float32
+            )
+
+        ops_s = ref_s.clone()
+
+        baseline(
+            input,
+            ref_q,
+            ref_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            column_major,
+            tma_aligned,
+        )
+        per_token_group_fp8_quant(
+            input,
+            ops_q,
+            ops_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            scale_ue8m0,
+            column_major,
+            tma_aligned,
+        )
+
+        assert torch.allclose(ref_s, ops_s)
+        # allow 1 ULP difference
+        assert (
+            ref_q.view(torch.uint8).to(torch.int16)
+            - ops_q.view(torch.uint8).to(torch.int16)
+        ).abs().max() <= 1
+
+
+class TestPerTokenGroupFp8QuantIntegration:
+    def test_kernel_registration_integration(self):
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        assert "per_token_group_fp8_quant" in registered_kernels
+
+        kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
+        assert kernel_wrapper.op_name == "per_token_group_fp8_quant"
+        assert kernel_wrapper._config_picker is not None
+        assert kernel_wrapper._mutates_args == ["output_q", "output_s"]
+
+    def test_fake_impl_functionality(self):
+        skip_if_platform_unsupported("per_token_group_fp8_quant")
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        kernel_wrapper = registered_kernels["per_token_group_fp8_quant"]
+        fake_impl = kernel_wrapper._fake_impl
+
+        args = _generate_fake_input(16, 4096, 128)
+        assert fake_impl(*args) is None

tests/kernels/helion/test_register.py

@@ -713,6 +713,7 @@ def default_picker(args, config_keys):
 
         new_op = Mock()
         registered_ops: dict[str, Mock] = {}
+        mutates_args = ["y"]
 
         class MockNamespace:
             def __getattr__(self, name):
@@ -748,13 +749,15 @@ def register_side_effect(op_name, op_func, **kwargs):
                 raw_kernel_func=sample_kernel,
                 op_name="test_kernel",
                 fake_impl=fake_impl,
+                mutates_args=mutates_args,
                 config_picker=default_picker,
             )
             result = wrapper._get_or_register_custom_op()
 
             mock_register.assert_called_once()
             assert result is new_op
             assert mock_register.call_args[1]["op_func"] is mock_decorated
+            assert mock_register.call_args[1]["mutates_args"] is mutates_args
 
 
 class TestKernelRegistry:

tests/kernels/helion/utils.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Helion Kernel test utils"""
+
+import pytest
+import torch
+
+from vllm.kernels.helion.config_manager import ConfigManager
+
+
+def skip_if_platform_unsupported(op_name: str):
+    try:
+        from vllm.kernels.helion.utils import get_canonical_gpu_name
+
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA not available")
+
+        platform = get_canonical_gpu_name()
+
+        try:
+            config_manager = ConfigManager.get_instance()
+        except RuntimeError:
+            config_manager = ConfigManager()
+
+        configs = config_manager.get_platform_configs(op_name, platform)
+        if len(configs) == 0:
+            pytest.skip(f"Current GPU platform not supported for {op_name} kernel")
+
+    except (ImportError, RuntimeError, KeyError):
+        pytest.skip(f"Error detecting platform support for {op_name} kernel")