Skip to content

vllm.kernels.helion

Helion integration for vLLM.

Modules:

Name Description
case_key

Structured key for identifying kernel config/autotune/benchmark cases.

config_manager

Configuration management for Helion kernels.

ops

Auto-import all Helion op modules to trigger kernel registration.

register

vLLM Helion kernel registration with pre-tuned config selection.

utils

Utility functions for Helion kernel management.

CaseKey

Bases: dict[str, Any]

Immutable, hashable dict for identifying kernel cases.

Used as the key for config lookup, autotuning, benchmarking, and input generation. Behaves like a read-only dict and can be used as a dict key or in sets.

The canonical string form (__str__) is stable JSON with sorted keys. Use CaseKey.default() for the default/fallback key. The regular constructor requires at least one key-value pair::

CaseKey({"intermediate": 2048, "numtokens": 256})
CaseKey.default()  # default/fallback
Source code in vllm/kernels/helion/case_key.py
class CaseKey(dict[str, Any]):
    """Immutable, hashable dict for identifying kernel cases.

    Used as the key for config lookup, autotuning, benchmarking, and
    input generation.  Behaves like a read-only dict and can be used
    as a dict key or in sets.

    The canonical string form (``__str__``) is stable JSON with sorted
    keys.  Use ``CaseKey.default()`` for the default/fallback key.
    The regular constructor requires at least one key-value pair::

        CaseKey({"intermediate": 2048, "numtokens": 256})
        CaseKey.default()  # default/fallback
    """

    def __init__(self, *args: Any, _allow_empty: bool = False, **kwargs: Any):
        super().__init__(*args, **kwargs)
        if not self and not _allow_empty:
            raise TypeError(
                "CaseKey requires at least one key-value pair. "
                "Use CaseKey.default() for the default config key."
            )
        self._str: str | None = None
        self._hash: int | None = None

    @classmethod
    def default(cls) -> CaseKey:
        """Create a default case key (empty)."""
        return cls(_allow_empty=True)

    def __hash__(self) -> int:  # type: ignore[override]
        if self._hash is None:
            self._hash = hash(str(self))
        return self._hash

    def __str__(self) -> str:
        if self._str is None:
            self._str = json.dumps(dict(self), sort_keys=True, separators=(",", ":"))
        return self._str

    def __repr__(self) -> str:
        if not self:
            return "CaseKey.default()"
        return f"CaseKey({dict(self)})"

    def is_default(self) -> bool:
        """Return True if this is the default case key (empty)."""
        return not self

    def _readonly(self, *args: Any, **kwargs: Any) -> Any:
        raise TypeError("CaseKey is immutable")

    __setitem__ = _readonly  # type: ignore[assignment]
    __delitem__ = _readonly  # type: ignore[assignment]
    __ior__ = _readonly  # type: ignore[assignment]
    update = _readonly  # type: ignore[assignment]
    pop = _readonly  # type: ignore[assignment]
    popitem = _readonly  # type: ignore[assignment]
    setdefault = _readonly  # type: ignore[assignment]
    clear = _readonly  # type: ignore[assignment]

default classmethod

default() -> CaseKey

Create a default case key (empty).

Source code in vllm/kernels/helion/case_key.py
@classmethod
def default(cls) -> CaseKey:
    """Create a default case key (empty)."""
    return cls(_allow_empty=True)

is_default

is_default() -> bool

Return True if this is the default case key (empty).

Source code in vllm/kernels/helion/case_key.py
def is_default(self) -> bool:
    """Return True if this is the default case key (empty)."""
    return not self

ConfigManager

File-level configuration management for Helion kernels (global singleton).

Source code in vllm/kernels/helion/config_manager.py
class ConfigManager:
    """File-level configuration management for Helion kernels (global singleton)."""

    _instance: ConfigManager | None = None
    _instance_base_dir: Path | None = None

    def __new__(cls, base_dir: str | Path | None = None) -> ConfigManager:
        resolved_base_dir = cls._resolve_base_dir(base_dir)

        if cls._instance is not None:
            if cls._instance_base_dir != resolved_base_dir:
                raise ValueError(
                    f"ConfigManager singleton already exists with base_dir "
                    f"'{cls._instance_base_dir}', cannot create with "
                    f"different base_dir '{resolved_base_dir}'"
                )
            return cls._instance

        instance = super().__new__(cls)
        cls._instance = instance
        cls._instance_base_dir = resolved_base_dir
        return instance

    def __init__(self, base_dir: str | Path | None = None):
        if hasattr(self, "_base_dir"):
            return

        self._base_dir = self._resolve_base_dir(base_dir)
        logger.debug("ConfigManager initialized with base_dir: %s", self._base_dir)

    @staticmethod
    def _resolve_base_dir(base_dir: str | Path | None) -> Path:
        if base_dir is not None:
            return Path(base_dir).resolve()
        return (Path(__file__).parent / "configs").resolve()

    @classmethod
    def get_instance(cls) -> ConfigManager:
        if cls._instance is None:
            raise RuntimeError(
                "ConfigManager instance has not been created. "
                "Call ConfigManager(base_dir=...) first to initialize."
            )
        return cls._instance

    @classmethod
    def reset_instance(cls) -> None:
        """For testing purposes only."""
        cls._instance = None
        cls._instance_base_dir = None

    def get_kernel_dir(self, kernel_name: str) -> Path:
        return self._base_dir / kernel_name

    def get_config_file_path(
        self, kernel_name: str, platform: str | None = None
    ) -> Path:
        if platform is not None:
            return self.get_kernel_dir(kernel_name) / f"{platform}.json"
        return self.get_kernel_dir(kernel_name)

    def ensure_base_dir_exists(self) -> Path:
        self._base_dir.mkdir(parents=True, exist_ok=True)
        return self._base_dir

    def ensure_base_dir_writable(self) -> None:
        self.ensure_base_dir_exists()
        test_file = self._base_dir / ".write_test"
        try:
            test_file.write_text("test")
            test_file.unlink()
        except OSError as e:
            raise OSError(
                f"Config directory '{self._base_dir}' is not writable: {e}"
            ) from e

    def _load_platform_file(self, kernel_name: str, platform: str) -> Any:
        config_path = self.get_config_file_path(kernel_name, platform)
        if not config_path.exists():
            return []
        try:
            with open(config_path) as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            logger.error("Failed to load config file %s: %s", config_path, e)
            return []

    def load_config_set(self, kernel_name: str) -> ConfigSet:
        kernel_dir = self.get_kernel_dir(kernel_name)
        if not kernel_dir.is_dir():
            return ConfigSet.from_dict(kernel_name, {})

        data: dict[str, Any] = {}
        for platform_file in sorted(kernel_dir.glob("*.json")):
            platform = platform_file.stem
            try:
                with open(platform_file) as f:
                    platform_data = json.load(f)
                data[platform] = platform_data
            except (json.JSONDecodeError, OSError) as e:
                logger.error(
                    "Failed to load config file %s: %s",
                    platform_file,
                    e,
                )

        return ConfigSet.from_dict(kernel_name, data)

    def get_platform_configs(
        self, kernel_name: str, platform: str
    ) -> dict[CaseKey, helion.Config]:
        platform_data = self._load_platform_file(kernel_name, platform)
        if not platform_data:
            return {}
        config_set = ConfigSet.from_dict(kernel_name, {platform: platform_data})
        return {
            k: config_set.get_config(platform, k)
            for k in config_set.get_config_keys(platform)
        }

    def save_config_set(self, config_set: ConfigSet) -> Path:
        kernel_dir = self.get_kernel_dir(config_set.kernel_name)
        kernel_dir.mkdir(parents=True, exist_ok=True)

        full_data = config_set.to_config_entries()
        for platform, pairs in full_data.items():
            platform_path = kernel_dir / f"{platform}.json"
            with open(platform_path, "w") as f:
                json.dump(pairs, f, indent=2)
                f.write("\n")
            logger.info("Saved config to: %s", platform_path)

        return kernel_dir

    def save_configs(
        self,
        kernel_name: str,
        platform: str,
        configs: dict[CaseKey, helion.Config],
    ) -> Path:
        """Save configs for a kernel/platform, merging with existing."""
        config_set = ConfigSet.from_dict(
            kernel_name,
            {platform: self._load_platform_file(kernel_name, platform)},
        )
        for key, config in configs.items():
            config_set.set_config(platform, key, config)

        pairs = config_set.to_config_entries().get(platform, [])
        platform_path = self.get_config_file_path(kernel_name, platform)
        platform_path.parent.mkdir(parents=True, exist_ok=True)
        with open(platform_path, "w") as f:
            json.dump(pairs, f, indent=2)
            f.write("\n")

        logger.info("Saved config to: %s", platform_path)
        return platform_path

    def config_exists(
        self,
        kernel_name: str,
        platform: str,
        config_key: CaseKey,
    ) -> bool:
        platform_data = self._load_platform_file(kernel_name, platform)
        if not platform_data:
            return False
        target = dict(config_key)
        return any(entry["key"] == target for entry in platform_data)

reset_instance classmethod

reset_instance() -> None

For testing purposes only.

Source code in vllm/kernels/helion/config_manager.py
@classmethod
def reset_instance(cls) -> None:
    """For testing purposes only."""
    cls._instance = None
    cls._instance_base_dir = None

save_configs

save_configs(
    kernel_name: str,
    platform: str,
    configs: dict[CaseKey, Config],
) -> Path

Save configs for a kernel/platform, merging with existing.

Source code in vllm/kernels/helion/config_manager.py
def save_configs(
    self,
    kernel_name: str,
    platform: str,
    configs: dict[CaseKey, helion.Config],
) -> Path:
    """Save configs for a kernel/platform, merging with existing."""
    config_set = ConfigSet.from_dict(
        kernel_name,
        {platform: self._load_platform_file(kernel_name, platform)},
    )
    for key, config in configs.items():
        config_set.set_config(platform, key, config)

    pairs = config_set.to_config_entries().get(platform, [])
    platform_path = self.get_config_file_path(kernel_name, platform)
    platform_path.parent.mkdir(parents=True, exist_ok=True)
    with open(platform_path, "w") as f:
        json.dump(pairs, f, indent=2)
        f.write("\n")

    logger.info("Saved config to: %s", platform_path)
    return platform_path

ConfigSet

In-memory collection of Helion configs with lookup/query capabilities.

Configs are stored keyed by CaseKey. The default config uses CaseKey.default() as its key.

Source code in vllm/kernels/helion/config_manager.py
class ConfigSet:
    """In-memory collection of Helion configs with lookup/query capabilities.

    Configs are stored keyed by ``CaseKey``.  The default config
    uses ``CaseKey.default()`` as its key.
    """

    _ConfigDict = dict[str, dict[CaseKey, "helion.Config"]]

    def __init__(self, kernel_name: str):
        self._kernel_name = kernel_name
        self._configs: ConfigSet._ConfigDict = {}

    @property
    def kernel_name(self) -> str:
        return self._kernel_name

    def get_config(self, platform: str, config_key: CaseKey) -> helion.Config:
        platform_dict = self._configs.get(platform)
        if platform_dict is None:
            avail_platforms = self.get_platforms()
            # TODO(@gmagogsfm): add a CLI/env override flag so users can
            # directly specify a platform name instead of relying on
            # auto-detection, and suggest it in this error message.
            raise KeyError(
                f"Config not found for kernel '{self._kernel_name}': "
                f"platform '{platform}' not found. "
                f"Available platforms: {avail_platforms or '(none)'}. "
                f"If your GPU is a variant of a supported platform, "
                f"consider adding a mapping in _GPU_NAME_ALIASES in "
                f"vllm/kernels/helion/utils.py, or run "
                f"scripts/autotune_helion_kernels.py to generate configs "
                f"for your platform."
            )

        config = platform_dict.get(config_key)
        if config is None:
            avail_keys = self.get_config_keys(platform)
            raise KeyError(
                f"Config not found for kernel '{self._kernel_name}': "
                f"config_key '{config_key}' not found for "
                f"platform '{platform}'. "
                f"Available config_keys: {avail_keys or '(none)'}"
            )

        return config

    def get_platforms(self) -> list[str]:
        return sorted(self._configs.keys())

    def get_config_keys(self, platform: str) -> list[CaseKey]:
        platform_dict = self._configs.get(platform.lower())
        if platform_dict is None:
            return []
        return sorted(platform_dict.keys(), key=str)

    def to_config_entries(self) -> dict[str, list[dict[str, Any]]]:
        """Serialize to config entries format for JSON output."""
        result: dict[str, list[dict[str, Any]]] = {}
        for platform, config_dict in self._configs.items():
            pairs: list[dict[str, Any]] = []
            for config_key, config in config_dict.items():
                config_data = json.loads(config.to_json())
                pairs.append({"key": dict(config_key), "config": config_data})
            result[platform] = pairs
        return result

    def to_dict(self) -> dict[str, dict[CaseKey, Any]]:
        """Return configs as a nested dict (platform -> key -> config)."""
        result: dict[str, dict[CaseKey, Any]] = {}
        for platform, config_dict in self._configs.items():
            result[platform] = {
                k: json.loads(v.to_json()) for k, v in config_dict.items()
            }
        return result

    @classmethod
    def from_dict(cls, kernel_name: str, data: dict[str, Any]) -> ConfigSet:
        config_set = cls(kernel_name)
        count = 0

        for platform, platform_data in data.items():
            if platform not in config_set._configs:
                config_set._configs[platform] = {}

            for entry in platform_data:
                raw_key = entry["key"]
                key = CaseKey.default() if not raw_key else CaseKey(raw_key)
                config = helion.Config(**entry["config"])
                config_set._configs[platform][key] = config
                count += 1

        if count > 0:
            logger.debug(
                "Loaded %d configs for kernel '%s'",
                count,
                kernel_name,
            )

        return config_set

    def set_config(
        self,
        platform: str,
        config_key: CaseKey,
        config: helion.Config,
    ) -> None:
        platform = platform.lower()
        if platform not in self._configs:
            self._configs[platform] = {}
        self._configs[platform][config_key] = config
        logger.debug(
            "Set config for kernel '%s': platform='%s', key='%s'",
            self._kernel_name,
            platform,
            config_key,
        )

    def has_config(self, platform: str, config_key: CaseKey) -> bool:
        platform = platform.lower()
        platform_dict = self._configs.get(platform)
        if platform_dict is None:
            return False
        return config_key in platform_dict

to_config_entries

to_config_entries() -> dict[str, list[dict[str, Any]]]

Serialize to config entries format for JSON output.

Source code in vllm/kernels/helion/config_manager.py
def to_config_entries(self) -> dict[str, list[dict[str, Any]]]:
    """Serialize to config entries format for JSON output."""
    result: dict[str, list[dict[str, Any]]] = {}
    for platform, config_dict in self._configs.items():
        pairs: list[dict[str, Any]] = []
        for config_key, config in config_dict.items():
            config_data = json.loads(config.to_json())
            pairs.append({"key": dict(config_key), "config": config_data})
        result[platform] = pairs
    return result

to_dict

to_dict() -> dict[str, dict[CaseKey, Any]]

Return configs as a nested dict (platform -> key -> config).

Source code in vllm/kernels/helion/config_manager.py
def to_dict(self) -> dict[str, dict[CaseKey, Any]]:
    """Return configs as a nested dict (platform -> key -> config)."""
    result: dict[str, dict[CaseKey, Any]] = {}
    for platform, config_dict in self._configs.items():
        result[platform] = {
            k: json.loads(v.to_json()) for k, v in config_dict.items()
        }
    return result

ConfiguredHelionKernel

A configured Helion kernel bound to a specific platform.

Source code in vllm/kernels/helion/register.py
class ConfiguredHelionKernel:
    """A configured Helion kernel bound to a specific platform."""

    def __init__(
        self,
        op_name: str,
        config_picker: ConfigPicker | None,
        raw_kernel_func: Callable,
        helion_settings: helion.Settings | None = None,
    ):
        self.op_name = op_name
        self.config_picker = config_picker
        self.raw_kernel_func = raw_kernel_func
        self.helion_settings = helion_settings
        self._decorated_kernel = self._create_decorated_kernel()

    def __call__(self, *args, **kwargs):
        return self._decorated_kernel(*args, **kwargs)

    def _create_key_computer(self):
        """
        Create a key computer function derived from the config picker.

        The returned function receives kernel arguments unpacked (*args) to match
        Helion's key signature (called as self._key_fn(*args)).
        """
        if self.config_picker is None:
            raise RuntimeError(
                f"No config picker registered for kernel '{self.op_name}'. "
                f"A config_picker must be provided to register_kernel()."
            )

        picker = self.config_picker
        all_keys = list(self.configs.keys())
        default = CaseKey.default()
        has_default = default in self.configs

        def key_computer(*args):
            selected = picker(args, all_keys)
            if selected is not None:
                return str(selected)
            if has_default:
                return str(default)
            return None

        return key_computer

    def _create_config_selector(self, key_computer):
        str_to_key = {str(k): k for k in self.configs}

        def config_selector(args):
            selected_str = key_computer(*args)

            if selected_str is None:
                raise ValueError(
                    f"Config picker returned None for kernel "
                    f"'{self.op_name}' with available config keys: "
                    f"{list(self.configs.keys())}"
                )

            config_key = str_to_key.get(selected_str)
            if config_key is None:
                raise ValueError(
                    f"Config picker returned invalid config key "
                    f"'{selected_str}' for kernel "
                    f"'{self.op_name}'. "
                    f"Available keys: {list(self.configs.keys())}"
                )

            return self.configs[config_key]

        return config_selector

    def _load_platform_configs(self) -> None:
        from vllm.kernels.helion.config_manager import ConfigManager
        from vllm.kernels.helion.utils import get_canonical_gpu_name

        self.platform = get_canonical_gpu_name()
        config_manager = ConfigManager()
        self.configs = config_manager.get_platform_configs(self.op_name, self.platform)

        if not self.configs:
            raise ValueError(
                f"No configs available for kernel '{self.op_name}' "
                f"on platform '{self.platform}'"
            )

    def _create_decorated_kernel(self) -> Callable[..., Any]:
        self._load_platform_configs()

        key_computer = self._create_key_computer()
        config_selector = self._create_config_selector(key_computer)

        extra_kwargs = {
            "autotuner_fn": lambda _, args: PresetConfigSearch(args, config_selector),
            "key": key_computer,
        }

        logger.debug(
            "Creating decorated kernel %s with custom autotuner on platform %s",
            self.op_name,
            self.platform,
        )
        return create_helion_decorated_kernel(
            self.raw_kernel_func, self.helion_settings, extra_kwargs
        )

_create_key_computer

_create_key_computer()

Create a key computer function derived from the config picker.

The returned function receives kernel arguments unpacked (args) to match Helion's key signature (called as self._key_fn(args)).

Source code in vllm/kernels/helion/register.py
def _create_key_computer(self):
    """
    Create a key computer function derived from the config picker.

    The returned function receives kernel arguments unpacked (*args) to match
    Helion's key signature (called as self._key_fn(*args)).
    """
    if self.config_picker is None:
        raise RuntimeError(
            f"No config picker registered for kernel '{self.op_name}'. "
            f"A config_picker must be provided to register_kernel()."
        )

    picker = self.config_picker
    all_keys = list(self.configs.keys())
    default = CaseKey.default()
    has_default = default in self.configs

    def key_computer(*args):
        selected = picker(args, all_keys)
        if selected is not None:
            return str(selected)
        if has_default:
            return str(default)
        return None

    return key_computer

HelionKernelWrapper

Wrapper for Helion kernels with pre-tuned config selection and HOP support.

Source code in vllm/kernels/helion/register.py
class HelionKernelWrapper:
    """Wrapper for Helion kernels with pre-tuned config selection and HOP support."""

    def __init__(
        self,
        raw_kernel_func: Callable,
        op_name: str,
        fake_impl: Callable,
        config_picker: ConfigPicker,
        helion_settings: helion.Settings | None = None,
        input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None,
    ):
        # Validate helion_settings doesn't conflict with our custom autotuner
        validate_helion_settings(helion_settings, op_name)

        self.raw_kernel_func = raw_kernel_func
        self.op_name = op_name
        self._fake_impl = fake_impl
        self.helion_settings = helion_settings
        self._config_picker = config_picker
        self._input_generator = input_generator
        self._configured_kernel: ConfiguredHelionKernel | None = None
        # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR,
        # which handles op enablement/disablement.
        self._disabled = False
        self._disabled_reason: str | None = None

        try:
            if not _HOP_AVAILABLE:
                self._get_or_register_custom_op()
            else:
                self.get_configured_op()
        except ValueError as e:
            self._disabled = True
            self._disabled_reason = str(e)
            logger.warning(
                "Helion kernel '%s' is disabled: %s",
                op_name,
                self._disabled_reason,
            )

    def __call__(self, *args, **kwargs):
        if self._disabled:
            raise RuntimeError(
                f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
            )
        if not _HOP_AVAILABLE:
            op = getattr(torch.ops.vllm_helion, self.op_name)
            return op(*args, **kwargs)
        assert self._configured_kernel is not None, (
            f"Kernel '{self.op_name}' was not initialized. "
            "Please open an issue on GitHub."
        )

        # During Dynamo tracing, this call will be intercepted by our custom
        # HelionKernelWrapperVariable and handled via proper HOP emission.
        # During eager execution, call the kernel directly.
        return self._configured_kernel(*args, **kwargs)

    def get_inputs(self) -> dict[CaseKey, tuple[Any, ...]]:
        if self._input_generator is None:
            raise NotImplementedError(
                f"No input generator registered for kernel '{self.op_name}'. "
                f"Use register_kernel(..., input_generator=...) to register one."
            )
        return self._input_generator()

    def run_autotune(
        self,
        inputs: tuple[Any, ...],
        autotune_effort: str = "quick",
    ) -> Config:
        """Run autotuning for a single input configuration."""
        extra_kwargs = {
            "autotune_effort": autotune_effort,
            "autotune_ignore_errors": True,
        }
        autotune_kernel = create_helion_decorated_kernel(
            self.raw_kernel_func, self.helion_settings, extra_kwargs
        )
        return autotune_kernel.autotune(inputs)

    def get_configured_op(self) -> ConfiguredHelionKernel:
        if self._disabled:
            raise RuntimeError(
                f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}"
            )
        if self._configured_kernel is None:
            self._configured_kernel = ConfiguredHelionKernel(
                op_name=self.op_name,
                config_picker=self._config_picker,
                raw_kernel_func=self.raw_kernel_func,
                helion_settings=self.helion_settings,
            )
        return self._configured_kernel

    def _get_or_register_custom_op(self) -> Any:
        if hasattr(torch.ops.vllm_helion, self.op_name):
            return getattr(torch.ops.vllm_helion, self.op_name)

        configured_kernel = self.get_configured_op()

        logger.info("Registering op: vllm_helion::%s", self.op_name)
        direct_register_custom_op(
            op_name=self.op_name,
            op_func=configured_kernel._decorated_kernel,
            mutates_args=None,
            fake_impl=self._fake_impl,
            target_lib=vllm_helion_lib,
        )
        return getattr(torch.ops.vllm_helion, self.op_name)

run_autotune

run_autotune(
    inputs: tuple[Any, ...], autotune_effort: str = "quick"
) -> Config

Run autotuning for a single input configuration.

Source code in vllm/kernels/helion/register.py
def run_autotune(
    self,
    inputs: tuple[Any, ...],
    autotune_effort: str = "quick",
) -> Config:
    """Run autotuning for a single input configuration."""
    extra_kwargs = {
        "autotune_effort": autotune_effort,
        "autotune_ignore_errors": True,
    }
    autotune_kernel = create_helion_decorated_kernel(
        self.raw_kernel_func, self.helion_settings, extra_kwargs
    )
    return autotune_kernel.autotune(inputs)

canonicalize_gpu_name

canonicalize_gpu_name(name: str) -> str

Canonicalize GPU name for use as a platform identifier.

Converts to lowercase, replaces spaces and hyphens with underscores, and maps known variant names to their canonical form via _GPU_NAME_ALIASES. e.g., "NVIDIA H100 80GB HBM3" -> "nvidia_h100" "NVIDIA A100-SXM4-80GB" -> "nvidia_a100" "AMD Instinct MI300X" -> "amd_instinct_mi300x"

Source code in vllm/kernels/helion/utils.py
def canonicalize_gpu_name(name: str) -> str:
    """
    Canonicalize GPU name for use as a platform identifier.

    Converts to lowercase, replaces spaces and hyphens with underscores,
    and maps known variant names to their canonical form via _GPU_NAME_ALIASES.
    e.g., "NVIDIA H100 80GB HBM3" -> "nvidia_h100"
          "NVIDIA A100-SXM4-80GB" -> "nvidia_a100"
          "AMD Instinct MI300X"   -> "amd_instinct_mi300x"
    """
    if not name or not name.strip():
        raise ValueError("GPU name cannot be empty")
    name = name.lower()
    name = name.replace(" ", "_")
    name = name.replace("-", "_")
    if name in _GPU_NAME_ALIASES:
        return _GPU_NAME_ALIASES[name]
    return name

register_kernel

register_kernel(
    op_name: str | None = None,
    *,
    config_picker: ConfigPicker,
    fake_impl: Callable | None = None,
    helion_settings: Settings | None = None,
    input_generator: Callable[
        [], dict[CaseKey, tuple[Any, ...]]
    ]
    | None = None,
) -> Callable[[Callable], HelionKernelWrapper]

Register a Helion kernel with pre-tuned config selection.

Parameters:

Name Type Description Default
config_picker ConfigPicker

Required. Receives (args, config_keys) where each config key is a dict[str, Any] mapping parameter names to values. Return the best-matching dict, or None to fall back to the default config.

Example::

def pick_config(args, config_keys):
    x = args[0]
    best = min(config_keys, key=lambda k: abs(k["size"] - x.shape[0]))
    return best
required
input_generator Callable[[], dict[CaseKey, tuple[Any, ...]]] | None

Optional. Returns dict[str, tuple] where each key is a serialized config key and each value is a tuple of arguments to pass to the kernel.

Example::

def generate_inputs():
    return {
        "4096": (torch.randn(4096, device="cuda"), 0.5),
        "8192": (torch.randn(8192, device="cuda"), 0.5),
    }
None
Source code in vllm/kernels/helion/register.py
def register_kernel(
    op_name: str | None = None,
    *,
    config_picker: ConfigPicker,
    fake_impl: Callable | None = None,
    helion_settings: helion.Settings | None = None,
    input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None,
) -> Callable[[Callable], HelionKernelWrapper]:
    """Register a Helion kernel with pre-tuned config selection.

    Args:
        config_picker: Required. Receives ``(args, config_keys)``
            where each config key is a ``dict[str, Any]`` mapping
            parameter names to values.  Return the best-matching
            dict, or ``None`` to fall back to the default config.

            Example::

                def pick_config(args, config_keys):
                    x = args[0]
                    best = min(config_keys, key=lambda k: abs(k["size"] - x.shape[0]))
                    return best

        input_generator: Optional. Returns ``dict[str, tuple]`` where
            each key is a serialized config key and each value is a
            tuple of arguments to pass to the kernel.

            Example::

                def generate_inputs():
                    return {
                        "4096": (torch.randn(4096, device="cuda"), 0.5),
                        "8192": (torch.randn(8192, device="cuda"), 0.5),
                    }
    """

    def decorator(kernel_func: Callable) -> HelionKernelWrapper:
        final_op_name = op_name if op_name else kernel_func.__name__

        if final_op_name in _REGISTERED_KERNELS:
            raise ValueError(
                f"Helion kernel '{final_op_name}' is already registered. "
                f"Use a different op_name or check for duplicate registrations."
            )

        final_fake_impl = fake_impl
        if final_fake_impl is None:
            final_fake_impl = infer_fake_impl(kernel_func, helion_settings)
            logger.debug(
                "Auto-generated fake_impl for Helion kernel '%s'",
                kernel_func.__name__,
            )

        kernel_wrapper = HelionKernelWrapper(
            raw_kernel_func=kernel_func,
            op_name=final_op_name,
            fake_impl=final_fake_impl,
            config_picker=config_picker,
            helion_settings=helion_settings,
            input_generator=input_generator,
        )

        _REGISTERED_KERNELS[final_op_name] = kernel_wrapper

        logger.info(
            "Registered Helion kernel '%s' as HelionKernelWrapper",
            kernel_func.__name__,
        )

        return kernel_wrapper

    return decorator