| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Configuration management for Helion kernels. |
| 4 | |
| 5 | This module provides centralized configuration file management for Helion custom |
| 6 | operations, including naming conventions, directory resolution, and file I/O. |
| 7 | |
| 8 | Config File Structure |
| 9 | --------------------- |
| 10 | Each kernel has a directory: {kernel_name}/ |
| 11 | Inside, each GPU platform has its own JSON file: {kernel_name}/{platform}.json |
| 12 | |
| 13 | Platform files store config entries as a JSON array:: |
| 14 | |
| 15 | [ |
| 16 | {"key": {}, "config": {...}}, |
| 17 | {"key": {"intermediate": 2048, "numtokens": 256}, "config": {...}}, |
| 18 | ..., |
| 19 | ] |
| 20 | |
| 21 | Config keys are ``CaseKey`` instances mapping parameter names to |
| 22 | values. The default config uses ``CaseKey.default()``. |
| 23 | """ |
| 24 | |
| 25 | from __future__ import annotations |
| 26 | |
| 27 | import json |
| 28 | from pathlib import Path |
| 29 | from typing import Any |
| 30 | |
| 31 | from vllm.kernels.helion.case_key import CaseKey |
| 32 | from vllm.logger import init_logger |
| 33 | from vllm.utils.import_utils import has_helion |
| 34 | |
| 35 | if not has_helion(): |
| 36 | raise ImportError( |
| 37 | "ConfigManager requires helion to be installed. " |
| 38 | "Install it with: pip install helion" |
| 39 | ) |
| 40 | |
| 41 | import helion |
| 42 | |
| 43 | logger = init_logger(__name__) |
| 44 | |
| 45 | |
| 46 | class ConfigSet: |
| 47 | """In-memory collection of Helion configs with lookup/query capabilities. |
| 48 | |
| 49 | Configs are stored keyed by ``CaseKey``. The default config |
| 50 | uses ``CaseKey.default()`` as its key. |
| 51 | """ |
| 52 | |
| 53 | _ConfigDict = dict[str, dict[CaseKey, "helion.Config"]] |
| 54 | |
| 55 | def __init__(self, kernel_name: str): |
| 56 | self._kernel_name = kernel_name |
| 57 | self._configs: ConfigSet._ConfigDict = {} |
| 58 | |
| 59 | @property |
| 60 | def kernel_name(self) -> str: |
| 61 | return self._kernel_name |
| 62 | |
| 63 | def get_config(self, platform: str, config_key: CaseKey) -> helion.Config: |
| 64 | platform_dict = self._configs.get(platform) |
| 65 | if platform_dict is None: |
| 66 | avail_platforms = self.get_platforms() |
| 67 | # TODO(@gmagogsfm): add a CLI/env override flag so users can |
| 68 | # directly specify a platform name instead of relying on |
| 69 | # auto-detection, and suggest it in this error message. |
| 70 | raise KeyError( |
| 71 | f"Config not found for kernel '{self._kernel_name}': " |
| 72 | f"platform '{platform}' not found. " |
| 73 | f"Available platforms: {avail_platforms or '(none)'}. " |
| 74 | f"If your GPU is a variant of a supported platform, " |
| 75 | f"consider adding a mapping in _GPU_NAME_ALIASES in " |
| 76 | f"vllm/kernels/helion/utils.py, or run " |
| 77 | f"tools/autotune_helion_kernels.py to generate configs " |
| 78 | f"for your platform." |
| 79 | ) |
| 80 | |
| 81 | config = platform_dict.get(config_key) |
| 82 | if config is None: |
| 83 | avail_keys = self.get_config_keys(platform) |
| 84 | raise KeyError( |
| 85 | f"Config not found for kernel '{self._kernel_name}': " |
| 86 | f"config_key '{config_key}' not found for " |
| 87 | f"platform '{platform}'. " |
| 88 | f"Available config_keys: {avail_keys or '(none)'}" |
| 89 | ) |
| 90 | |
| 91 | return config |
| 92 | |
| 93 | def get_platforms(self) -> list[str]: |
| 94 | return sorted(self._configs.keys()) |
| 95 | |
| 96 | def get_config_keys(self, platform: str) -> list[CaseKey]: |
| 97 | platform_dict = self._configs.get(platform.lower()) |
| 98 | if platform_dict is None: |
| 99 | return [] |
| 100 | return sorted(platform_dict.keys(), key=str) |
| 101 | |
| 102 | def to_config_entries(self) -> dict[str, list[dict[str, Any]]]: |
| 103 | """Serialize to config entries format for JSON output.""" |
| 104 | result: dict[str, list[dict[str, Any]]] = {} |
| 105 | for platform, config_dict in self._configs.items(): |
| 106 | pairs: list[dict[str, Any]] = [] |
| 107 | for config_key, config in config_dict.items(): |
| 108 | config_data = json.loads(config.to_json()) |
| 109 | pairs.append({"key": dict(config_key), "config": config_data}) |
| 110 | result[platform] = pairs |
| 111 | return result |
| 112 | |
| 113 | def to_dict(self) -> dict[str, dict[CaseKey, Any]]: |
| 114 | """Return configs as a nested dict (platform -> key -> config).""" |
| 115 | result: dict[str, dict[CaseKey, Any]] = {} |
| 116 | for platform, config_dict in self._configs.items(): |
| 117 | result[platform] = { |
| 118 | k: json.loads(v.to_json()) for k, v in config_dict.items() |
| 119 | } |
| 120 | return result |
| 121 | |
| 122 | @classmethod |
| 123 | def from_dict(cls, kernel_name: str, data: dict[str, Any]) -> ConfigSet: |
| 124 | config_set = cls(kernel_name) |
| 125 | count = 0 |
| 126 | |
| 127 | for platform, platform_data in data.items(): |
| 128 | if platform not in config_set._configs: |
| 129 | config_set._configs[platform] = {} |
| 130 | |
| 131 | for entry in platform_data: |
| 132 | raw_key = entry["key"] |
| 133 | key = CaseKey.default() if not raw_key else CaseKey(raw_key) |
| 134 | config = helion.Config(**entry["config"]) |
| 135 | config_set._configs[platform][key] = config |
| 136 | count += 1 |
| 137 | |
| 138 | if count > 0: |
| 139 | logger.debug( |
| 140 | "Loaded %d configs for kernel '%s'", |
| 141 | count, |
| 142 | kernel_name, |
| 143 | ) |
| 144 | |
| 145 | return config_set |
| 146 | |
| 147 | def set_config( |
| 148 | self, |
| 149 | platform: str, |
| 150 | config_key: CaseKey, |
| 151 | config: helion.Config, |
| 152 | ) -> None: |
| 153 | platform = platform.lower() |
| 154 | if platform not in self._configs: |
| 155 | self._configs[platform] = {} |
| 156 | self._configs[platform][config_key] = config |
| 157 | logger.debug( |
| 158 | "Set config for kernel '%s': platform='%s', key='%s'", |
| 159 | self._kernel_name, |
| 160 | platform, |
| 161 | config_key, |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | class ConfigManager: |
| 166 | """File-level configuration management for Helion kernels (global singleton).""" |
| 167 | |
| 168 | _instance: ConfigManager | None = None |
| 169 | _instance_base_dir: Path | None = None |
| 170 | |
| 171 | def __new__(cls, base_dir: str | Path | None = None) -> ConfigManager: |
| 172 | resolved_base_dir = cls._resolve_base_dir(base_dir) |
| 173 | |
| 174 | if cls._instance is not None: |
| 175 | if cls._instance_base_dir != resolved_base_dir: |
| 176 | raise ValueError( |
| 177 | f"ConfigManager singleton already exists with base_dir " |
| 178 | f"'{cls._instance_base_dir}', cannot create with " |
| 179 | f"different base_dir '{resolved_base_dir}'" |
| 180 | ) |
| 181 | return cls._instance |
| 182 | |
| 183 | instance = super().__new__(cls) |
| 184 | cls._instance = instance |
| 185 | cls._instance_base_dir = resolved_base_dir |
| 186 | return instance |
| 187 | |
| 188 | def __init__(self, base_dir: str | Path | None = None): |
| 189 | if hasattr(self, "_base_dir"): |
| 190 | return |
| 191 | |
| 192 | self._base_dir = self._resolve_base_dir(base_dir) |
| 193 | logger.debug("ConfigManager initialized with base_dir: %s", self._base_dir) |
| 194 | |
| 195 | @staticmethod |
| 196 | def _resolve_base_dir(base_dir: str | Path | None) -> Path: |
| 197 | if base_dir is not None: |
| 198 | return Path(base_dir).resolve() |
| 199 | return (Path(__file__).parent / "configs").resolve() |
| 200 | |
| 201 | @classmethod |
| 202 | def get_instance(cls) -> ConfigManager: |
| 203 | if cls._instance is None: |
| 204 | raise RuntimeError( |
| 205 | "ConfigManager instance has not been created. " |
| 206 | "Call ConfigManager(base_dir=...) first to initialize." |
| 207 | ) |
| 208 | return cls._instance |
| 209 | |
| 210 | @classmethod |
| 211 | def reset_instance(cls) -> None: |
| 212 | """For testing purposes only.""" |
| 213 | cls._instance = None |
| 214 | cls._instance_base_dir = None |
| 215 | |
| 216 | def get_kernel_dir(self, kernel_name: str) -> Path: |
| 217 | return self._base_dir / kernel_name |
| 218 | |
| 219 | def get_config_file_path( |
| 220 | self, kernel_name: str, platform: str | None = None |
| 221 | ) -> Path: |
| 222 | if platform is not None: |
| 223 | return self.get_kernel_dir(kernel_name) / f"{platform}.json" |
| 224 | return self.get_kernel_dir(kernel_name) |
| 225 | |
| 226 | def ensure_base_dir_exists(self) -> Path: |
| 227 | self._base_dir.mkdir(parents=True, exist_ok=True) |
| 228 | return self._base_dir |
| 229 | |
| 230 | def ensure_base_dir_writable(self) -> None: |
| 231 | self.ensure_base_dir_exists() |
| 232 | test_file = self._base_dir / ".write_test" |
| 233 | try: |
| 234 | test_file.write_text("test") |
| 235 | test_file.unlink() |
| 236 | except OSError as e: |
| 237 | raise OSError( |
| 238 | f"Config directory '{self._base_dir}' is not writable: {e}" |
| 239 | ) from e |
| 240 | |
| 241 | def _load_platform_file(self, kernel_name: str, platform: str) -> Any: |
| 242 | config_path = self.get_config_file_path(kernel_name, platform) |
| 243 | if not config_path.exists(): |
| 244 | return [] |
| 245 | try: |
| 246 | with open(config_path) as f: |
| 247 | return json.load(f) |
| 248 | except (json.JSONDecodeError, OSError) as e: |
| 249 | logger.error("Failed to load config file %s: %s", config_path, e) |
| 250 | return [] |
| 251 | |
| 252 | def load_config_set(self, kernel_name: str) -> ConfigSet: |
| 253 | kernel_dir = self.get_kernel_dir(kernel_name) |
| 254 | if not kernel_dir.is_dir(): |
| 255 | return ConfigSet.from_dict(kernel_name, {}) |
| 256 | |
| 257 | data: dict[str, Any] = {} |
| 258 | for platform_file in sorted(kernel_dir.glob("*.json")): |
| 259 | platform = platform_file.stem |
| 260 | try: |
| 261 | with open(platform_file) as f: |
| 262 | platform_data = json.load(f) |
| 263 | data[platform] = platform_data |
| 264 | except (json.JSONDecodeError, OSError) as e: |
| 265 | logger.error( |
| 266 | "Failed to load config file %s: %s", |
| 267 | platform_file, |
| 268 | e, |
| 269 | ) |
| 270 | |
| 271 | return ConfigSet.from_dict(kernel_name, data) |
| 272 | |
| 273 | def get_platform_configs( |
| 274 | self, kernel_name: str, platform: str |
| 275 | ) -> dict[CaseKey, helion.Config]: |
| 276 | platform_data = self._load_platform_file(kernel_name, platform) |
| 277 | if not platform_data: |
| 278 | return {} |
| 279 | config_set = ConfigSet.from_dict(kernel_name, {platform: platform_data}) |
| 280 | return { |
| 281 | k: config_set.get_config(platform, k) |
| 282 | for k in config_set.get_config_keys(platform) |
| 283 | } |
| 284 | |
| 285 | def save_config_set(self, config_set: ConfigSet) -> Path: |
| 286 | kernel_dir = self.get_kernel_dir(config_set.kernel_name) |
| 287 | kernel_dir.mkdir(parents=True, exist_ok=True) |
| 288 | |
| 289 | full_data = config_set.to_config_entries() |
| 290 | for platform, pairs in full_data.items(): |
| 291 | platform_path = kernel_dir / f"{platform}.json" |
| 292 | with open(platform_path, "w") as f: |
| 293 | json.dump(pairs, f, indent=2) |
| 294 | f.write("\n") |
| 295 | logger.info("Saved config to: %s", platform_path) |
| 296 | |
| 297 | return kernel_dir |
| 298 | |
| 299 | def save_configs( |
| 300 | self, |
| 301 | kernel_name: str, |
| 302 | platform: str, |
| 303 | configs: dict[CaseKey, helion.Config], |
| 304 | ) -> Path: |
| 305 | """Save configs for a kernel/platform, merging with existing.""" |
| 306 | config_set = ConfigSet.from_dict( |
| 307 | kernel_name, |
| 308 | {platform: self._load_platform_file(kernel_name, platform)}, |
| 309 | ) |
| 310 | for key, config in configs.items(): |
| 311 | config_set.set_config(platform, key, config) |
| 312 | |
| 313 | pairs = config_set.to_config_entries().get(platform, []) |
| 314 | platform_path = self.get_config_file_path(kernel_name, platform) |
| 315 | platform_path.parent.mkdir(parents=True, exist_ok=True) |
| 316 | with open(platform_path, "w") as f: |
| 317 | json.dump(pairs, f, indent=2) |
| 318 | f.write("\n") |
| 319 | |
| 320 | logger.info("Saved config to: %s", platform_path) |
| 321 | return platform_path |