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 6, 20:30 UTC
[Kernel] [Helion] [12/N] Use FakeTensorMode to avoid GPU allocation during config key computation (#36563)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Yanan Cao committed
Mar 11, 2026 at 14:25 UTC
a3774a819897ff60ab12a7622f587452f6208680
1 file changed
+12
-7
scripts/autotune_helion_kernels.py
+12
-7
@@ -27,6 +27,7 @@ import time
27
from dataclasses import dataclass
28
29
import torch
30
+from torch._subclasses.fake_tensor import FakeTensorMode
31
32
try:
33
import helion
@@ -109,7 +110,8 @@ def autotune_kernel(
110
)
111
112
try:
112
- inputs_dict = kernel_wrapper.get_inputs()
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)
@@ -126,15 +128,15 @@ def autotune_kernel(
128
"Autotuning kernel '%s' for platform '%s' with %d configs",
129
kernel_name,
130
platform,
129
- len(inputs_dict),
131
+ len(all_config_keys),
132
)
133
132
- configs_to_autotune = {}
134
if not force:
135
existing_configs = config_manager.get_platform_configs(
136
kernel_name, platform
137
)
137
- for config_key, inputs in inputs_dict.items():
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",
@@ -142,12 +144,12 @@ def autotune_kernel(
144
platform,
145
)
146
else:
145
- configs_to_autotune[config_key] = inputs
147
+ keys_to_autotune.append(config_key)
148
else:
149
logger.debug("Force mode enabled, will re-autotune all configs")
148
- configs_to_autotune = inputs_dict
150
+ keys_to_autotune = all_config_keys
151
150
- if not configs_to_autotune:
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.",
@@ -162,6 +164,9 @@ def autotune_kernel(
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 = []