@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 685074cd61

[Bugfix][V2] Warm up kernels before capturing CUDA graphs (#55341)

Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

aoshen02 committed Sep 5, 2026 at 06:44 UTC 685074cd61d6b8432575d11b38341559992cf9dd

7 files changed +129 -16

tests/v1/worker/test_gpu_model_runner_v2.py

+82

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import contextlib

5 from types import SimpleNamespace

7 import pytest

@@ -213,3 +214,84 @@ def test_append_block_ids_rejects_write_past_row_capacity():

214 )

215

216 assert block_tables.num_blocks.np[0, 1] == 3

217 +

218 +

219 +def _make_capture_runner(captured: bool) -> GPUModelRunner:

220 + """Minimal V2 runner for capture_model: fakes everything except the

221 + cudagraph_manager's needs_capture decision."""

222 + runner = GPUModelRunner.__new__(GPUModelRunner)

223 + runner.model_state = SimpleNamespace(supports_mm_inputs=False)

224 + runner.cudagraph_manager = SimpleNamespace(

225 + needs_capture=lambda: captured,

226 + capture=lambda *args, **kwargs: None,

227 + )

228 + runner.lora_config = None

229 + runner.maybe_setup_dummy_loras = lambda _cfg: contextlib.nullcontext()

230 + runner.speculator = None

231 + runner.adaptive_verification = None

232 + runner.model = None

233 + runner.input_buffers = None

234 + runner.pcp_manager = None

235 + runner.intermediate_tensors = None

236 + runner.block_tables = None

237 + runner.attn_groups = None

238 + runner.kv_cache_config = None

239 + runner.use_aux_hidden_state_outputs = False

240 + return runner

241 +

242 +

243 +def test_capture_model_locks_workspace_after_capture(monkeypatch):

244 + """A workspace resize after capture frees the buffer the captured graphs

245 + baked in, so capture_model must lock the workspace before returning

246 + (https://github.com/vllm-project/vllm/issues/55336)."""

247 + runner = _make_capture_runner(captured=True)

248 + monkeypatch.setattr(

249 + model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext

250 + )

251 + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)

252 + monkeypatch.setattr(

253 + torch.accelerator, "get_memory_info", lambda: (1 << 30, 1 << 30)

254 + )

255 + lock_calls = []

256 + monkeypatch.setattr(

257 + model_runner_module, "lock_workspace", lambda: lock_calls.append("lock")

258 + )

259 +

260 + runner.capture_model()

261 +

262 + assert lock_calls == ["lock"]

263 +

264 +

265 +def test_capture_model_skips_lock_when_nothing_captured(monkeypatch):

266 + """With no graphs to capture (e.g. enforce_eager) there is nothing baked

267 + into the workspace, so the early return must not lock it."""

268 + runner = _make_capture_runner(captured=False)

269 + lock_calls = []

270 + monkeypatch.setattr(

271 + model_runner_module, "lock_workspace", lambda: lock_calls.append("lock")

272 + )

273 +

274 + assert runner.capture_model() == 0

275 + assert lock_calls == []

276 +

277 +

278 +def test_capture_model_profile_only_skips_lock(monkeypatch):

279 + """The memory-profiling capture pass runs before kernel warmup and the

280 + real capture; locking there would stop the warmup from growing the

281 + workspace to its scheduler-realistic size."""

282 + runner = _make_capture_runner(captured=True)

283 + monkeypatch.setattr(

284 + model_runner_module, "freeze_gc_for_cudagraph_capture", contextlib.nullcontext

285 + )

286 + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)

287 + monkeypatch.setattr(

288 + torch.accelerator, "get_memory_info", lambda: (1 << 30, 1 << 30)

289 + )

290 + lock_calls = []

291 + monkeypatch.setattr(

292 + model_runner_module, "lock_workspace", lambda: lock_calls.append("lock")

293 + )

294 +

295 + runner.capture_model(profile_only=True)

296 +

297 + assert lock_calls == []

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py

+9 -8

@@ -67,7 +67,8 @@ def _make_profiling_runner(

67 events: list[str] = []

68 runner.events = events

70 - def _capture_model() -> int:

70 + def _capture_model(*, profile_only: bool = False) -> int:

71 + assert profile_only

72 events.append("capture")

73 # Simulate the manager's per-FULL-graph memory sampling.

74 samples = runner.cudagraph_manager._capture_mem_samples

@@ -183,7 +184,7 @@ def test_profile_cudagraph_memory_tears_down_on_capture_error(monkeypatch):

184 _patch_module(monkeypatch)

185 runner = _make_profiling_runner(CUDAGraphMode.FULL)

186

186 - def _boom() -> int:

187 + def _boom(*, profile_only: bool = False) -> int:

188 runner.events.append("capture")

189 raise RuntimeError("capture failed")

190

@@ -204,7 +205,7 @@ def test_profile_cudagraph_memory_restores_compilation_counters(monkeypatch):

205 _patch_module(monkeypatch)

206 runner = _make_profiling_runner(CUDAGraphMode.FULL)

207

207 - def _capture_model() -> int:

208 + def _capture_model(*, profile_only: bool = False) -> int:

209 compilation_counter.num_cudagraph_captured += 5

210 compilation_counter.num_gpu_runner_capture_triggers += 1

211 return 1 << 30

@@ -288,9 +289,9 @@ def test_profile_cudagraph_memory_redirects_wrapper_pools(monkeypatch):

289 try:

290 capture_model = runner.capture_model

291

291 - def _capture_model() -> int:

292 + def _capture_model(*, profile_only: bool = False) -> int:

293 wrapper.pool_during_capture = wrapper.graph_pool

293 - return capture_model()

294 + return capture_model(profile_only=profile_only)

295

296 runner.capture_model = _capture_model

297

@@ -328,9 +329,9 @@ def test_profile_cudagraph_memory_swaps_and_drops_speculator_managers(monkeypatc

329 pools_seen: list[Any] = []

330 capture_model = runner.capture_model

331

331 - def _capture_model() -> int:

332 + def _capture_model(*, profile_only: bool = False) -> int:

333 pools_seen.append(runner.speculator.cudagraph_manager.pool)

333 - return capture_model()

334 + return capture_model(profile_only=profile_only)

335

336 runner.capture_model = _capture_model

337

@@ -379,7 +380,7 @@ def test_profile_cudagraph_memory_frees_throwaway_pool(monkeypatch):

380 manager.pool = cgu.current_platform.get_global_graph_pool()

381 r.speculator = SimpleNamespace(cudagraph_manager=manager)

382

382 - def _capture_model() -> int:

383 + def _capture_model(*, profile_only: bool = False) -> int:

384 for owner in (

385 runner.cudagraph_manager,

386 runner.speculator.cudagraph_manager,

tests/v1/worker/test_workspace.py

+24

@@ -82,6 +82,30 @@ def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None:

82 assert len(pointers) == 4

85 +def test_workspace_lock_blocks_growth_and_unlock_restores(monkeypatch) -> None:

86 + """Once locked, oversized requests fail loudly instead of reallocating the

87 + buffer that captured CUDA graphs point at; unlock restores growth."""

88 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)

89 + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1)

90 +

91 + (buf,) = manager.get_simultaneous(((256,), torch.uint8))

92 + manager.lock()

93 + assert manager.is_locked()

94 +

95 + # Requests within the reserved size still reuse the same buffer.

96 + (same,) = manager.get_simultaneous(((256,), torch.uint8))

97 + (smaller,) = manager.get_simultaneous(((8,), torch.uint8))

98 + assert same.data_ptr() == buf.data_ptr()

99 + assert smaller.data_ptr() == buf.data_ptr()

100 +

101 + with pytest.raises(AssertionError, match="Workspace is locked"):

102 + manager.get_simultaneous(((512,), torch.uint8))

103 +

104 + manager.unlock()

105 + (grown,) = manager.get_simultaneous(((512,), torch.uint8))

106 + assert grown.numel() == 512

107 +

108 +

109 def test_workspace_lane_validation(monkeypatch) -> None:

110 monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)

111 manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1)

vllm/v1/worker/gpu/cudagraph_utils.py

+1 -1

@@ -828,7 +828,7 @@ def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:

828 mem_samples: list[int] = []

829 manager._capture_mem_samples = mem_samples

830

831 - measured = int(runner.capture_model())

831 + measured = int(runner.capture_model(profile_only=True))

832

833 # The measured delta covers PIECEWISE, encoder and speculator graphs

834 # plus the sampled FULL graphs; swap the sampled FULL cost for the

vllm/v1/worker/gpu/model_runner.py

+7 -2

@@ -173,7 +173,7 @@ from vllm.v1.worker.utils import (

173 copy_kv_cache_blocks_inplace,

174 get_uniform_decode_token_count,

175 )

176 -from vllm.v1.worker.workspace import use_workspace_lane

176 +from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane

177

178 logger = init_logger(__name__)

179

@@ -928,7 +928,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

928 return _profile_cudagraph_memory(self)

929

930 @torch.inference_mode()

931 - def capture_model(self) -> int:

931 + def capture_model(self, *, profile_only: bool = False) -> int:

932 assert self.cudagraph_manager is not None

933 capture_encoder = (

934 self.model_state.supports_mm_inputs

@@ -986,6 +986,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

986

987 end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

988

989 + if not profile_only:

990 + # Lock workspace to prevent resizing during execution. A resize after

991 + # capture frees the static cuda graph buffer.

992 + lock_workspace()

993 +

994 end_time = time.perf_counter()

995 elapsed_time = end_time - start_time

996 cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory

vllm/v1/worker/gpu_worker.py

+5 -4

@@ -790,6 +790,10 @@ class Worker(WorkerBase):

790 # cuda graph capture.

791 kernel_warmup(self)

792

793 + if self.use_v2_model_runner:

794 + # A workspace resize after capture frees what the graphs point at.

795 + warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)

796 +

797 cuda_graph_memory_bytes = 0

798 if not self.model_config.enforce_eager:

799 cuda_graph_memory_bytes = self.model_runner.capture_model()

@@ -868,10 +872,7 @@ class Worker(WorkerBase):

872

873 maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit)

874

871 - if self.use_v2_model_runner:

872 - # V2: Run full execute_model + sample_tokens to JIT compile triton kernels.

873 - warmup_kernels(self.model_runner, self.execute_model, self.sample_tokens)

874 - elif get_pp_group().is_last_rank:

875 + if not self.use_v2_model_runner and get_pp_group().is_last_rank:

876 # V1: Warm up sampler and preallocate memory buffer for logits and other

877 # sampling related tensors of max possible shape to avoid memory

878 # fragmentation issue.

vllm/v1/worker/mm_encoder_model_runner.py

+1 -1

@@ -67,7 +67,7 @@ class MMEncoderModelRunner(GPUModelRunner):

67 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

68 return {}

70 - def capture_model(self) -> int:

70 + def capture_model(self, *, profile_only: bool = False) -> int:

71 return 0

73 def _dummy_run(