A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 08:45 UTC
[Kernel] [Helion] [5/N] Add Helion Autotuning infrastructure (#34025)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com>
Yanan Cao committed
Feb 12, 2026 at 18:21 UTC
de13dd781f1bb18fb5bbaf4535389053d98780f8
3 files changed
+551
-18
scripts/autotune_helion_kernels.py
new
+430
@@ -0,0 +1,430 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""
5
+Autotune registered Helion kernels for optimal configurations.
6
+
7
+Usage:
8
+ # Autotune all registered kernels
9
+ python scripts/autotune_helion_kernels.py
10
+
11
+ # Autotune specific kernel
12
+ python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8
13
+
14
+ # Autotune multiple kernels
15
+ python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8 rms_norm_fp8
16
+
17
+ # Force re-autotuning
18
+ python scripts/autotune_helion_kernels.py --force
19
+
20
+ # List available kernels
21
+ python scripts/autotune_helion_kernels.py --list
22
+"""
23
+
24
+import argparse
25
+import sys
26
+import time
27
+from dataclasses import dataclass
28
+
29
+import torch
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.utils import get_canonical_gpu_name
40
+ from vllm.logger import init_logger
41
+ from vllm.utils.import_utils import has_helion
42
+except ImportError as e:
43
+ print(f"Error importing vLLM: {e}")
44
+ print("Please ensure vLLM is installed and in your Python path")
45
+ sys.exit(1)
46
+
47
+logger = init_logger("vllm.scripts.autotune_helion_kernels")
48
+
49
+
50
+@dataclass
51
+class AutotuneResult:
52
+ status: str # "success" | "partial" | "error" | "skipped"
53
+ successful: int
54
+ failed: int
55
+ configs: dict[str, "helion.Config"]
56
+ message: str = ""
57
+
58
+
59
+def list_kernels() -> None:
60
+ kernels = get_registered_kernels()
61
+
62
+ if not kernels:
63
+ print("No Helion kernels found in registry.")
64
+ return
65
+
66
+ print("Available Helion kernels:")
67
+ print("=" * 50)
68
+
69
+ for name in sorted(kernels.keys()):
70
+ print(f" {name}")
71
+
72
+ print(f"\nTotal: {len(kernels)} kernels")
73
+
74
+
75
+def check_requirements() -> bool:
76
+ if not torch.cuda.is_available():
77
+ logger.error("CUDA is not available. Helion autotuning requires GPU.")
78
+ return False
79
+
80
+ if not has_helion():
81
+ logger.error("Helion is not installed. Please install Helion package.")
82
+ return False
83
+
84
+ return True
85
+
86
+
87
+def autotune_kernel(
88
+ kernel_name: str,
89
+ platform: str,
90
+ config_manager: ConfigManager,
91
+ force: bool = False,
92
+ autotune_effort: str = "quick",
93
+) -> AutotuneResult:
94
+ logger.debug(
95
+ "Starting autotune for kernel '%s' with effort='%s'",
96
+ kernel_name,
97
+ autotune_effort,
98
+ )
99
+ kernel_wrapper = get_kernel_by_name(kernel_name)
100
+ if kernel_wrapper is None:
101
+ error_msg = f"Kernel '{kernel_name}' not found in registry"
102
+ logger.error(error_msg)
103
+ return AutotuneResult(
104
+ status="error",
105
+ message=error_msg,
106
+ successful=0,
107
+ failed=0,
108
+ configs={},
109
+ )
110
+
111
+ try:
112
+ inputs_dict = kernel_wrapper.get_inputs()
113
+ except NotImplementedError:
114
+ error_msg = f"Kernel '{kernel_name}' has no input generator registered"
115
+ logger.error(error_msg)
116
+ return AutotuneResult(
117
+ status="error",
118
+ message=error_msg,
119
+ successful=0,
120
+ failed=0,
121
+ configs={},
122
+ )
123
+
124
+ try:
125
+ logger.info(
126
+ "Autotuning kernel '%s' for platform '%s' with %d configs",
127
+ kernel_name,
128
+ platform,
129
+ len(inputs_dict),
130
+ )
131
+
132
+ configs_to_autotune = {}
133
+ if not force:
134
+ existing_configs = config_manager.get_platform_configs(
135
+ kernel_name, platform
136
+ )
137
+ for config_key, inputs in inputs_dict.items():
138
+ if config_key in existing_configs:
139
+ logger.debug(
140
+ "Config '%s' already exists for platform '%s', skipping",
141
+ config_key,
142
+ platform,
143
+ )
144
+ else:
145
+ configs_to_autotune[config_key] = inputs
146
+ else:
147
+ logger.debug("Force mode enabled, will re-autotune all configs")
148
+ configs_to_autotune = inputs_dict
149
+
150
+ if not configs_to_autotune:
151
+ logger.info(
152
+ "All configs already exist for kernel '%s' on platform '%s'. "
153
+ "Use --force to re-autotune.",
154
+ kernel_name,
155
+ platform,
156
+ )
157
+ return AutotuneResult(
158
+ status="skipped",
159
+ message="All configs already exist",
160
+ successful=0,
161
+ failed=0,
162
+ configs={},
163
+ )
164
+
165
+ total_start_time = time.time()
166
+ autotuned_configs = {}
167
+ failed_configs = []
168
+
169
+ for config_key, inputs in configs_to_autotune.items():
170
+ logger.info("Autotuning config: %s", config_key)
171
+ logger.debug(
172
+ "Input shapes: %s",
173
+ [getattr(inp, "shape", type(inp).__name__) for inp in inputs],
174
+ )
175
+
176
+ try:
177
+ config_start_time = time.time()
178
+ config = kernel_wrapper.run_autotune(inputs, autotune_effort)
179
+ config_duration = time.time() - config_start_time
180
+
181
+ # Save immediately for checkpointing
182
+ config_manager.save_configs(kernel_name, platform, {config_key: config})
183
+
184
+ autotuned_configs[config_key] = config
185
+ logger.debug("Config details: %s", config)
186
+
187
+ logger.info(
188
+ "✓ Autotuned and saved config '%s' (%.2fs)",
189
+ config_key,
190
+ config_duration,
191
+ )
192
+
193
+ except (RuntimeError, ValueError, OSError) as e:
194
+ logger.exception(
195
+ "Failed to autotune config '%s': %s",
196
+ config_key,
197
+ e,
198
+ )
199
+ failed_configs.append(config_key)
200
+
201
+ total_duration = time.time() - total_start_time
202
+ successful = len(autotuned_configs)
203
+ failed = len(failed_configs)
204
+
205
+ logger.info(
206
+ "Completed autotuning for kernel '%s': %d successful, %d failed (%.2fs)",
207
+ kernel_name,
208
+ successful,
209
+ failed,
210
+ total_duration,
211
+ )
212
+
213
+ status = "success" if failed == 0 else "partial"
214
+ return AutotuneResult(
215
+ status=status,
216
+ successful=successful,
217
+ failed=failed,
218
+ configs=autotuned_configs,
219
+ )
220
+
221
+ except (KeyError, RuntimeError, ValueError, OSError) as e:
222
+ error_msg = f"Unexpected error: {e}"
223
+ logger.exception("Failed to autotune kernel '%s': %s", kernel_name, e)
224
+ return AutotuneResult(
225
+ status="error",
226
+ message=error_msg,
227
+ successful=0,
228
+ failed=0,
229
+ configs={},
230
+ )
231
+
232
+
233
+def summarize_results(results: dict[str, AutotuneResult]) -> bool:
234
+ logger.info("=" * 50)
235
+ logger.info("Autotuning Results Summary")
236
+ logger.info("=" * 50)
237
+
238
+ total_successful = 0
239
+ total_failed = 0
240
+ success_kernels = []
241
+ partial_kernels = []
242
+ error_kernels = []
243
+ skipped_kernels = []
244
+
245
+ for kernel_name, result in results.items():
246
+ total_successful += result.successful
247
+ total_failed += result.failed
248
+
249
+ if result.status == "success":
250
+ success_kernels.append(f"{kernel_name} ({result.successful} configs)")
251
+ logger.info("✓ %s: %d configs successful", kernel_name, result.successful)
252
+ elif result.status == "partial":
253
+ partial_kernels.append(
254
+ f"{kernel_name} ({result.successful} ok, {result.failed} failed)"
255
+ )
256
+ logger.warning(
257
+ "⚠ %s: %d successful, %d failed",
258
+ kernel_name,
259
+ result.successful,
260
+ result.failed,
261
+ )
262
+ elif result.status == "error":
263
+ error_kernels.append(f"{kernel_name}: {result.message or 'Unknown error'}")
264
+ logger.error("✗ %s: %s", kernel_name, result.message or "Unknown error")
265
+ elif result.status == "skipped":
266
+ skipped_kernels.append(f"{kernel_name}: {result.message or 'Skipped'}")
267
+ logger.info("- %s: %s", kernel_name, result.message or "Skipped")
268
+
269
+ logger.info("=" * 50)
270
+ logger.info(
271
+ "Summary: %d total configs (%d successful, %d failed)",
272
+ total_successful + total_failed,
273
+ total_successful,
274
+ total_failed,
275
+ )
276
+ logger.info(
277
+ "Kernels: %d success, %d partial, %d error, %d skipped",
278
+ len(success_kernels),
279
+ len(partial_kernels),
280
+ len(error_kernels),
281
+ len(skipped_kernels),
282
+ )
283
+
284
+ has_failures = bool(error_kernels or partial_kernels)
285
+
286
+ if not has_failures:
287
+ if total_successful > 0:
288
+ logger.info("All configs autotuned successfully!")
289
+ else:
290
+ logger.info("No new configs were generated (all may already exist)")
291
+
292
+ return not has_failures
293
+
294
+
295
+def get_kernels_to_autotune(requested_kernels: list[str] | None) -> list[str]:
296
+ all_kernels = get_registered_kernels()
297
+ if not all_kernels:
298
+ logger.error("No Helion kernels found in registry")
299
+ sys.exit(1)
300
+
301
+ if not requested_kernels:
302
+ return list(all_kernels.keys())
303
+
304
+ if len(requested_kernels) != len(set(requested_kernels)):
305
+ duplicates = [
306
+ k for k in set(requested_kernels) if requested_kernels.count(k) > 1
307
+ ]
308
+ logger.error("Duplicate kernel names in --kernels flag: %s", duplicates)
309
+ sys.exit(1)
310
+
311
+ kernels_to_autotune = []
312
+ missing_kernels = []
313
+
314
+ for kernel_name in requested_kernels:
315
+ if kernel_name in all_kernels:
316
+ kernels_to_autotune.append(kernel_name)
317
+ else:
318
+ missing_kernels.append(kernel_name)
319
+
320
+ if missing_kernels:
321
+ logger.error("Kernel(s) not found: %s", missing_kernels)
322
+ logger.error("Available kernels: %s", list(all_kernels.keys()))
323
+ sys.exit(1)
324
+
325
+ return kernels_to_autotune
326
+
327
+
328
+def main():
329
+ parser = argparse.ArgumentParser(
330
+ description="Autotune Helion kernels",
331
+ formatter_class=argparse.RawDescriptionHelpFormatter,
332
+ epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "",
333
+ )
334
+
335
+ parser.add_argument(
336
+ "--kernels",
337
+ nargs="+",
338
+ help="Kernel(s) to autotune (default: all kernels)",
339
+ )
340
+
341
+ parser.add_argument(
342
+ "--config-dir",
343
+ type=str,
344
+ help="Config directory for config files (default: vLLM helion configs dir)",
345
+ )
346
+
347
+ parser.add_argument(
348
+ "--list",
349
+ action="store_true",
350
+ help="List available Helion kernels and exit",
351
+ )
352
+
353
+ parser.add_argument(
354
+ "--force",
355
+ action="store_true",
356
+ help=(
357
+ "Force re-autotuning even if configs already exist for the "
358
+ "platform and config keys"
359
+ ),
360
+ )
361
+
362
+ parser.add_argument(
363
+ "--autotune-effort",
364
+ type=str,
365
+ default="quick",
366
+ help=(
367
+ "Helion autotune effort level: 'quick' (smaller search) or "
368
+ "'full' (full search budget) (default: quick)"
369
+ ),
370
+ )
371
+
372
+ parser.add_argument(
373
+ "--verbose",
374
+ action="store_true",
375
+ help="Enable verbose logging",
376
+ )
377
+
378
+ args = parser.parse_args()
379
+
380
+ import logging
381
+
382
+ if args.verbose:
383
+ logging.getLogger("vllm").setLevel(logging.DEBUG)
384
+ logger.debug("Verbose mode enabled")
385
+ logger.debug("Arguments: %s", vars(args))
386
+ else:
387
+ logging.getLogger("vllm").setLevel(logging.INFO)
388
+
389
+ if args.list:
390
+ list_kernels()
391
+ return
392
+
393
+ if not check_requirements():
394
+ sys.exit(1)
395
+
396
+ platform = get_canonical_gpu_name()
397
+ logger.info("Detected GPU platform: %s", platform)
398
+
399
+ config_manager = (
400
+ ConfigManager(args.config_dir) if args.config_dir else ConfigManager()
401
+ )
402
+
403
+ try:
404
+ config_manager.ensure_base_dir_writable()
405
+ except OSError as e:
406
+ logger.error("Failed to access config directory: %s", e)
407
+ sys.exit(1)
408
+
409
+ kernels_to_autotune = get_kernels_to_autotune(args.kernels)
410
+
411
+ logger.info(
412
+ "Will autotune %d kernel(s) for platform '%s': %s",
413
+ len(kernels_to_autotune),
414
+ platform,
415
+ kernels_to_autotune,
416
+ )
417
+
418
+ results = {}
419
+ for kernel_name in kernels_to_autotune:
420
+ result = autotune_kernel(
421
+ kernel_name, platform, config_manager, args.force, args.autotune_effort
422
+ )
423
+ results[kernel_name] = result
424
+
425
+ success = summarize_results(results)
426
+ sys.exit(0 if success else 1)
427
+
428
+
429
+if __name__ == "__main__":
430
+ main()
vllm/kernels/helion/config_manager.py
+48
-3
@@ -131,6 +131,27 @@ class ConfigSet:
131
132
return config_set
133
134
+ def set_config(
135
+ self, platform: str, config_key: str, config: "helion.Config"
136
+ ) -> None:
137
+ platform = platform.lower()
138
+ if platform not in self._configs:
139
+ self._configs[platform] = {}
140
+ self._configs[platform][config_key] = config
141
+ logger.debug(
142
+ "Set config for kernel '%s': platform='%s', key='%s'",
143
+ self._kernel_name,
144
+ platform,
145
+ config_key,
146
+ )
147
+
148
+ def has_config(self, platform: str, config_key: str) -> bool:
149
+ platform = platform.lower()
150
+ platform_dict = self._configs.get(platform)
151
+ if platform_dict is None:
152
+ return False
153
+ return config_key in platform_dict
154
+
155
156
class ConfigManager:
157
"""File-level configuration management for Helion kernels (global singleton)."""
@@ -142,7 +163,6 @@ class ConfigManager:
163
resolved_base_dir = cls._resolve_base_dir(base_dir)
164
165
if cls._instance is not None:
145
- # Instance already exists - check for base_dir mismatch
166
if cls._instance_base_dir != resolved_base_dir:
167
raise ValueError(
168
f"ConfigManager singleton already exists with base_dir "
@@ -151,14 +171,12 @@ class ConfigManager:
171
)
172
return cls._instance
173
154
- # Create new instance
174
instance = super().__new__(cls)
175
cls._instance = instance
176
cls._instance_base_dir = resolved_base_dir
177
return instance
178
179
def __init__(self, base_dir: str | Path | None = None):
161
- # Only initialize if not already initialized
180
if hasattr(self, "_base_dir"):
181
return
182
@@ -193,6 +211,17 @@ class ConfigManager:
211
self._base_dir.mkdir(parents=True, exist_ok=True)
212
return self._base_dir
213
214
+ def ensure_base_dir_writable(self) -> None:
215
+ self.ensure_base_dir_exists()
216
+ test_file = self._base_dir / ".write_test"
217
+ try:
218
+ test_file.write_text("test")
219
+ test_file.unlink()
220
+ except OSError as e:
221
+ raise OSError(
222
+ f"Config directory '{self._base_dir}' is not writable: {e}"
223
+ ) from e
224
+
225
def load_config_set(self, kernel_name: str) -> ConfigSet:
226
config_path = self.get_config_file_path(kernel_name)
227
if not config_path.exists():
@@ -226,3 +255,19 @@ class ConfigManager:
255
256
logger.info("Saved config to: %s", config_path)
257
return config_path
258
+
259
+ def save_configs(
260
+ self,
261
+ kernel_name: str,
262
+ platform: str,
263
+ configs: dict[str, "helion.Config"],
264
+ ) -> Path:
265
+ """Save configs for a kernel/platform, merging with existing."""
266
+ config_set = self.load_config_set(kernel_name)
267
+ for config_key, config in configs.items():
268
+ config_set.set_config(platform, config_key, config)
269
+ return self.save_config_set(config_set)
270
+
271
+ def config_exists(self, kernel_name: str, platform: str, config_key: str) -> bool:
272
+ config_set = self.load_config_set(kernel_name)
273
+ return config_set.has_config(platform, config_key)
vllm/kernels/helion/register.py
+73
-15
@@ -65,7 +65,6 @@ vllm_helion_lib = Library("vllm_helion", "FRAGMENT") # noqa
65
def validate_helion_settings(
66
helion_settings: "helion.Settings | None", op_name: str
67
) -> None:
68
- """Validate that helion_settings doesn't contain conflicting options."""
68
if helion_settings is None:
69
return
70
@@ -93,6 +92,26 @@ def validate_helion_settings(
92
)
93
94
95
+def create_helion_decorated_kernel(
96
+ raw_kernel_func: Callable,
97
+ helion_settings: "helion.Settings | None" = None,
98
+ extra_kwargs: dict[str, Any] | None = None,
99
+) -> Any:
100
+ kernel_kwargs: dict[str, Any] = {}
101
+ if helion_settings:
102
+ kernel_kwargs.update(helion_settings.to_dict())
103
+
104
+ # Set static_shapes=False by default if user didn't explicitly set it
105
+ # This is needed for dynamic batch sizes and sequence lengths in vLLM
106
+ if kernel_kwargs.get("static_shapes") is not True:
107
+ kernel_kwargs["static_shapes"] = False
108
+
109
+ if extra_kwargs:
110
+ kernel_kwargs.update(extra_kwargs)
111
+
112
+ return helion.kernel(**kernel_kwargs)(raw_kernel_func)
113
+
114
+
115
class PresetConfigSearch(BaseAutotuner):
116
"""Custom autotuner that uses a preset config selector instead of autotuning."""
117
@@ -198,26 +217,19 @@ class ConfiguredHelionKernel:
217
key_computer = self._create_key_computer()
218
config_selector = self._create_config_selector(key_computer)
219
201
- kernel_kwargs = {}
202
- if self.helion_settings:
203
- kernel_kwargs.update(self.helion_settings.to_dict())
204
-
205
- # Set static_shapes=False by default if user didn't explicitly set it to True
206
- # This is needed for dynamic batch sizes and sequence lengths in vLLM
207
- if kernel_kwargs.get("static_shapes") is not True:
208
- kernel_kwargs["static_shapes"] = False
209
-
210
- kernel_kwargs["autotuner_fn"] = lambda _, args: PresetConfigSearch(
211
- args, config_selector
212
- )
213
- kernel_kwargs["key"] = key_computer
220
+ extra_kwargs = {
221
+ "autotuner_fn": lambda _, args: PresetConfigSearch(args, config_selector),
222
+ "key": key_computer,
223
+ }
224
225
logger.debug(
226
"Creating decorated kernel %s with custom autotuner on platform %s",
227
self.op_name,
228
self.platform,
229
)
220
- return helion.kernel(**kernel_kwargs)(self.raw_kernel_func)
230
+ return create_helion_decorated_kernel(
231
+ self.raw_kernel_func, self.helion_settings, extra_kwargs
232
+ )
233
234
235
class HelionKernelWrapper:
@@ -240,6 +252,7 @@ class HelionKernelWrapper:
252
self._config_picker: (
253
Callable[[tuple[Any, ...], list[str]], str | None] | None
254
) = None
255
+ self._input_generator: Callable[[], dict[str, tuple[Any, ...]]] | None = None
256
257
def __call__(self, *args, **kwargs):
258
configured_op = self.get_configured_op()
@@ -251,6 +264,51 @@ class HelionKernelWrapper:
264
self._config_picker = picker_func
265
return picker_func
266
267
+ def register_input_generator(
268
+ self, generator_func: Callable[[], dict[str, tuple[Any, ...]]]
269
+ ) -> Callable[[], dict[str, tuple[Any, ...]]]:
270
+ """
271
+ Register a function to generate inputs for autotuning and benchmarking.
272
+
273
+ Args:
274
+ generator_func: Function that returns dict[str, tuple] where:
275
+ - key: Configuration identifier (e.g., "4096", "hidden_4096")
276
+ - value: Tuple of arguments to pass to the kernel
277
+
278
+ Returns:
279
+ The registered function (for decorator usage)
280
+
281
+ Example:
282
+ @kernel_wrapper.register_input_generator
283
+ def generate_inputs():
284
+ return {
285
+ "4096": (torch.randn(4096, device="cuda"), 0.5),
286
+ "8192": (torch.randn(8192, device="cuda"), 0.5),
287
+ }
288
+ """
289
+ self._input_generator = generator_func
290
+ return generator_func
291
+
292
+ def get_inputs(self) -> dict[str, tuple[Any, ...]]:
293
+ if self._input_generator is None:
294
+ raise NotImplementedError(
295
+ f"No input generator registered for kernel '{self.op_name}'. "
296
+ f"Use @{self.op_name}.register_input_generator to register one."
297
+ )
298
+ return self._input_generator()
299
+
300
+ def run_autotune(
301
+ self,
302
+ inputs: tuple[Any, ...],
303
+ autotune_effort: str = "quick",
304
+ ) -> Config:
305
+ """Run autotuning for a single input configuration."""
306
+ extra_kwargs = {"autotune_effort": autotune_effort}
307
+ autotune_kernel = create_helion_decorated_kernel(
308
+ self.raw_kernel_func, self.helion_settings, extra_kwargs
309
+ )
310
+ return autotune_kernel.autotune(inputs)
311
+
312
def get_configured_op(self) -> Any:
313
assert self._config_picker is not None, (
314
f"No config picker registered for kernel '{self.op_name}'. "