biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel][Helion] [16/N] Refactor register_kernel API to be more Dynamo-friendly (#36705)

Browse Source 
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Yanan Cao
2026-03-17 18:23:35 -07:00
committed by GitHub
parent e6c4797704
commit ff9fbc9aff
6 changed files with 661 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
tests/kernels/helion/helpers.py Normal file
View File

@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import tempfile
from collections.abc import Callable
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import helion
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.register import register_kernel
from vllm.kernels.helion.utils import get_canonical_gpu_name
GPU_PLATFORM = get_canonical_gpu_name()
DEFAULT_CONFIGS: dict[str, helion.Config] = {
"default": helion.Config(block_sizes=[32]),
}
@contextmanager
def dummy_kernel_registry(
configs: dict[str, helion.Config] | None = None,
):
"""Context manager providing a register function with automatic config setup.
Yields a ``register`` callable with the same signature as
``register_kernel``. Before applying the real decorator it writes a
config JSON for the kernel name (from ``op_name`` or ``fn.__name__``)
into a temporary directory backed by a fresh ``ConfigManager``.
"""
if configs is None:
configs = DEFAULT_CONFIGS
config_data = {k: v.__dict__["config"] for k, v in configs.items()}
with tempfile.TemporaryDirectory() as tmpdir:
config_dir = Path(tmpdir)
ConfigManager.reset_instance()
cm = ConfigManager(base_dir=config_dir)
with patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=cm,
):
def register(
op_name: str | None = None,
**kwargs,
) -> Callable:
def decorator(fn: Callable) -> Callable:
name = op_name or fn.__name__
kernel_dir = config_dir / name
kernel_dir.mkdir(parents=True, exist_ok=True)
(kernel_dir / f"{GPU_PLATFORM}.json").write_text(
json.dumps(config_data)
)
return register_kernel(op_name, **kwargs)(fn)
return decorator
try:
yield register
finally:
ConfigManager.reset_instance()

91
tests/kernels/helion/test_autotune.py Normal file
View File

@@ -0,0 +1,91 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for autotuning Helion kernels, including disabled kernels with no configs."""
import pytest
import torch
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,
)
import helion
import helion.language as hl
from helion.autotuner.base_search import BaseSearch
from tests.kernels.helion.helpers import dummy_kernel_registry
from vllm.kernels.helion.register import create_helion_decorated_kernel
def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
for tile in hl.tile(x.size()):
out[tile] = x[tile] + y[tile]
return out
class NoCompileSearch(BaseSearch):
"""Autotuner that returns the default config without GPU compilation.
Modeled after helion's test BasicSearch (pytorch/helion#1649).
"""
def autotune(self, *, skip_cache: bool = False):
return self.config_spec.default_config()
def _no_compile_autotuner_fn(bound_kernel, args, **kwargs):
return NoCompileSearch(bound_kernel, args, **kwargs)
class TestAutotuneDisabledKernel:
"""Test autotuning flow on disabled kernels (no platform configs)."""
def setup_method(self):
from vllm.kernels.helion.register import _REGISTERED_KERNELS
self._saved_registry = dict(_REGISTERED_KERNELS)
_REGISTERED_KERNELS.clear()
def teardown_method(self):
from vllm.kernels.helion.register import _REGISTERED_KERNELS
_REGISTERED_KERNELS.clear()
_REGISTERED_KERNELS.update(self._saved_registry)
def test_autotune_disabled_kernel_produces_valid_config(self):
"""Register a kernel with no configs (disabled), run autotune,
verify it produces a valid helion.Config."""
with dummy_kernel_registry(configs={}) as register:
wrapper = register(
"autotune_test_kernel",
config_picker=lambda args, keys: "default",
fake_impl=lambda *a, **kw: None,
input_generator=lambda: {
"small": (
torch.randn(4, 4, device="cuda"),
torch.randn(4, 4, device="cuda"),
),
},
)(_add_kernel)
assert wrapper._disabled is True
inputs = wrapper.get_inputs()
assert "small" in inputs
settings = helion.Settings()
settings.autotuner_fn = _no_compile_autotuner_fn
wrapper.helion_settings = settings
config = wrapper.run_autotune(inputs["small"])
expected_default = (
create_helion_decorated_kernel(_add_kernel, helion_settings=settings)
.bind(inputs["small"])
.config_spec.default_config()
)
assert config == expected_default

6
tests/kernels/helion/test_pattern_matching.py
View File

@@ -52,7 +52,7 @@ def _helion_mock_context():
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -87,8 +87,8 @@ class TestMakeFxHop:
raw_kernel_func=raw_add_scale,
op_name="test_make_fx",
fake_impl=lambda *a, **kw: None,
config_picker=lambda args, keys: "default",
)
wrapper.register_config_picker(lambda args, keys: "default")
def fn(x, y):
return wrapper(x, y, scale)
@@ -143,8 +143,8 @@ class TestMakeFxHop:
raw_kernel_func=raw_silu_mul,
op_name="test_pm_silu_mul",
fake_impl=lambda *a, **kw: None,
config_picker=lambda args, keys: "default",
)
wrapper.register_config_picker(lambda args, keys: "default")
def pattern(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
return torch.nn.functional.silu(x) * y

568
tests/kernels/helion/test_register.py
View File

@@ -21,7 +21,9 @@ if not has_helion():
)
import helion
import helion.language as hl
from tests.kernels.helion.helpers import dummy_kernel_registry
from vllm.kernels.helion.config_manager import ConfigManager
from vllm.kernels.helion.register import (
_HOP_AVAILABLE,
@@ -34,6 +36,13 @@ from vllm.kernels.helion.register import (
)
def _add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
for tile in hl.tile(x.size()):
out[tile] = x[tile] + y[tile]
return out
@pytest.fixture
def sample_configs():
"""Create real Helion config objects for testing."""
@@ -90,7 +99,7 @@ def configured_kernel(sample_kernel, sample_configs, config_manager_with_test_co
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=config_manager_with_test_configs,
),
patch(
@@ -158,7 +167,7 @@ def create_configured_kernel_with_configs(
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -189,7 +198,7 @@ class TestConfiguredHelionKernel:
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -266,7 +275,7 @@ class TestConfiguredHelionKernel:
with (
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -310,7 +319,7 @@ class TestConfiguredHelionKernel:
with (
patch("vllm.kernels.helion.register.helion.kernel") as mock_helion_kernel,
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -346,23 +355,15 @@ class TestConfiguredHelionKernel:
class TestHelionKernelWrapper:
"""Test suite for HelionKernelWrapper."""
def test_get_configured_op_validates_configs_available(self, sample_kernel):
"""Test get_configured_op validates configs are available."""
def test_init_disables_on_missing_configs(self, sample_kernel):
"""Test __init__ marks wrapper as disabled when configs are missing."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
def default_picker(args, config_keys):
return "default"
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(
return_value={}
@@ -370,72 +371,7 @@ class TestHelionKernelWrapper:
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
pytest.raises(ValueError, match="No configs available"),
):
wrapper.get_configured_op()
def test_get_configured_op_validates_config_picker(
self, sample_kernel, sample_configs
):
"""Test get_configured_op validates config picker."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
# Don't set config picker - should raise assertion error
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
pytest.raises(AssertionError, match="No config picker registered"),
):
wrapper.get_configured_op()
def test_get_configured_op_returns_cached_kernel(
self, sample_kernel, sample_configs
):
"""Test get_configured_op returns cached ConfiguredHelionKernel."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -444,13 +380,269 @@ class TestHelionKernelWrapper:
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
assert wrapper._disabled is True
assert "No configs available" in wrapper._disabled_reason
def test_disabled_wrapper_raises_on_call(self, sample_kernel):
"""Test __call__ raises RuntimeError on a disabled wrapper."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value={})
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
with pytest.raises(RuntimeError, match="is disabled"):
wrapper(torch.randn(4, 4), torch.randn(4, 4))
def test_disabled_wrapper_get_configured_op_raises(self, sample_kernel):
"""Test get_configured_op raises RuntimeError on a disabled wrapper."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value={})
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
with pytest.raises(RuntimeError, match="is disabled"):
wrapper.get_configured_op()
def test_disabled_wrapper_supports_get_inputs(self, sample_kernel):
"""Test get_inputs works on a disabled wrapper."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
expected_inputs = {"key1": (torch.randn(4),)}
input_gen = Mock(return_value=expected_inputs)
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value={})
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
input_generator=input_gen,
)
assert wrapper._disabled is True
result = wrapper.get_inputs()
assert result is expected_inputs
def test_disabled_wrapper_supports_run_autotune(self, sample_kernel):
"""Test run_autotune works on a disabled wrapper."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value={})
mock_config = Mock()
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
assert wrapper._disabled is True
with patch(
"vllm.kernels.helion.register.create_helion_decorated_kernel"
) as mock_create:
mock_autotune_kernel = Mock()
mock_autotune_kernel.autotune.return_value = mock_config
mock_create.return_value = mock_autotune_kernel
inputs = (torch.randn(4, 4),)
result = wrapper.run_autotune(inputs)
assert result is mock_config
def test_init_caches_configured_kernel(self, sample_kernel, sample_configs):
"""Test __init__ eagerly builds and caches ConfiguredHelionKernel."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
def default_picker(args, config_keys):
return "default"
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=sample_kernel)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
assert wrapper._configured_kernel is not None
result1 = wrapper.get_configured_op()
result2 = wrapper.get_configured_op()
assert result1 is result2
@pytest.mark.skipif(
not _HOP_AVAILABLE, reason="HOP path only used when HOP available"
)
def test_init_eagerly_initializes_hop_path(self):
"""Test that register_kernel eagerly builds the configured kernel
on the HOP path (no custom op registration needed)."""
from vllm.kernels.helion.utils import get_canonical_gpu_name
configs = {"default": helion.Config(block_sizes=[4, 4])}
with (
dummy_kernel_registry(configs=configs) as register,
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
wraps=get_canonical_gpu_name,
) as mock_gpu,
):
wrapper = register(
config_picker=lambda args, keys: "default",
)(_add_kernel)
mock_gpu.assert_called_once()
assert wrapper._configured_kernel is not None
with patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
side_effect=AssertionError("get_canonical_gpu_name called during __call__"),
):
x = torch.randn(4, 4, device="cuda")
y = torch.randn(4, 4, device="cuda")
result = wrapper(x, y)
expected = x + y
assert torch.allclose(result, expected)
@pytest.mark.skipif(
_HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
)
def test_init_eagerly_initializes(self):
"""Test that register_kernel eagerly loads configs and detects GPU
during construction so __call__ needs no further initialization."""
from vllm.kernels.helion.utils import get_canonical_gpu_name
with (
dummy_kernel_registry() as register,
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
wraps=get_canonical_gpu_name,
) as mock_gpu,
):
wrapper = register(
config_picker=lambda args, keys: "default",
)(_add_kernel)
# Init must have detected GPU and built the kernel
mock_gpu.assert_called_once()
assert wrapper._configured_kernel is not None
assert hasattr(torch.ops.vllm_helion, wrapper.op_name)
@pytest.mark.skipif(
_HOP_AVAILABLE, reason="CustomOp path not used when HOP available"
)
@@ -463,13 +655,6 @@ class TestHelionKernelWrapper:
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
@@ -479,7 +664,7 @@ class TestHelionKernelWrapper:
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -491,6 +676,13 @@ class TestHelionKernelWrapper:
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
result = wrapper._get_or_register_custom_op()
assert result is existing_op
@@ -506,13 +698,6 @@ class TestHelionKernelWrapper:
def default_picker(args, config_keys):
return "default"
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
)
wrapper._config_picker = default_picker
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
@@ -532,7 +717,7 @@ class TestHelionKernelWrapper:
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
@@ -548,6 +733,13 @@ class TestHelionKernelWrapper:
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
wrapper = HelionKernelWrapper(
raw_kernel_func=sample_kernel,
op_name="test_kernel",
fake_impl=fake_impl,
config_picker=default_picker,
)
result = wrapper._get_or_register_custom_op()
mock_register.assert_called_once()
@@ -584,11 +776,10 @@ class TestKernelRegistry:
def test_get_kernel_by_name_returns_kernel(self):
"""Test get_kernel_by_name returns registered kernel."""
wrapper = HelionKernelWrapper(
raw_kernel_func=Mock(),
op_name="test_kernel",
fake_impl=Mock(),
)
with dummy_kernel_registry() as register:
wrapper = register(
"test_kernel", config_picker=lambda args, keys: "default"
)(_add_kernel)
from vllm.kernels.helion.register import _REGISTERED_KERNELS
@@ -604,112 +795,87 @@ class TestKernelRegistry:
def test_register_kernel_auto_generates_fake_impl(self):
"""Test register_kernel auto-generates fake_impl when not provided."""
with patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer:
with (
dummy_kernel_registry() as register,
patch("vllm.kernels.helion.register.infer_fake_impl") as mock_infer,
):
mock_fake = Mock()
mock_infer.return_value = mock_fake
wrapper = register(
config_picker=lambda args, keys: "default",
)(_add_kernel)
def original_kernel(x):
return x
wrapper = register_kernel(original_kernel)
mock_infer.assert_called_once_with(original_kernel, None)
assert wrapper._fake_impl is mock_fake
mock_infer.assert_called_once_with(_add_kernel, None)
assert wrapper._fake_impl is mock_fake
def test_register_kernel_creates_wrapper(self):
"""Test register_kernel creates HelionKernelWrapper."""
def test_kernel(x):
return x
result = register_kernel("test_name")(test_kernel)
with dummy_kernel_registry() as register:
result = register("test_name", config_picker=lambda args, keys: "default")(
_add_kernel
)
assert isinstance(result, HelionKernelWrapper)
assert result.op_name == "test_name"
assert result.raw_kernel_func is test_kernel
assert result.raw_kernel_func is _add_kernel
def test_register_kernel_auto_detects_name(self):
"""Test register_kernel uses function name when no name provided."""
with dummy_kernel_registry() as register:
wrapper = register(config_picker=lambda args, keys: "default")(_add_kernel)
@register_kernel
def my_test_kernel(x):
return x
assert my_test_kernel.op_name == "my_test_kernel"
assert wrapper.op_name == "_add_kernel"
def test_register_kernel_registers_in_global_registry(self):
"""Test register_kernel adds wrapper to global registry."""
@register_kernel
def test_kernel(x):
return x
with dummy_kernel_registry() as register:
wrapper = register(
"test_kernel", config_picker=lambda args, keys: "default"
)(_add_kernel)
registered_kernels = get_registered_kernels()
assert "test_kernel" in registered_kernels
assert registered_kernels["test_kernel"] is test_kernel
assert registered_kernels["test_kernel"] is wrapper
def test_register_kernel_passes_helion_settings(self):
"""Test register_kernel passes helion_settings to wrapper."""
mock_settings = Mock()
mock_settings.to_dict.return_value = {"debug": True}
settings = helion.Settings()
settings.print_output_code = True
@register_kernel("test_name", helion_settings=mock_settings)
def test_kernel(x):
return x
with dummy_kernel_registry() as register:
result = register(
"test_name",
config_picker=lambda args, keys: "default",
helion_settings=settings,
)(_add_kernel)
assert test_kernel.helion_settings is mock_settings
assert result.helion_settings is settings
def test_register_kernel_supports_decorator_syntax(self):
"""Test register_kernel works with decorator arguments."""
mock_fake = Mock()
wrapper = register_kernel("custom_name", fake_impl=mock_fake)
def test_kernel(x):
return x
result = wrapper(test_kernel)
with dummy_kernel_registry() as register:
result = register(
"custom_name",
config_picker=lambda args, keys: "default",
fake_impl=mock_fake,
)(_add_kernel)
assert result.op_name == "custom_name"
assert result._fake_impl is mock_fake
def test_register_kernel_bare_decorator(self):
"""Test register_kernel works as bare decorator."""
@register_kernel
def test_kernel(x):
return x
assert isinstance(test_kernel, HelionKernelWrapper)
assert test_kernel.op_name == "test_kernel"
def test_registered_wrapper_can_register_config_picker(self):
"""Test that registered wrapper can register config picker."""
@register_kernel
def test_kernel(x):
return x
def my_picker(args, config_keys):
return "default"
result = test_kernel.register_config_picker(my_picker)
assert result is my_picker
assert test_kernel._config_picker is my_picker
def test_register_kernel_raises_on_duplicate_registration(self):
"""Test register_kernel raises error on duplicate names."""
with dummy_kernel_registry() as register:
register("duplicate_name", config_picker=lambda args, keys: "default")(
_add_kernel
)
@register_kernel("duplicate_name")
def kernel1(x):
return x
with pytest.raises(ValueError, match="already registered"):
@register_kernel("duplicate_name")
def kernel2(x):
return x
with pytest.raises(ValueError, match="already registered"):
register("duplicate_name", config_picker=lambda args, keys: "default")(
_add_kernel
)
def test_register_kernel_rejects_autotuner_fn_in_settings(self):
"""Test register_kernel rejects conflicting autotuner_fn."""
@@ -718,7 +884,11 @@ class TestKernelRegistry:
with pytest.raises(ValueError, match="uses a custom autotuner"):
@register_kernel("test", helion_settings=mock_settings)
@register_kernel(
"test",
config_picker=lambda args, keys: "default",
helion_settings=mock_settings,
)
def test_kernel(x):
return x
@@ -727,11 +897,47 @@ class TestKernelRegistry:
mock_settings = Mock()
mock_settings.to_dict.return_value = {"static_shapes": False}
with patch("vllm.kernels.helion.register.logger") as mock_logger:
with (
dummy_kernel_registry() as register,
patch("vllm.kernels.helion.register.logger") as mock_logger,
):
register(
"test",
config_picker=lambda args, keys: "default",
helion_settings=mock_settings,
)(_add_kernel)
@register_kernel("test", helion_settings=mock_settings)
def test_kernel(x):
return x
mock_logger.warning.assert_not_called()
# Should not call warning
mock_logger.warning.assert_not_called()
def test_disabled_kernel_appears_in_registry(self):
"""Test that a disabled wrapper is still in the global registry."""
def fake_impl(*args, **kwargs):
return torch.zeros_like(args[0])
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value={})
with (
patch(
"vllm.kernels.helion.config_manager.ConfigManager",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
):
mock_kernel.return_value = Mock(return_value=_add_kernel)
wrapper = register_kernel(
"disabled_kernel",
config_picker=lambda args, keys: "default",
fake_impl=fake_impl,
)(_add_kernel)
assert wrapper._disabled is True
registered = get_registered_kernels()
assert "disabled_kernel" in registered
assert registered["disabled_kernel"] is wrapper