| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | import tempfile |
| 6 | from collections.abc import Callable |
| 7 | from contextlib import contextmanager |
| 8 | from pathlib import Path |
| 9 | from typing import Any |
| 10 | from unittest.mock import patch |
| 11 | |
| 12 | import helion |
| 13 | |
| 14 | from vllm.kernels.helion.case_key import CaseKey |
| 15 | from vllm.kernels.helion.config_manager import ConfigManager |
| 16 | from vllm.kernels.helion.register import register_kernel |
| 17 | from vllm.kernels.helion.utils import get_canonical_gpu_name |
| 18 | |
| 19 | GPU_PLATFORM = get_canonical_gpu_name() |
| 20 | |
| 21 | DEFAULT_CONFIGS: dict[CaseKey, helion.Config] = { |
| 22 | CaseKey.default(): helion.Config(block_sizes=[32]), |
| 23 | } |
| 24 | |
| 25 | |
| 26 | @contextmanager |
| 27 | def dummy_kernel_registry( |
| 28 | configs: dict[CaseKey, helion.Config] | None = None, |
| 29 | ): |
| 30 | """Context manager providing a register function with automatic config setup. |
| 31 | |
| 32 | Yields a ``register`` callable with the same signature as |
| 33 | ``register_kernel``. Before applying the real decorator it writes a |
| 34 | config JSON for the kernel name (from ``op_name`` or ``fn.__name__``) |
| 35 | into a temporary directory backed by a fresh ``ConfigManager``. |
| 36 | """ |
| 37 | if configs is None: |
| 38 | configs = DEFAULT_CONFIGS |
| 39 | |
| 40 | def _to_config_entries(cfgs: dict) -> list[dict[str, Any]]: |
| 41 | pairs: list[dict[str, Any]] = [] |
| 42 | for k, v in cfgs.items(): |
| 43 | config_data = v.__dict__["config"] |
| 44 | pairs.append({"key": dict(k), "config": config_data}) |
| 45 | return pairs |
| 46 | |
| 47 | with tempfile.TemporaryDirectory() as tmpdir: |
| 48 | config_dir = Path(tmpdir) |
| 49 | ConfigManager.reset_instance() |
| 50 | cm = ConfigManager(base_dir=config_dir) |
| 51 | |
| 52 | with patch( |
| 53 | "vllm.kernels.helion.config_manager.ConfigManager", |
| 54 | return_value=cm, |
| 55 | ): |
| 56 | |
| 57 | def register( |
| 58 | op_name: str | None = None, |
| 59 | **kwargs, |
| 60 | ) -> Callable: |
| 61 | def decorator(fn: Callable) -> Callable: |
| 62 | name = op_name or fn.__name__ |
| 63 | kernel_dir = config_dir / name |
| 64 | kernel_dir.mkdir(parents=True, exist_ok=True) |
| 65 | (kernel_dir / f"{GPU_PLATFORM}.json").write_text( |
| 66 | json.dumps(_to_config_entries(configs)) |
| 67 | ) |
| 68 | return register_kernel(op_name, **kwargs)(fn) |
| 69 | |
| 70 | return decorator |
| 71 | |
| 72 | try: |
| 73 | yield register |
| 74 | finally: |
| 75 | ConfigManager.reset_instance() |