| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """vLLM Helion kernel registration with pre-tuned config selection. |
| 4 | |
| 5 | This module leverages Helion's internal config selection infrastructure to use |
| 6 | pre-tuned configs instead of runtime autotuning. |
| 7 | |
| 8 | How Helion Normally Works |
| 9 | ------------------------- |
| 10 | For each kernel invocation, Helion: |
| 11 | 1. Computes a cache key from input arguments |
| 12 | 2. Looks up the key in its internal compilation cache |
| 13 | 3. On cache miss, runs autotuning to find the best config |
| 14 | 4. Compiles and caches the kernel with that config |
| 15 | |
| 16 | How We Override It |
| 17 | ------------------ |
| 18 | We override two Helion hooks to use pre-tuned configs: |
| 19 | |
| 20 | 1. **key**: We provide a key function (derived from config_picker) that |
| 21 | computes cache keys matching our pre-tuned config keys. This ensures Helion's |
| 22 | internal cache uses keys that correspond to configs we've prepared. |
| 23 | |
| 24 | 2. **autotuner_fn**: We provide PresetConfigSearch which, instead of autotuning, |
| 25 | simply returns the pre-tuned config for the computed key. On cache miss, |
| 26 | Helion calls our autotuner which returns the author-prepared config. |
| 27 | |
| 28 | Both hooks use the same config_picker logic to ensure the cache key computed |
| 29 | by key matches the config returned by the autotuner. |
| 30 | |
| 31 | Key Classes |
| 32 | ----------- |
| 33 | - HelionKernelWrapper: Wraps raw kernel + config_picker, creates configured kernels |
| 34 | - ConfiguredHelionKernel: Platform-specific kernel with pre-tuned configs |
| 35 | - PresetConfigSearch: Custom autotuner that returns pre-tuned configs |
| 36 | """ |
| 37 | |
| 38 | from __future__ import annotations |
| 39 | |
| 40 | import inspect |
| 41 | from collections.abc import Callable |
| 42 | from typing import Any |
| 43 | |
| 44 | import torch |
| 45 | from torch.library import Library |
| 46 | |
| 47 | from vllm.kernels.helion.case_key import CaseKey |
| 48 | from vllm.logger import init_logger |
| 49 | from vllm.utils.import_utils import has_helion |
| 50 | from vllm.utils.torch_utils import direct_register_custom_op |
| 51 | |
| 52 | if not has_helion(): |
| 53 | raise ImportError( |
| 54 | "register module requires helion to be installed. " |
| 55 | "Install it with: pip install helion" |
| 56 | ) |
| 57 | |
| 58 | import helion |
| 59 | from helion.autotuner.base_search import BaseAutotuner |
| 60 | from helion.runtime.config import Config |
| 61 | from helion.runtime.settings import default_autotuner_fn |
| 62 | |
| 63 | # TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op, |
| 64 | # vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11. |
| 65 | # FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed. |
| 66 | # _HOP_AVAILABLE = requires_torch_version("2.11") |
| 67 | _HOP_AVAILABLE = False |
| 68 | |
| 69 | if _HOP_AVAILABLE: |
| 70 | from helion._compat import supports_torch_compile_fusion |
| 71 | from helion._compiler._dynamo.higher_order_ops import helion_kernel_side_table |
| 72 | from helion._compiler._dynamo.variables import HelionKernelVariable |
| 73 | from helion.runtime.kernel import Kernel |
| 74 | from torch._dynamo.guards import GuardBuilder |
| 75 | from torch._dynamo.variables.builder import VariableBuilder |
| 76 | |
| 77 | |
| 78 | logger = init_logger(__name__) |
| 79 | |
| 80 | vllm_helion_lib = Library("vllm_helion", "FRAGMENT") # noqa |
| 81 | |
| 82 | ConfigPicker = Callable[[tuple[Any, ...], list[CaseKey]], CaseKey | None] |
| 83 | |
| 84 | |
| 85 | def validate_helion_settings( |
| 86 | helion_settings: helion.Settings | None, op_name: str |
| 87 | ) -> None: |
| 88 | if helion_settings is None: |
| 89 | return |
| 90 | |
| 91 | settings_dict = helion_settings.to_dict() |
| 92 | |
| 93 | if ( |
| 94 | "autotuner_fn" in settings_dict |
| 95 | and settings_dict["autotuner_fn"] is not None |
| 96 | and settings_dict["autotuner_fn"] is not default_autotuner_fn |
| 97 | ): |
| 98 | raise ValueError( |
| 99 | f"HelionKernelWrapper for '{op_name}' uses a custom autotuner via " |
| 100 | f"config picker. Remove 'autotuner_fn' from helion_settings and use " |
| 101 | f"register_kernel(..., config_picker=...) instead." |
| 102 | ) |
| 103 | |
| 104 | if settings_dict.get("static_shapes") is True: |
| 105 | logger.warning( |
| 106 | "Kernel '%s' has static_shapes=True in helion_settings, " |
| 107 | "which will be overridden to False. vLLM requires dynamic " |
| 108 | "shapes for variable batch sizes and sequence lengths.", |
| 109 | op_name, |
| 110 | ) |
| 111 | |
| 112 | |
| 113 | def create_helion_decorated_kernel( |
| 114 | raw_kernel_func: Callable, |
| 115 | helion_settings: helion.Settings | None = None, |
| 116 | extra_kwargs: dict[str, Any] | None = None, |
| 117 | ) -> Any: |
| 118 | kernel_kwargs: dict[str, Any] = {} |
| 119 | if helion_settings: |
| 120 | kernel_kwargs.update(helion_settings.to_dict()) |
| 121 | |
| 122 | # vLLM requires dynamic shapes for variable batch sizes and sequence lengths |
| 123 | kernel_kwargs["static_shapes"] = False |
| 124 | |
| 125 | if extra_kwargs: |
| 126 | kernel_kwargs.update(extra_kwargs) |
| 127 | |
| 128 | return helion.kernel(**kernel_kwargs)(raw_kernel_func) |
| 129 | |
| 130 | |
| 131 | class PresetConfigSearch(BaseAutotuner): |
| 132 | """Custom autotuner that uses a preset config selector instead of autotuning.""" |
| 133 | |
| 134 | def __init__( |
| 135 | self, |
| 136 | args: tuple[Any, ...], |
| 137 | config_selector: Callable[[tuple[Any, ...]], Config], |
| 138 | ): |
| 139 | self.args = args |
| 140 | self.config_selector = config_selector |
| 141 | |
| 142 | def autotune(self, *, skip_cache: bool = False) -> Config: |
| 143 | return self.config_selector(self.args) |
| 144 | |
| 145 | |
| 146 | class ConfiguredHelionKernel: |
| 147 | """A configured Helion kernel bound to a specific platform.""" |
| 148 | |
| 149 | def __init__( |
| 150 | self, |
| 151 | op_name: str, |
| 152 | config_picker: ConfigPicker | None, |
| 153 | raw_kernel_func: Callable, |
| 154 | helion_settings: helion.Settings | None = None, |
| 155 | ): |
| 156 | self.op_name = op_name |
| 157 | self.config_picker = config_picker |
| 158 | self.raw_kernel_func = raw_kernel_func |
| 159 | self.helion_settings = helion_settings |
| 160 | self._decorated_kernel = self._create_decorated_kernel() |
| 161 | |
| 162 | def __call__(self, *args, **kwargs): |
| 163 | return self._decorated_kernel(*args, **kwargs) |
| 164 | |
| 165 | def _create_key_computer(self): |
| 166 | """Create a key computer function derived from the config picker. |
| 167 | |
| 168 | The returned function receives kernel arguments unpacked (*args) to match |
| 169 | Helion's key signature (called as self._key_fn(*args)). |
| 170 | """ |
| 171 | if self.config_picker is None: |
| 172 | raise RuntimeError( |
| 173 | f"No config picker registered for kernel '{self.op_name}'. " |
| 174 | f"A config_picker must be provided to register_kernel()." |
| 175 | ) |
| 176 | |
| 177 | picker = self.config_picker |
| 178 | all_keys = list(self.configs.keys()) |
| 179 | default = CaseKey.default() |
| 180 | has_default = default in self.configs |
| 181 | |
| 182 | def key_computer(*args): |
| 183 | selected = picker(args, all_keys) |
| 184 | if selected is not None: |
| 185 | return str(selected) |
| 186 | if has_default: |
| 187 | return str(default) |
| 188 | return None |
| 189 | |
| 190 | return key_computer |
| 191 | |
| 192 | def _create_config_selector(self, key_computer): |
| 193 | str_to_key = {str(k): k for k in self.configs} |
| 194 | |
| 195 | def config_selector(args): |
| 196 | selected_str = key_computer(*args) |
| 197 | |
| 198 | if selected_str is None: |
| 199 | raise ValueError( |
| 200 | f"Config picker returned None for kernel " |
| 201 | f"'{self.op_name}' with available config keys: " |
| 202 | f"{list(self.configs.keys())}" |
| 203 | ) |
| 204 | |
| 205 | config_key = str_to_key.get(selected_str) |
| 206 | if config_key is None: |
| 207 | raise ValueError( |
| 208 | f"Config picker returned invalid config key " |
| 209 | f"'{selected_str}' for kernel " |
| 210 | f"'{self.op_name}'. " |
| 211 | f"Available keys: {list(self.configs.keys())}" |
| 212 | ) |
| 213 | |
| 214 | return self.configs[config_key] |
| 215 | |
| 216 | return config_selector |
| 217 | |
| 218 | def _load_platform_configs(self) -> None: |
| 219 | from vllm.kernels.helion.config_manager import ConfigManager |
| 220 | from vllm.kernels.helion.utils import get_canonical_gpu_name |
| 221 | |
| 222 | self.platform = get_canonical_gpu_name() |
| 223 | config_manager = ConfigManager() |
| 224 | self.configs = config_manager.get_platform_configs(self.op_name, self.platform) |
| 225 | |
| 226 | if not self.configs: |
| 227 | raise ValueError( |
| 228 | f"No configs available for kernel '{self.op_name}' " |
| 229 | f"on platform '{self.platform}'" |
| 230 | ) |
| 231 | |
| 232 | def _create_decorated_kernel(self) -> Callable[..., Any]: |
| 233 | self._load_platform_configs() |
| 234 | |
| 235 | key_computer = self._create_key_computer() |
| 236 | config_selector = self._create_config_selector(key_computer) |
| 237 | |
| 238 | extra_kwargs = { |
| 239 | "autotuner_fn": lambda _, args: PresetConfigSearch(args, config_selector), |
| 240 | "key": key_computer, |
| 241 | } |
| 242 | |
| 243 | logger.debug( |
| 244 | "Creating decorated kernel %s with custom autotuner on platform %s", |
| 245 | self.op_name, |
| 246 | self.platform, |
| 247 | ) |
| 248 | return create_helion_decorated_kernel( |
| 249 | self.raw_kernel_func, self.helion_settings, extra_kwargs |
| 250 | ) |
| 251 | |
| 252 | |
| 253 | class HelionKernelWrapper: |
| 254 | """Wrapper for Helion kernels with pre-tuned config selection and HOP support.""" |
| 255 | |
| 256 | def __init__( |
| 257 | self, |
| 258 | raw_kernel_func: Callable, |
| 259 | op_name: str, |
| 260 | fake_impl: Callable | None, |
| 261 | config_picker: ConfigPicker, |
| 262 | mutates_args: list[str] | None = None, |
| 263 | helion_settings: helion.Settings | None = None, |
| 264 | input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, |
| 265 | ): |
| 266 | # Validate helion_settings doesn't conflict with our custom autotuner |
| 267 | validate_helion_settings(helion_settings, op_name) |
| 268 | |
| 269 | self.raw_kernel_func = raw_kernel_func |
| 270 | self.op_name = op_name |
| 271 | self._fake_impl = fake_impl |
| 272 | self.helion_settings = helion_settings |
| 273 | self._config_picker = config_picker |
| 274 | self._input_generator = input_generator |
| 275 | self._mutates_args = mutates_args |
| 276 | self._configured_kernel: ConfiguredHelionKernel | None = None |
| 277 | # TODO(@gmagogsfm): Remove this disable flag once integrated with vLLM IR, |
| 278 | # which handles op enablement/disablement. |
| 279 | self._disabled = False |
| 280 | self._disabled_reason: str | None = None |
| 281 | |
| 282 | try: |
| 283 | if not _HOP_AVAILABLE: |
| 284 | self._get_or_register_custom_op() |
| 285 | else: |
| 286 | self.get_configured_op() |
| 287 | except ValueError as e: |
| 288 | self._disabled = True |
| 289 | self._disabled_reason = str(e) |
| 290 | logger.warning( |
| 291 | "Helion kernel '%s' is disabled: %s", |
| 292 | op_name, |
| 293 | self._disabled_reason, |
| 294 | ) |
| 295 | |
| 296 | def __call__(self, *args, **kwargs): |
| 297 | if self._disabled: |
| 298 | raise RuntimeError( |
| 299 | f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}" |
| 300 | ) |
| 301 | if not _HOP_AVAILABLE: |
| 302 | op = getattr(torch.ops.vllm_helion, self.op_name) |
| 303 | return op(*args, **kwargs) |
| 304 | assert self._configured_kernel is not None, ( |
| 305 | f"Kernel '{self.op_name}' was not initialized. " |
| 306 | "Please open an issue on GitHub." |
| 307 | ) |
| 308 | |
| 309 | # During Dynamo tracing, this call will be intercepted by our custom |
| 310 | # HelionKernelWrapperVariable and handled via proper HOP emission. |
| 311 | # During eager execution, call the kernel directly. |
| 312 | return self._configured_kernel(*args, **kwargs) |
| 313 | |
| 314 | def get_inputs(self) -> dict[CaseKey, tuple[Any, ...]]: |
| 315 | if self._input_generator is None: |
| 316 | raise NotImplementedError( |
| 317 | f"No input generator registered for kernel '{self.op_name}'. " |
| 318 | f"Use register_kernel(..., input_generator=...) to register one." |
| 319 | ) |
| 320 | return self._input_generator() |
| 321 | |
| 322 | def run_autotune( |
| 323 | self, |
| 324 | inputs: tuple[Any, ...], |
| 325 | autotune_effort: str = "quick", |
| 326 | ) -> Config: |
| 327 | """Run autotuning for a single input configuration.""" |
| 328 | extra_kwargs = { |
| 329 | "autotune_effort": autotune_effort, |
| 330 | "autotune_ignore_errors": True, |
| 331 | } |
| 332 | autotune_kernel = create_helion_decorated_kernel( |
| 333 | self.raw_kernel_func, self.helion_settings, extra_kwargs |
| 334 | ) |
| 335 | return autotune_kernel.autotune(inputs) |
| 336 | |
| 337 | def get_configured_op(self) -> ConfiguredHelionKernel: |
| 338 | if self._disabled: |
| 339 | raise RuntimeError( |
| 340 | f"Helion kernel '{self.op_name}' is disabled: {self._disabled_reason}" |
| 341 | ) |
| 342 | if self._configured_kernel is None: |
| 343 | self._configured_kernel = ConfiguredHelionKernel( |
| 344 | op_name=self.op_name, |
| 345 | config_picker=self._config_picker, |
| 346 | raw_kernel_func=self.raw_kernel_func, |
| 347 | helion_settings=self.helion_settings, |
| 348 | ) |
| 349 | return self._configured_kernel |
| 350 | |
| 351 | def _get_or_register_custom_op(self) -> Any: |
| 352 | if hasattr(torch.ops.vllm_helion, self.op_name): |
| 353 | return getattr(torch.ops.vllm_helion, self.op_name) |
| 354 | |
| 355 | configured_kernel = self.get_configured_op() |
| 356 | |
| 357 | logger.info("Registering op: vllm_helion::%s", self.op_name) |
| 358 | direct_register_custom_op( |
| 359 | op_name=self.op_name, |
| 360 | op_func=configured_kernel._decorated_kernel, |
| 361 | mutates_args=self._mutates_args, |
| 362 | fake_impl=self._fake_impl, |
| 363 | target_lib=vllm_helion_lib, |
| 364 | ) |
| 365 | return getattr(torch.ops.vllm_helion, self.op_name) |
| 366 | |
| 367 | |
| 368 | # Global registry for tracking all registered HelionKernelWrapper instances |
| 369 | _REGISTERED_KERNELS: dict[str, HelionKernelWrapper] = {} |
| 370 | |
| 371 | |
| 372 | def get_registered_kernels() -> dict[str, HelionKernelWrapper]: |
| 373 | return _REGISTERED_KERNELS.copy() |
| 374 | |
| 375 | |
| 376 | def get_kernel_by_name(kernel_name: str) -> HelionKernelWrapper | None: |
| 377 | return _REGISTERED_KERNELS.get(kernel_name) |
| 378 | |
| 379 | |
| 380 | def infer_fake_impl( |
| 381 | kernel_func: Callable, |
| 382 | helion_settings: helion.Settings | None = None, |
| 383 | ) -> Callable: |
| 384 | def helion_fake_kernel(*args, **kwargs): |
| 385 | kernel_kwargs = {} |
| 386 | if helion_settings: |
| 387 | kernel_kwargs.update(helion_settings.to_dict()) |
| 388 | |
| 389 | temp_decorated_kernel = helion.kernel(**kernel_kwargs)(kernel_func) |
| 390 | |
| 391 | # Bind with args to get config_spec, then get a valid default config |
| 392 | bound = temp_decorated_kernel.bind(args) |
| 393 | default_config = bound.config_spec.default_config() |
| 394 | compiled_runner = bound.compile_config(default_config) |
| 395 | |
| 396 | return compiled_runner(*args, **kwargs, _launcher=lambda *a, **kw: None) |
| 397 | |
| 398 | return helion_fake_kernel |
| 399 | |
| 400 | |
| 401 | def register_kernel( |
| 402 | op_name: str | None = None, |
| 403 | *, |
| 404 | config_picker: ConfigPicker, |
| 405 | fake_impl: Callable | None = None, |
| 406 | mutates_args: list[str] | None = None, |
| 407 | helion_settings: helion.Settings | None = None, |
| 408 | input_generator: (Callable[[], dict[CaseKey, tuple[Any, ...]]] | None) = None, |
| 409 | ) -> Callable[[Callable], HelionKernelWrapper]: |
| 410 | """Register a Helion kernel with pre-tuned config selection. |
| 411 | |
| 412 | Args: |
| 413 | op_name: Name to register the op under. Defaults to the function name. |
| 414 | fake_impl: Optional meta/fake implementation for tracing. |
| 415 | mutates_args: Names of arguments the kernel mutates in place. |
| 416 | helion_settings: Optional Helion settings override. |
| 417 | config_picker: Required. Receives ``(args, config_keys)`` |
| 418 | where each config key is a ``dict[str, Any]`` mapping |
| 419 | parameter names to values. Return the best-matching |
| 420 | dict, or ``None`` to fall back to the default config. |
| 421 | |
| 422 | Example:: |
| 423 | |
| 424 | def pick_config(args, config_keys): |
| 425 | x = args[0] |
| 426 | best = min(config_keys, key=lambda k: abs(k["size"] - x.shape[0])) |
| 427 | return best |
| 428 | |
| 429 | input_generator: Optional. Returns ``dict[str, tuple]`` where |
| 430 | each key is a serialized config key and each value is a |
| 431 | tuple of arguments to pass to the kernel. |
| 432 | |
| 433 | Example:: |
| 434 | |
| 435 | def generate_inputs(): |
| 436 | return { |
| 437 | "4096": (torch.randn(4096, device="cuda"), 0.5), |
| 438 | "8192": (torch.randn(8192, device="cuda"), 0.5), |
| 439 | } |
| 440 | |
| 441 | """ |
| 442 | |
| 443 | def decorator(kernel_func: Callable) -> HelionKernelWrapper: |
| 444 | final_op_name = op_name if op_name else kernel_func.__name__ |
| 445 | |
| 446 | if final_op_name in _REGISTERED_KERNELS: |
| 447 | raise ValueError( |
| 448 | f"Helion kernel '{final_op_name}' is already registered. " |
| 449 | f"Use a different op_name or check for duplicate registrations." |
| 450 | ) |
| 451 | |
| 452 | # A void-returning kernel that mutates its args needs no fake impl: |
| 453 | # PyTorch auto-generates a trivial one for mutable, no-return schemas. |
| 454 | returns_none = inspect.signature(kernel_func).return_annotation is None |
| 455 | |
| 456 | final_fake_impl = fake_impl |
| 457 | if final_fake_impl is None and not (mutates_args and returns_none): |
| 458 | final_fake_impl = infer_fake_impl(kernel_func, helion_settings) |
| 459 | logger.debug( |
| 460 | "Auto-generated fake_impl for Helion kernel '%s'", |
| 461 | kernel_func.__name__, |
| 462 | ) |
| 463 | |
| 464 | kernel_wrapper = HelionKernelWrapper( |
| 465 | raw_kernel_func=kernel_func, |
| 466 | op_name=final_op_name, |
| 467 | fake_impl=final_fake_impl, |
| 468 | config_picker=config_picker, |
| 469 | mutates_args=mutates_args, |
| 470 | helion_settings=helion_settings, |
| 471 | input_generator=input_generator, |
| 472 | ) |
| 473 | |
| 474 | _REGISTERED_KERNELS[final_op_name] = kernel_wrapper |
| 475 | |
| 476 | logger.info( |
| 477 | "Registered Helion kernel '%s' as HelionKernelWrapper", |
| 478 | kernel_func.__name__, |
| 479 | ) |
| 480 | |
| 481 | return kernel_wrapper |
| 482 | |
| 483 | return decorator |
| 484 | |
| 485 | |
| 486 | # Register HelionKernelWrapper with Dynamo's variable tracker system |
| 487 | if _HOP_AVAILABLE: |
| 488 | |
| 489 | def _register_vllm_helion_dynamo_variable(): |
| 490 | """Register HelionKernelWrapper with Dynamo's VariableBuilder. |
| 491 | |
| 492 | When Dynamo encounters a HelionKernelWrapper during tracing, this |
| 493 | extracts the underlying Helion Kernel and delegates to Helion's own |
| 494 | registered Kernel handler, which handles HOP emission, side table |
| 495 | registration, and inductor lowering setup. |
| 496 | """ |
| 497 | |
| 498 | def wrap_helion_kernel_wrapper( |
| 499 | builder: VariableBuilder, value: HelionKernelWrapper |
| 500 | ): |
| 501 | kernel = value.get_configured_op()._decorated_kernel |
| 502 | if supports_torch_compile_fusion(): |
| 503 | helion_handler = VariableBuilder._type_dispatch()[Kernel] |
| 504 | return helion_handler(builder, kernel) |
| 505 | kernel_idx = helion_kernel_side_table.add_kernel(kernel) |
| 506 | builder.install_guards(GuardBuilder.ID_MATCH) |
| 507 | return HelionKernelVariable(kernel, kernel_idx, source=builder.source) |
| 508 | |
| 509 | dispatch = VariableBuilder._type_dispatch() |
| 510 | dispatch[HelionKernelWrapper] = wrap_helion_kernel_wrapper |
| 511 | |
| 512 | # Register immediately when the module is imported |
| 513 | _register_vllm_helion_dynamo_variable() |