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:
            # Instance already exists - check for base_dir mismatch
            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

        # Create new 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):
        # Only initialize if not already initialized
        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 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

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()
            raise KeyError(
                f"Config not found for kernel '{self._kernel_name}': "
                f"platform '{platform}' not found. "
                f"Available platforms: {avail_platforms or '(none)'}"
            )

        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():
                # Convert helion.Config to dict using to_json() + json.loads()
                import json

                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

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)

        kernel_kwargs = {}
        if self.helion_settings:
            kernel_kwargs.update(self.helion_settings.to_dict())

        # Set static_shapes=False by default if user didn't explicitly set it to True
        # This is needed for dynamic batch sizes and sequence lengths in vLLM
        if kernel_kwargs.get("static_shapes") is not True:
            kernel_kwargs["static_shapes"] = False

        kernel_kwargs["autotuner_fn"] = lambda _, args: PresetConfigSearch(
            args, config_selector
        )
        kernel_kwargs["key"] = key_computer

        logger.debug(
            "Creating decorated kernel %s with custom autotuner on platform %s",
            self.op_name,
            self.platform,
        )
        return helion.kernel(**kernel_kwargs)(self.raw_kernel_func)

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

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