minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix flashinfer autotune to only wrap run_once() (#19004)

Browse Source
This commit is contained in:
Cheng Wan
2026-02-19 20:02:21 -08:00
committed by GitHub
parent 64bca5315f
commit 13a4a0406e
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+7 -3
View File
@@ -152,6 +152,7 @@ from sglang.srt.utils import (
MultiprocessingSerializer, MultiprocessingSerializer,
cpu_has_amx_support, cpu_has_amx_support,
dynamic_import, dynamic_import,
empty_context,
enable_show_time_cost, enable_show_time_cost,
get_available_gpu_memory, get_available_gpu_memory,
get_cpu_ids_by_node, get_cpu_ids_by_node,
@@ -1861,12 +1862,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
logger.info("Running FlashInfer autotune...") logger.info("Running FlashInfer autotune...")
with torch.inference_mode(), autotune(): self._dummy_run(
self._dummy_run(batch_size=self.req_to_token_pool.size) batch_size=self.req_to_token_pool.size,
run_ctx=autotune(),
)
logger.info("FlashInfer autotune completed.") logger.info("FlashInfer autotune completed.")
def _dummy_run(self, batch_size: int): def _dummy_run(self, batch_size: int, run_ctx=None):
"""Run a dummy forward pass for warmup/profiling.""" """Run a dummy forward pass for warmup/profiling."""
if self.is_generation: if self.is_generation:
capture_forward_mode = ForwardMode.DECODE capture_forward_mode = ForwardMode.DECODE
@@ -2106,6 +2109,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
torch.get_device_module(self.device).synchronize() torch.get_device_module(self.device).synchronize()
self.tp_group.barrier() self.tp_group.barrier()
with torch.inference_mode(), run_ctx or empty_context():
run_once() run_once()
def init_device_graphs(self): def init_device_graphs(self):