vllm.kernels.helion ¶

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 different "
                    f"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_config_file_path(self, kernel_name: str) -> Path:
        return self._base_dir / f"{kernel_name}.json"

    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_config_set(self, kernel_name: str) -> ConfigSet:
        config_path = self.get_config_file_path(kernel_name)
        if not config_path.exists():
            return ConfigSet.from_dict(kernel_name, {})

        try:
            with open(config_path) as f:
                data = json.load(f)
            return ConfigSet.from_dict(kernel_name, data)
        except (json.JSONDecodeError, OSError) as e:
            logger.error("Failed to load config file %s: %s", config_path, e)
            return ConfigSet.from_dict(kernel_name, {})

    def get_platform_configs(
        self, kernel_name: str, platform: str
    ) -> dict[str, helion.Config]:
        config_set = self.load_config_set(kernel_name)
        config_keys = config_set.get_config_keys(platform)

        return {
            config_key: config_set.get_config(platform, config_key)
            for config_key in config_keys
        }

    def save_config_set(self, config_set: ConfigSet) -> Path:
        config_path = self.get_config_file_path(config_set.kernel_name)
        config_path.parent.mkdir(parents=True, exist_ok=True)

        with open(config_path, "w") as f:
            json.dump(config_set.to_dict(), f, indent=2)

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

    def save_configs(
        self,
        kernel_name: str,
        platform: str,
        configs: dict[str, "helion.Config"],
    ) -> Path:
        """Save configs for a kernel/platform, merging with existing."""
        config_set = self.load_config_set(kernel_name)
        for config_key, config in configs.items():
            config_set.set_config(platform, config_key, config)
        return self.save_config_set(config_set)

    def config_exists(self, kernel_name: str, platform: str, config_key: str) -> bool:
        config_set = self.load_config_set(kernel_name)
        return config_set.has_config(platform, config_key)

class ConfigSet:
    """In-memory collection of Helion configs with lookup/query capabilities."""

    # Type alias for nested config structure:
    # platform -> config_key -> helion.Config
    _ConfigDict = dict[str, dict[str, "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: str) -> 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 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[str]:
        platform_dict = self._configs.get(platform.lower())
        if platform_dict is None:
            return []
        return sorted(platform_dict.keys())

    def to_dict(self) -> dict[str, Any]:
        result: dict[str, Any] = {}

        for platform, config_keys_dict in self._configs.items():
            result[platform] = {}

            for config_key, config in config_keys_dict.items():
                result[platform][config_key] = json.loads(config.to_json())

        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 config_key, config_data in platform_data.items():
                config = helion.Config(**config_data)
                config_set._configs[platform][config_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: str, 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: str) -> bool:
        platform = platform.lower()
        platform_dict = self._configs.get(platform)
        if platform_dict is None:
            return False
        return config_key in platform_dict

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 that creates config-specific PyTorch custom ops."""

    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._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None

    def __call__(self, *args, **kwargs):
        configured_op = self.get_configured_op()
        return configured_op(*args, **kwargs)

    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) -> Any:
        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 hasattr(torch.ops.vllm_helion, self.op_name):
            logger.debug("Op vllm_helion::%s already registered", self.op_name)
            return getattr(torch.ops.vllm_helion, self.op_name)

        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,
        )

        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,  # Register decorated kernel
            # TODO(gmagogsfm): Implement automatic mutation/aliasing detection
            # for Helion kernels.
            mutates_args=None,
            fake_impl=self._fake_impl,
            target_lib=vllm_helion_lib,
        )
        return getattr(torch.ops.vllm_helion, self.op_name)