[Kernel][Helion][1/N] Add Helion kernel for dynamic_per_token_scaled_fp8_quant (#33790)

Signed-off-by: Sean Chen <seachen@redhat.com>
Co-authored-by: Yanan Cao <gmagogsfm@gmail.com>

tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel
+
+Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_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.dynamic_per_token_scaled_fp8_quant import (
+    _pick_cache,
+    baseline,
+    dynamic_per_token_scaled_fp8_quant,
+    pick_config,
+)
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_helion
+from vllm.utils.torch_utils import set_random_seed
+
+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) -> tuple[Any, ...]:
+    with FakeTensorMode():
+        input = torch.randn(
+            num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
+        )
+        result = torch.empty(
+            input.shape, device=input.device, dtype=current_platform.fp8_dtype()
+        )
+        scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
+        scale_ub = torch.mean(input).to(torch.float32)
+        args = (result, input, scale, scale_ub)
+        return args
+
+
+@pytest.fixture(autouse=True)
+def reset_config_manager_singleton():
+    ConfigManager.reset_instance()
+    ConfigManager()
+    yield
+    ConfigManager.reset_instance()
+
+
+class TestDynamicPerTokenScaledFp8QuantConfigPicker:
+    def setup_method(self):
+        _pick_cache.clear()
+
+    def test_config_picker_exact_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "num_tokens": 16}),
+        ]
+
+        args = _generate_fake_input(16, 4096)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
+
+    def test_config_picker_closest_match(self):
+        config_keys = [
+            CaseKey({"hidden_size": 2048, "num_tokens": 16}),
+            CaseKey({"hidden_size": 2048, "num_tokens": 32}),
+            CaseKey({"hidden_size": 4096, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "num_tokens": 32}),
+        ]
+
+        args = _generate_fake_input(20, 3000)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
+
+    def test_config_picker_no_configs(self):
+        config_keys: list[dict] = []
+
+        args = _generate_fake_input(16, 4096)
+        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, "num_tokens": 16}),
+            CaseKey({"hidden_size": 4096, "num_tokens": 16}),
+        ]
+
+        args = _generate_fake_input(32, 8192)
+        selected_key = pick_config(args, config_keys)
+        assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
+
+
+class TestDynamicPerTokenScaledFp8QuantCorrectness:
+    @pytest.mark.parametrize("num_tokens", [1, 7, 4096])
+    @pytest.mark.parametrize("hidden_size", [17, 1024, 1025, 1026, 5137, 8193])
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float])
+    @pytest.mark.parametrize("has_scale_ub", [True, False])
+    @pytest.mark.parametrize("seed", [0])
+    def test_dynamic_per_token_fp8_quant(
+        self,
+        num_tokens: int,
+        hidden_size: int,
+        dtype: torch.dtype,
+        has_scale_ub: bool,
+        seed: int,
+    ) -> None:
+        skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
+        set_random_seed(seed)
+
+        x = (
+            torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
+        )  # avoid nans
+
+        scale_ub = (
+            torch.mean(x).to(dtype=torch.float32, device="cuda")
+            if has_scale_ub
+            else None
+        )
+
+        ref_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
+        ref_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
+        baseline(ref_out, x, ref_scales, scale_ub)
+
+        ops_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
+        ops_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
+        dynamic_per_token_scaled_fp8_quant(ops_out, x, ops_scales, scale_ub)
+
+        torch.testing.assert_close(ref_scales, ops_scales)
+        # allow 1 ULP difference
+        assert (
+            ref_out.view(torch.uint8).to(torch.int16)
+            - ops_out.view(torch.uint8).to(torch.int16)
+        ).abs().max() <= 1
+
+
+class TestDynamicPerTokenScaledFp8QuantIntegration:
+    def test_kernel_registration_integration(self):
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        assert "dynamic_per_token_scaled_fp8_quant" in registered_kernels
+
+        kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
+        assert kernel_wrapper.op_name == "dynamic_per_token_scaled_fp8_quant"
+        assert kernel_wrapper._config_picker is not None
+        assert kernel_wrapper._mutates_args == ["result", "scale"]
+
+    def test_fake_impl_functionality(self):
+        skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
+        from vllm.kernels.helion.register import get_registered_kernels
+
+        registered_kernels = get_registered_kernels()
+        kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
+        fake_impl = kernel_wrapper._fake_impl
+
+        args = _generate_fake_input(16, 4096)
+        assert fake_impl(*args) is None

vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from itertools import product
+from typing import Any
+
+import torch
+
+from vllm.kernels.helion.case_key import CaseKey
+from vllm.kernels.helion.register import register_kernel
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
+)
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_helion
+
+if not has_helion():
+    raise ImportError(
+        "Helion kernel requires helion to be installed. "
+        "Install it with: pip install helion"
+    )
+
+import helion
+import helion.language as hl
+
+logger = init_logger(__name__)
+
+
+def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
+    # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input
+    # property combination. Currently, dtypes are fixed. We need optimization to
+    # bucket/skip some combinations
+    num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
+    hidden_size_list = [2048, 4096, 5120]
+    in_dtype: torch.dtype = torch.bfloat16
+    out_dtype: torch.dtype = current_platform.fp8_dtype()
+    scale_dtype: torch.dtype = torch.float32
+    inputs = {}
+    for num_tokens, hidden_size in product(num_tokens_list, hidden_size_list):
+        input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
+        result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
+        scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype)
+        scale_ub = torch.mean(input).to(scale_dtype)
+
+        config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens})
+        inputs[config_key] = (result, input, scale, scale_ub)
+
+    return inputs
+
+
+_pick_cache: dict[tuple[int, int], CaseKey | None] = {}
+
+
+def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
+    """Pick the best pre-tuned config for the given input shape.
+
+    Selection strategy:
+      1. Find the closest hidden_size among available configs
+         (exact match preferred).
+      2. Among the num_tokens values tuned for that hidden_size, pick
+         the smallest num_tokens >= the input's num_tokens. If the input is
+         larger than all available num_tokens, fall back to the largest.
+    """
+
+    if not config_keys:
+        return None
+
+    _, input, *_ = args
+    num_tokens, hidden_size = input.shape
+
+    cache_key = (num_tokens, hidden_size)
+    cached = _pick_cache.get(cache_key)
+    if cached is not None:
+        return cached
+
+    configs: dict[int, list[int]] = {}
+    for key in config_keys:
+        if key.is_default():
+            continue
+        configs.setdefault(key["hidden_size"], []).append(key["num_tokens"])
+
+    if not configs:
+        return None
+
+    best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size))
+    available_num_tokens = sorted(configs[best_hidden_size])
+    best_num_tokens = next(
+        (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
+    )
+
+    result = CaseKey({"hidden_size": best_hidden_size, "num_tokens": best_num_tokens})
+    _pick_cache[cache_key] = result
+    return result
+
+
+def fake_impl(
+    result: torch.Tensor,  # [num_tokens, hidden_size]
+    input: torch.Tensor,  # [num_tokens, hidden_size]
+    scale: torch.Tensor,  # [num_tokens, 1]
+    scale_ub: torch.Tensor | None = None,  # scalar tensor
+) -> None:
+    return
+
+
+def baseline(
+    result: torch.Tensor,  # [num_tokens, hidden_size]
+    input: torch.Tensor,  # [num_tokens, hidden_size]
+    scale: torch.Tensor,  # [num_tokens, 1]
+    scale_ub: torch.Tensor | None = None,  # scalar tensor
+) -> None:
+    torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub)
+
+
+# Overwrite autotune_baseline_atol and autotune_baseline_rtol
+# if too many configs failed due to baseline check during autotuning
+@register_kernel(
+    mutates_args=["result", "scale"],
+    config_picker=pick_config,
+    input_generator=generate_inputs,
+    fake_impl=fake_impl,
+    helion_settings=helion.Settings(
+        autotune_baseline_fn=baseline,
+        ignore_warnings=[helion.exc.TensorOperationInWrapper],
+    ),
+)
+def dynamic_per_token_scaled_fp8_quant(
+    result: torch.Tensor,  # [num_tokens, hidden_size]
+    input: torch.Tensor,  # [num_tokens, hidden_size]
+    scale: torch.Tensor,  # [num_tokens, 1]
+    scale_ub: torch.Tensor | None = None,  # scalar tensor
+) -> None:
+    # This code assumes batch_dim and num_tokens are flattened
+    assert input.ndim == 2
+    num_tokens, hidden_size = input.shape
+    hl.specialize(hidden_size)
+
+    assert result.shape == input.shape
+    assert scale.shape[0] == num_tokens
+    assert scale.dtype == torch.float32
+    assert input.stride()[-1] == 1
+    assert result.stride()[-1] == 1
+
+    fp8_min, fp8_max = get_fp8_min_max()
+    min_scaling_factor = 1.0 / (fp8_max * 512.0)
+
+    for tile_m in hl.tile(num_tokens, block_size=1):
+        s_blk = hl.zeros([tile_m], dtype=torch.float32)
+        for tile_n in hl.tile(hidden_size):
+            x_blk = input[tile_m, tile_n].to(dtype=torch.float32)
+            tmp_blk = torch.amax(torch.abs(x_blk), dim=-1)
+            s_blk = torch.maximum(s_blk, tmp_blk)
+
+        if scale_ub is not None:
+            scale_ub_s = hl.load(scale_ub, [])
+            s_blk = s_blk.clamp(max=scale_ub_s)
+        s_blk = s_blk * (1.0 / fp8_max)
+        s_blk = s_blk.clamp(min=min_scaling_factor)
+        scale[tile_m, 0] = s_blk
+
+        for tile_n in hl.tile(hidden_size):
+            x_blk = input[tile_m, tile_n].to(torch.float32)
+            y_blk = x_blk * (1.0 / s_blk[:, None])
+
+            result[tile_m, tile_n] = y_blk.clamp(fp8_min, fp8_max).to(result.dtype)