vllm.kernels.helion.register ¶

class ConfiguredHelionKernel:
    """A configured Helion kernel bound to a specific platform."""

    def __init__(
        self,
        op_name: str,
        config_picker: Callable[[tuple[Any, ...], list[str]], str | None] | 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"Use @{self.op_name}.register_config_picker to register one."
            )

        # After None check, config_picker is guaranteed to be non-None
        assert self.config_picker is not None

        def key_computer(*args):
            config_keys = list(self.configs.keys())
            # Cast is safe because we checked for None above
            config_picker = cast(
                Callable[[tuple[Any, ...], list[str]], str | None], self.config_picker
            )
            selected_key = config_picker(args, config_keys)
            if selected_key:
                return selected_key
            return "default" if "default" in self.configs else None

        return key_computer

    def _create_config_selector(self, key_computer):
        def config_selector(args):
            # args is a tuple; key_computer expects unpacked args
            selected_config_key = key_computer(*args)

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

            if selected_config_key not in self.configs:
                raise ValueError(
                    f"Config picker returned invalid config key "
                    f"'{selected_config_key}' for kernel '{self.op_name}'. "
                    f"Available keys: {list(self.configs.keys())}"
                )

            return self.configs[selected_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.get_instance()
        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
        )

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,
        helion_settings: "helion.Settings | 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: (
            Callable[[tuple[Any, ...], list[str]], str | None] | None
        ) = None
        self._configured_kernel: ConfiguredHelionKernel | None = None
        self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None

    def __call__(self, *args, **kwargs):
        # CustomOp fallback: register as torch custom op for torch.compile
        # compatibility on older PyTorch lacking HOP/EffectType support
        if not _HOP_AVAILABLE:
            custom_op = self._get_or_register_custom_op()
            return custom_op(*args, **kwargs)
        # HOP tracing: record HigherOrderOp in the FX graph
        if get_proxy_mode() is not None:
            return self._call_via_hop(args, kwargs)
        # Eager: run the configured kernel directly
        return self.get_configured_op()(*args, **kwargs)

    def _call_via_hop(
        self,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> Any:
        kernel = self.get_configured_op()._decorated_kernel
        kernel_idx = helion_kernel_side_table.add_kernel(kernel)

        constant_args, tensor_args = self._partition_args(kernel, args, kwargs)

        all_named = {**constant_args, **tensor_args}
        full_args = tuple(
            all_named.get(n, p.default)
            for n, p in kernel.signature.parameters.items()  # type: ignore[attr-defined]
            if n in all_named or p.default is not p.empty
        )

        with disable_proxy_modes_tracing():
            output_spec = infer_output_spec(kernel, full_args)

        hop_result = helion_kernel_wrapper_mutation(
            kernel_idx=kernel_idx,
            constant_args=constant_args,
            tensor_args=tensor_args,
            output_spec=output_spec,
        )

        tree_spec_str = output_spec.get("tree_spec_str")
        if tree_spec_str is None:
            return None
        tree_spec = pytree.treespec_loads(tree_spec_str)

        hop_iter = iter(hop_result)
        reconstructed = []
        for spec in output_spec["leaf_specs"]:
            is_constant_scalar = spec["type"] == "scalar" and not isinstance(
                spec.get("scalar_value"), torch.SymInt
            )
            if is_constant_scalar:
                reconstructed.append(spec["scalar_value"])
            else:
                reconstructed.append(next(hop_iter))
        return pytree.tree_unflatten(reconstructed, tree_spec)

    @staticmethod
    def _partition_args(
        kernel: Any,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        constant_args: dict[str, Any] = {}
        tensor_args: dict[str, Any] = {}
        params = list(kernel.signature.parameters.keys())
        for i, val in enumerate(args):
            name = params[i]
            if isinstance(val, torch.Tensor):
                tensor_args[name] = val
            else:
                constant_args[name] = val
        for name, val in kwargs.items():
            if isinstance(val, torch.Tensor):
                tensor_args[name] = val
            else:
                constant_args[name] = val
        return constant_args, tensor_args

    def register_config_picker(
        self, picker_func: Callable[[tuple[Any, ...], list[str]], str | None]
    ) -> Callable[[tuple[Any, ...], list[str]], str | None]:
        self._config_picker = picker_func
        return picker_func

    def register_input_generator(
        self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
    ) -> Callable[[], dict[str, tuple[Any, ...]]]:
        """
        Register a function to generate inputs for autotuning and benchmarking.

        Args:
            generator_func: Function that returns dict[str, tuple] where:
                - key: Configuration identifier (e.g., "4096", "hidden_4096")
                - value: Tuple of arguments to pass to the kernel

        Returns:
            The registered function (for decorator usage)

        Example:
            @kernel_wrapper.register_input_generator
            def generate_inputs():
                return {
                    "4096": (torch.randn(4096, device="cuda"), 0.5),
                    "8192": (torch.randn(8192, device="cuda"), 0.5),
                }
        """
        self._input_generator = generator_func
        return generator_func

    def get_inputs(self) -> dict[str, tuple[Any, ...]]:
        if self._input_generator is None:
            raise NotImplementedError(
                f"No input generator registered for kernel '{self.op_name}'. "
                f"Use @{self.op_name}.register_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_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:
        assert self._config_picker is not None, (
            f"No config picker registered for kernel '{self.op_name}'. "
            f"Use @{self.op_name}.register_config_picker to register one."
        )

        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)

class PresetConfigSearch(BaseAutotuner):
    """Custom autotuner that uses a preset config selector instead of autotuning."""

    def __init__(
        self,
        args: tuple[Any, ...],
        config_selector: Callable[[tuple[Any, ...]], Config],
    ):
        self.args = args
        self.config_selector = config_selector

    def autotune(self, *, skip_cache: bool = False) -> Config:
        return self.config_selector(self.args)