| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | """Autotune registered Helion kernels for optimal configurations. |
| 5 | |
| 6 | Usage: |
| 7 | # Autotune all registered kernels |
| 8 | python tools/autotune_helion_kernels.py |
| 9 | |
| 10 | # Autotune specific kernel |
| 11 | python tools/autotune_helion_kernels.py --kernels silu_mul_fp8 |
| 12 | |
| 13 | # Autotune multiple kernels |
| 14 | python tools/autotune_helion_kernels.py --kernels silu_mul_fp8 rms_norm_fp8 |
| 15 | |
| 16 | # Force re-autotuning |
| 17 | python tools/autotune_helion_kernels.py --force |
| 18 | |
| 19 | # List available kernels |
| 20 | python tools/autotune_helion_kernels.py --list |
| 21 | """ |
| 22 | |
| 23 | import argparse |
| 24 | import sys |
| 25 | import time |
| 26 | from dataclasses import dataclass |
| 27 | |
| 28 | import torch |
| 29 | from torch._subclasses.fake_tensor import FakeTensorMode |
| 30 | |
| 31 | try: |
| 32 | import helion |
| 33 | |
| 34 | from vllm.kernels.helion import ( |
| 35 | ConfigManager, |
| 36 | get_kernel_by_name, |
| 37 | get_registered_kernels, |
| 38 | ) |
| 39 | from vllm.kernels.helion.ops import import_all_kernels |
| 40 | from vllm.kernels.helion.utils import get_canonical_gpu_name |
| 41 | from vllm.logger import init_logger |
| 42 | from vllm.utils.import_utils import has_helion |
| 43 | except ImportError as e: |
| 44 | print(f"Error importing vLLM: {e}") |
| 45 | print("Please ensure vLLM is installed and in your Python path") |
| 46 | sys.exit(1) |
| 47 | |
| 48 | logger = init_logger("vllm.scripts.autotune_helion_kernels") |
| 49 | |
| 50 | |
| 51 | @dataclass |
| 52 | class AutotuneResult: |
| 53 | status: str # "success" | "partial" | "error" | "skipped" |
| 54 | successful: int |
| 55 | failed: int |
| 56 | configs: dict[str, "helion.Config"] |
| 57 | message: str = "" |
| 58 | |
| 59 | |
| 60 | def list_kernels() -> None: |
| 61 | kernels = get_registered_kernels() |
| 62 | |
| 63 | if not kernels: |
| 64 | print("No Helion kernels found in registry.") |
| 65 | return |
| 66 | |
| 67 | print("Available Helion kernels:") |
| 68 | print("=" * 50) |
| 69 | |
| 70 | for name in sorted(kernels.keys()): |
| 71 | print(f" {name}") |
| 72 | |
| 73 | print(f"\nTotal: {len(kernels)} kernels") |
| 74 | |
| 75 | |
| 76 | def check_requirements() -> bool: |
| 77 | if not torch.cuda.is_available(): |
| 78 | logger.error("CUDA is not available. Helion autotuning requires GPU.") |
| 79 | return False |
| 80 | |
| 81 | if not has_helion(): |
| 82 | logger.error("Helion is not installed. Please install Helion package.") |
| 83 | return False |
| 84 | |
| 85 | return True |
| 86 | |
| 87 | |
| 88 | def autotune_kernel( |
| 89 | kernel_name: str, |
| 90 | platform: str, |
| 91 | config_manager: ConfigManager, |
| 92 | force: bool = False, |
| 93 | autotune_effort: str = "quick", |
| 94 | ) -> AutotuneResult: |
| 95 | logger.debug( |
| 96 | "Starting autotune for kernel '%s' with effort='%s'", |
| 97 | kernel_name, |
| 98 | autotune_effort, |
| 99 | ) |
| 100 | kernel_wrapper = get_kernel_by_name(kernel_name) |
| 101 | if kernel_wrapper is None: |
| 102 | error_msg = f"Kernel '{kernel_name}' not found in registry" |
| 103 | logger.error(error_msg) |
| 104 | return AutotuneResult( |
| 105 | status="error", |
| 106 | message=error_msg, |
| 107 | successful=0, |
| 108 | failed=0, |
| 109 | configs={}, |
| 110 | ) |
| 111 | |
| 112 | try: |
| 113 | with FakeTensorMode(): |
| 114 | all_config_keys = list(kernel_wrapper.get_inputs().keys()) |
| 115 | except NotImplementedError: |
| 116 | error_msg = f"Kernel '{kernel_name}' has no input generator registered" |
| 117 | logger.error(error_msg) |
| 118 | return AutotuneResult( |
| 119 | status="error", |
| 120 | message=error_msg, |
| 121 | successful=0, |
| 122 | failed=0, |
| 123 | configs={}, |
| 124 | ) |
| 125 | |
| 126 | try: |
| 127 | logger.info( |
| 128 | "Autotuning kernel '%s' for platform '%s' with %d configs", |
| 129 | kernel_name, |
| 130 | platform, |
| 131 | len(all_config_keys), |
| 132 | ) |
| 133 | |
| 134 | if not force: |
| 135 | existing_configs = config_manager.get_platform_configs( |
| 136 | kernel_name, platform |
| 137 | ) |
| 138 | keys_to_autotune = [] |
| 139 | for config_key in all_config_keys: |
| 140 | if config_key in existing_configs: |
| 141 | logger.debug( |
| 142 | "Config '%s' already exists for platform '%s', skipping", |
| 143 | config_key, |
| 144 | platform, |
| 145 | ) |
| 146 | else: |
| 147 | keys_to_autotune.append(config_key) |
| 148 | else: |
| 149 | logger.debug("Force mode enabled, will re-autotune all configs") |
| 150 | keys_to_autotune = all_config_keys |
| 151 | |
| 152 | if not keys_to_autotune: |
| 153 | logger.info( |
| 154 | "All configs already exist for kernel '%s' on platform '%s'. " |
| 155 | "Use --force to re-autotune.", |
| 156 | kernel_name, |
| 157 | platform, |
| 158 | ) |
| 159 | return AutotuneResult( |
| 160 | status="skipped", |
| 161 | message="All configs already exist", |
| 162 | successful=0, |
| 163 | failed=0, |
| 164 | configs={}, |
| 165 | ) |
| 166 | |
| 167 | inputs_dict = kernel_wrapper.get_inputs() |
| 168 | configs_to_autotune = {k: inputs_dict[k] for k in keys_to_autotune} |
| 169 | |
| 170 | total_start_time = time.time() |
| 171 | autotuned_configs = {} |
| 172 | failed_configs = [] |
| 173 | |
| 174 | for config_key, inputs in configs_to_autotune.items(): |
| 175 | logger.info("Autotuning config: %s", config_key) |
| 176 | logger.debug( |
| 177 | "Input shapes: %s", |
| 178 | [getattr(inp, "shape", type(inp).__name__) for inp in inputs], |
| 179 | ) |
| 180 | |
| 181 | try: |
| 182 | config_start_time = time.time() |
| 183 | config = kernel_wrapper.run_autotune(inputs, autotune_effort) |
| 184 | config_duration = time.time() - config_start_time |
| 185 | |
| 186 | # Save immediately for checkpointing |
| 187 | config_manager.save_configs(kernel_name, platform, {config_key: config}) |
| 188 | |
| 189 | autotuned_configs[config_key] = config |
| 190 | logger.debug("Config details: %s", config) |
| 191 | |
| 192 | logger.info( |
| 193 | "✓ Autotuned and saved config '%s' (%.2fs)", |
| 194 | config_key, |
| 195 | config_duration, |
| 196 | ) |
| 197 | |
| 198 | except (RuntimeError, ValueError, OSError) as e: |
| 199 | logger.exception( |
| 200 | "Failed to autotune config '%s': %s", |
| 201 | config_key, |
| 202 | e, |
| 203 | ) |
| 204 | failed_configs.append(config_key) |
| 205 | |
| 206 | total_duration = time.time() - total_start_time |
| 207 | successful = len(autotuned_configs) |
| 208 | failed = len(failed_configs) |
| 209 | |
| 210 | logger.info( |
| 211 | "Completed autotuning for kernel '%s': %d successful, %d failed (%.2fs)", |
| 212 | kernel_name, |
| 213 | successful, |
| 214 | failed, |
| 215 | total_duration, |
| 216 | ) |
| 217 | |
| 218 | status = "success" if failed == 0 else "partial" |
| 219 | return AutotuneResult( |
| 220 | status=status, |
| 221 | successful=successful, |
| 222 | failed=failed, |
| 223 | configs=autotuned_configs, |
| 224 | ) |
| 225 | |
| 226 | except (KeyError, RuntimeError, ValueError, OSError) as e: |
| 227 | error_msg = f"Unexpected error: {e}" |
| 228 | logger.exception("Failed to autotune kernel '%s': %s", kernel_name, e) |
| 229 | return AutotuneResult( |
| 230 | status="error", |
| 231 | message=error_msg, |
| 232 | successful=0, |
| 233 | failed=0, |
| 234 | configs={}, |
| 235 | ) |
| 236 | |
| 237 | |
| 238 | def summarize_results(results: dict[str, AutotuneResult]) -> bool: |
| 239 | logger.info("=" * 50) |
| 240 | logger.info("Autotuning Results Summary") |
| 241 | logger.info("=" * 50) |
| 242 | |
| 243 | total_successful = 0 |
| 244 | total_failed = 0 |
| 245 | success_kernels = [] |
| 246 | partial_kernels = [] |
| 247 | error_kernels = [] |
| 248 | skipped_kernels = [] |
| 249 | |
| 250 | for kernel_name, result in results.items(): |
| 251 | total_successful += result.successful |
| 252 | total_failed += result.failed |
| 253 | |
| 254 | if result.status == "success": |
| 255 | success_kernels.append(f"{kernel_name} ({result.successful} configs)") |
| 256 | logger.info("✓ %s: %d configs successful", kernel_name, result.successful) |
| 257 | elif result.status == "partial": |
| 258 | partial_kernels.append( |
| 259 | f"{kernel_name} ({result.successful} ok, {result.failed} failed)" |
| 260 | ) |
| 261 | logger.warning( |
| 262 | "⚠ %s: %d successful, %d failed", |
| 263 | kernel_name, |
| 264 | result.successful, |
| 265 | result.failed, |
| 266 | ) |
| 267 | elif result.status == "error": |
| 268 | error_kernels.append(f"{kernel_name}: {result.message or 'Unknown error'}") |
| 269 | logger.error("✗ %s: %s", kernel_name, result.message or "Unknown error") |
| 270 | elif result.status == "skipped": |
| 271 | skipped_kernels.append(f"{kernel_name}: {result.message or 'Skipped'}") |
| 272 | logger.info("- %s: %s", kernel_name, result.message or "Skipped") |
| 273 | |
| 274 | logger.info("=" * 50) |
| 275 | logger.info( |
| 276 | "Summary: %d total configs (%d successful, %d failed)", |
| 277 | total_successful + total_failed, |
| 278 | total_successful, |
| 279 | total_failed, |
| 280 | ) |
| 281 | logger.info( |
| 282 | "Kernels: %d success, %d partial, %d error, %d skipped", |
| 283 | len(success_kernels), |
| 284 | len(partial_kernels), |
| 285 | len(error_kernels), |
| 286 | len(skipped_kernels), |
| 287 | ) |
| 288 | |
| 289 | has_failures = bool(error_kernels or partial_kernels) |
| 290 | |
| 291 | if not has_failures: |
| 292 | if total_successful > 0: |
| 293 | logger.info("All configs autotuned successfully!") |
| 294 | else: |
| 295 | logger.info("No new configs were generated (all may already exist)") |
| 296 | |
| 297 | return not has_failures |
| 298 | |
| 299 | |
| 300 | def get_kernels_to_autotune(requested_kernels: list[str] | None) -> list[str]: |
| 301 | all_kernels = get_registered_kernels() |
| 302 | if not all_kernels: |
| 303 | logger.error("No Helion kernels found in registry") |
| 304 | sys.exit(1) |
| 305 | |
| 306 | if not requested_kernels: |
| 307 | return list(all_kernels.keys()) |
| 308 | |
| 309 | if len(requested_kernels) != len(set(requested_kernels)): |
| 310 | duplicates = [ |
| 311 | k for k in set(requested_kernels) if requested_kernels.count(k) > 1 |
| 312 | ] |
| 313 | logger.error("Duplicate kernel names in --kernels flag: %s", duplicates) |
| 314 | sys.exit(1) |
| 315 | |
| 316 | kernels_to_autotune = [] |
| 317 | missing_kernels = [] |
| 318 | |
| 319 | for kernel_name in requested_kernels: |
| 320 | if kernel_name in all_kernels: |
| 321 | kernels_to_autotune.append(kernel_name) |
| 322 | else: |
| 323 | missing_kernels.append(kernel_name) |
| 324 | |
| 325 | if missing_kernels: |
| 326 | logger.error("Kernel(s) not found: %s", missing_kernels) |
| 327 | logger.error("Available kernels: %s", list(all_kernels.keys())) |
| 328 | sys.exit(1) |
| 329 | |
| 330 | return kernels_to_autotune |
| 331 | |
| 332 | |
| 333 | def main(): |
| 334 | parser = argparse.ArgumentParser( |
| 335 | description="Autotune Helion kernels", |
| 336 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 337 | epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "", |
| 338 | ) |
| 339 | |
| 340 | parser.add_argument( |
| 341 | "--kernels", |
| 342 | nargs="+", |
| 343 | help="Kernel(s) to autotune (default: all kernels)", |
| 344 | ) |
| 345 | |
| 346 | parser.add_argument( |
| 347 | "--config-dir", |
| 348 | type=str, |
| 349 | help="Config directory for config files (default: vLLM helion configs dir)", |
| 350 | ) |
| 351 | |
| 352 | parser.add_argument( |
| 353 | "--list", |
| 354 | action="store_true", |
| 355 | help="List available Helion kernels and exit", |
| 356 | ) |
| 357 | |
| 358 | parser.add_argument( |
| 359 | "--force", |
| 360 | action="store_true", |
| 361 | help=( |
| 362 | "Force re-autotuning even if configs already exist for the " |
| 363 | "platform and config keys" |
| 364 | ), |
| 365 | ) |
| 366 | |
| 367 | parser.add_argument( |
| 368 | "--autotune-effort", |
| 369 | type=str, |
| 370 | default="quick", |
| 371 | help=( |
| 372 | "Helion autotune effort level: 'quick' (smaller search) or " |
| 373 | "'full' (full search budget) (default: quick)" |
| 374 | ), |
| 375 | ) |
| 376 | |
| 377 | parser.add_argument( |
| 378 | "--verbose", |
| 379 | action="store_true", |
| 380 | help="Enable verbose logging", |
| 381 | ) |
| 382 | |
| 383 | args = parser.parse_args() |
| 384 | |
| 385 | # import all helion kernels |
| 386 | import_all_kernels() |
| 387 | |
| 388 | import logging |
| 389 | |
| 390 | if args.verbose: |
| 391 | logging.getLogger("vllm").setLevel(logging.DEBUG) |
| 392 | logger.debug("Verbose mode enabled") |
| 393 | logger.debug("Arguments: %s", vars(args)) |
| 394 | else: |
| 395 | logging.getLogger("vllm").setLevel(logging.INFO) |
| 396 | |
| 397 | if args.list: |
| 398 | list_kernels() |
| 399 | return |
| 400 | |
| 401 | if not check_requirements(): |
| 402 | sys.exit(1) |
| 403 | |
| 404 | platform = get_canonical_gpu_name() |
| 405 | logger.info("Detected GPU platform: %s", platform) |
| 406 | |
| 407 | config_manager = ( |
| 408 | ConfigManager(args.config_dir) if args.config_dir else ConfigManager() |
| 409 | ) |
| 410 | |
| 411 | try: |
| 412 | config_manager.ensure_base_dir_writable() |
| 413 | except OSError as e: |
| 414 | logger.error("Failed to access config directory: %s", e) |
| 415 | sys.exit(1) |
| 416 | |
| 417 | kernels_to_autotune = get_kernels_to_autotune(args.kernels) |
| 418 | |
| 419 | logger.info( |
| 420 | "Will autotune %d kernel(s) for platform '%s': %s", |
| 421 | len(kernels_to_autotune), |
| 422 | platform, |
| 423 | kernels_to_autotune, |
| 424 | ) |
| 425 | |
| 426 | results = {} |
| 427 | for kernel_name in kernels_to_autotune: |
| 428 | result = autotune_kernel( |
| 429 | kernel_name, platform, config_manager, args.force, args.autotune_effort |
| 430 | ) |
| 431 | results[kernel_name] = result |
| 432 | |
| 433 | success = summarize_results(results) |
| 434 | sys.exit(0 if success else 1) |
| 435 | |
| 436 | |
| 437 | if __name__ == "__main__": |
| 438 | main() |