[Kernel][Helion] Helion kernel lazy registration (#48264)
Signed-off-by: Sean Chen <seachen@redhat.com>
Xiaohong (Sean) Chen committed
Jul 15, 2026 at 11:42 UTC
2fa63e0fff41e2f6187a6d3b9bf49b25f9e065f8
3 files changed
+31
-6
scripts/autotune_helion_kernels.py
+4
@@ -37,6 +37,7 @@ try:
37
get_kernel_by_name,
38
get_registered_kernels,
39
)
40
+ from vllm.kernels.helion.ops import import_all_kernels
41
from vllm.kernels.helion.utils import get_canonical_gpu_name
42
from vllm.logger import init_logger
43
from vllm.utils.import_utils import has_helion
@@ -382,6 +383,9 @@ def main():
383
384
args = parser.parse_args()
385
386
+ # import all helion kernels
387
+ import_all_kernels()
388
+
389
import logging
390
391
if args.verbose:
vllm/kernels/helion/__init__.py
-1
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Helion integration for vLLM."""
4
5
-import vllm.kernels.helion.ops # noqa: F401 Auto-register all Helion ops
5
from vllm.kernels.helion.case_key import CaseKey
6
from vllm.kernels.helion.config_manager import (
7
ConfigManager,
vllm/kernels/helion/ops/__init__.py
+27
-5
@@ -1,11 +1,33 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Auto-import all Helion op modules to trigger kernel registration."""
3
+"""Helion kernel implementation.
4
+
5
+Importing this package does NOT register any kernels. Runtime code imports the
6
+specific op module it needs, e.g.::
7
+
8
+ from vllm.kernels.helion.ops import scaled_mm # noqa: F401
9
+
10
+which triggers that op's ``@register_kernel`` as an import side effect.
11
+
12
+Tools that need the full registry (e.g. scripts/autotune_helion_kernels.py)
13
+call ``import_all_ops()`` to force every op module to register.
14
+"""
15
16
import importlib
17
import pkgutil
18
8
-# Automatically import all submodules so that @register_kernel
9
-# decorators execute and register ops with torch.ops.vllm_helion.
10
-for _module_info in pkgutil.iter_modules(__path__):
11
- importlib.import_module(f"{__name__}.{_module_info.name}")
19
+
20
+def import_all_kernels() -> list[str]:
21
+ """Import every kernel submodule so all ``@register_kernel`` decorators run.
22
+
23
+ Returns:
24
+ The fully-qualified module names that were imported.
25
+ """
26
+ imported: list[str] = []
27
+ for module_info in pkgutil.iter_modules(__path__):
28
+ if module_info.ispkg:
29
+ continue
30
+ module_name = f"{__name__}.{module_info.name}"
31
+ importlib.import_module(module_name)
32
+ imported.append(module_name)
33
+ return imported