@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a4d70bef37

[Model Runner V2] Reserve CUDA graph memory (#53306)

Signed-off-by: Anh Tran <anh.tran.3889@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Anh Tran <anh.tran.3889@gmail.com> Co-authored-by: anhtra3889 <anhtra@nvidia.com>

Nick Hill committed Aug 23, 2026 at 23:11 UTC a4d70bef3724edb068c8206804154065acaa4cd4

13 files changed +573 -31

tests/test_config.py

+35

@@ -347,6 +347,41 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(

347 assert compilation_config.cudagraph_capture_sizes == expected_capture_sizes

348

349

350 +def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

351 + """Cudagraph memory profiling uses a minimal KV cache, so the Mamba

352 + block-count guard must only fire for the real cache sizing."""

353 + kv_cache_config = SimpleNamespace(has_mamba_layers=True, num_blocks=4)

354 +

355 + compilation_config = CompilationConfig(

356 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

357 + )

358 + with pytest.raises(ValueError, match="exceeds available Mamba cache blocks"):

359 + compilation_config.resolve_cudagraph_mode_and_sizes(

360 + AttentionCGSupport.ALWAYS,

361 + "FakeAttentionBackend",

362 + uniform_decode_query_len=1,

363 + use_v2_model_runner=True,

364 + tensor_parallel_size=1,

365 + kv_cache_config=kv_cache_config,

366 + max_num_reqs=256,

367 + )

368 +

369 + compilation_config = CompilationConfig(

370 + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,

371 + )

372 + cudagraph_mode = compilation_config.resolve_cudagraph_mode_and_sizes(

373 + AttentionCGSupport.ALWAYS,

374 + "FakeAttentionBackend",

375 + uniform_decode_query_len=1,

376 + use_v2_model_runner=True,

377 + tensor_parallel_size=1,

378 + kv_cache_config=kv_cache_config,

379 + max_num_reqs=256,

380 + is_profiling=True,

381 + )

382 + assert cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE

383 +

384 +

385 @pytest.mark.parametrize(

386 ("model_config", "expected"),

387 [

tests/v1/sample/test_logprobs.py

+4 -3

@@ -535,8 +535,9 @@ def test_logprobs_mode(logprobs_mode: LogprobsMode):

535 "facebook/opt-125m",

536 max_logprobs=5,

537 enable_prefix_caching=False,

538 - # 2 other llms alive during whole session

539 - gpu_memory_utilization=0.05,

538 + # 2 other llms alive during whole session; must also cover the

539 + # cudagraph memory reservation from startup profiling.

540 + gpu_memory_utilization=0.1,

541 max_model_len=16,

542 logprobs_mode=logprobs_mode,

543 )

@@ -575,7 +576,7 @@ def test_prompt_logprobs_mode():

576 llm = LLM(

577 "facebook/opt-125m",

578 enable_prefix_caching=False,

578 - gpu_memory_utilization=0.05,

579 + gpu_memory_utilization=0.1,

580 max_model_len=16,

581 logprobs_mode=mode,

582 )

tests/v1/worker/test_gpu_model_runner_v2_cudagraph_profiling.py new

+286

@@ -0,0 +1,286 @@

1 +#!/usr/bin/env python3

2 +# SPDX-License-Identifier: Apache-2.0

3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +"""Unit tests for GPUModelRunner (V2) CUDA graph memory profiling.

5 +

6 +These exercise the orchestration of ``profile_cudagraph_memory`` on CPU by

7 +building a runner via ``__new__`` and faking the GPU-only helpers, so the

8 +control flow (bootstrap -> sample FULL graphs into a throwaway pool ->

9 +extrapolate -> teardown) is covered without a GPU.

10 +See https://github.com/vllm-project/vllm/issues/49224.

11 +"""

12 +

13 +import contextlib

14 +from types import SimpleNamespace

15 +from typing import Any

16 +

17 +from vllm.compilation.counter import compilation_counter

18 +from vllm.config.compilation import CUDAGraphMode

19 +from vllm.v1.worker.gpu import cudagraph_utils as cgu

20 +from vllm.v1.worker.gpu import model_runner as mrv2

21 +

22 +GLOBAL_POOL = "global-pool"

23 +THROWAWAY_POOL = "throwaway-pool"

24 +

25 +

26 +class _FakeCudaGraphManager:

27 + def __init__(

28 + self, needs_capture: bool, num_full_descs: int, piecewise_only: bool = False

29 + ) -> None:

30 + self._needs_capture = needs_capture

31 + self.pool: Any = GLOBAL_POOL

32 + descs = [object() for _ in range(num_full_descs)]

33 + if piecewise_only:

34 + self._capture_descs = {CUDAGraphMode.PIECEWISE: descs}

35 + else:

36 + self._capture_descs = {CUDAGraphMode.FULL: descs} if needs_capture else {}

37 + # Profiling hooks set by profile_cudagraph_memory.

38 + self._max_full_descs_to_capture: int | None = None

39 + self._capture_mem_samples: list[int] | None = None

40 + self.use_breakable_cg = False

41 +

42 + def needs_capture(self) -> bool:

43 + return self._needs_capture

44 +

45 +

46 +def _make_profiling_runner(

47 + cudagraph_mode: CUDAGraphMode,

48 + *,

49 + needs_capture: bool = True,

50 + num_full_descs: int = 3,

51 + piecewise_only: bool = False,

52 + captured_bytes: int = 7 << 30,

53 + mem_samples: list[int] | None = None,

54 +) -> Any:

55 + runner: Any = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)

56 + runner.compilation_config = SimpleNamespace(cudagraph_mode=cudagraph_mode)

57 + runner.cudagraph_manager = _FakeCudaGraphManager(

58 + needs_capture, num_full_descs, piecewise_only

59 + )

60 + runner.vllm_config = SimpleNamespace()

61 +

62 + events: list[str] = []

63 + runner.events = events

64 +

65 + def _capture_model() -> int:

66 + events.append("capture")

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

68 + samples = runner.cudagraph_manager._capture_mem_samples

69 + if samples is not None:

70 + samples.extend(mem_samples or [])

71 + return captured_bytes

72 +

73 + runner.capture_model = _capture_model

74 + return runner

75 +

76 +

77 +def _patch_module(monkeypatch) -> None:

78 + @contextlib.contextmanager

79 + def _fake_set_current_vllm_config(_cfg):

80 + yield

81 +

82 + monkeypatch.setattr(cgu, "set_current_vllm_config", _fake_set_current_vllm_config)

83 + monkeypatch.setattr(

84 + cgu,

85 + "current_platform",

86 + SimpleNamespace(graph_pool_handle=lambda: THROWAWAY_POOL),

87 + )

88 + monkeypatch.setattr(

89 + cgu, "_init_minimal_kv_cache_for_profiling", lambda r: r.events.append("init")

90 + )

91 + monkeypatch.setattr(

92 + cgu, "_teardown_profiling_state", lambda r: r.events.append("teardown")

93 + )

94 + # The profiler reads free GPU memory before/after to compute what it

95 + # retained; default to a constant (nothing retained).

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

97 + monkeypatch.setattr(

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

99 + )

100 +

101 +

102 +def test_profile_cudagraph_memory_disabled_returns_zero(monkeypatch):

103 + _patch_module(monkeypatch)

104 + runner = _make_profiling_runner(CUDAGraphMode.NONE)

105 +

106 + result = cgu.profile_cudagraph_memory(runner)

107 +

108 + assert result == 0

109 + # No KV-cache bootstrap or teardown when cudagraphs are disabled.

110 + assert runner.events == []

111 +

112 +

113 +def test_profile_cudagraph_memory_no_graphs_tears_down(monkeypatch):

114 + _patch_module(monkeypatch)

115 + runner = _make_profiling_runner(CUDAGraphMode.FULL, needs_capture=False)

116 +

117 + result = cgu.profile_cudagraph_memory(runner)

118 +

119 + assert result == 0

120 + # Bootstrapped then cleaned up, without capturing or touching the pool.

121 + assert runner.events == ["init", "teardown"]

122 + assert runner.cudagraph_manager.pool == GLOBAL_POOL

123 +

124 +

125 +def test_profile_cudagraph_memory_samples_and_extrapolates(monkeypatch):

126 + _patch_module(monkeypatch)

127 + gib = 1 << 30

128 + # Measured delta 1000 MiB includes the sampled FULL graphs (100 + 20 MiB).

129 + # Extrapolated FULL cost for 3 graphs: 100 + 2 * 20 = 140 MiB.

130 + runner = _make_profiling_runner(

131 + CUDAGraphMode.FULL,

132 + num_full_descs=3,

133 + captured_bytes=1000 * gib,

134 + mem_samples=[100 * gib, 20 * gib],

135 + )

136 +

137 + result = cgu.profile_cudagraph_memory(runner)

138 +

139 + assert result == (1000 - (100 + 20) + (100 + 2 * 20)) * gib

140 + # Bootstrap, capture, and teardown run in order.

141 + assert runner.events == ["init", "capture", "teardown"]

142 + # Capture must use a throwaway pool, not the persistent global pool.

143 + assert runner.cudagraph_manager.pool == THROWAWAY_POOL

144 + # FULL capture must be limited to the largest few graphs.

145 + assert (

146 + runner.cudagraph_manager._max_full_descs_to_capture

147 + == cgu._FULL_GRAPH_PROFILING_SAMPLES

148 + )

149 +

150 +

151 +def test_profile_cudagraph_memory_piecewise_only_returns_measured(monkeypatch):

152 + _patch_module(monkeypatch)

153 + captured_bytes = 5 << 30

154 + runner = _make_profiling_runner(

155 + CUDAGraphMode.FULL_AND_PIECEWISE,

156 + piecewise_only=True,

157 + captured_bytes=captured_bytes,

158 + )

159 +

160 + result = cgu.profile_cudagraph_memory(runner)

161 +

162 + # No FULL graphs to sample or extrapolate: the measured delta is exact.

163 + assert result == captured_bytes

164 +

165 +

166 +def test_profile_cudagraph_memory_tears_down_on_capture_error(monkeypatch):

167 + _patch_module(monkeypatch)

168 + runner = _make_profiling_runner(CUDAGraphMode.FULL)

169 +

170 + def _boom() -> int:

171 + runner.events.append("capture")

172 + raise RuntimeError("capture failed")

173 +

174 + runner.capture_model = _boom

175 +

176 + try:

177 + cgu.profile_cudagraph_memory(runner)

178 + except RuntimeError:

179 + pass

180 + else:

181 + raise AssertionError("expected capture error to propagate")

182 +

183 + # Teardown still runs even if capture raises.

184 + assert runner.events == ["init", "capture", "teardown"]

185 +

186 +

187 +def test_profile_cudagraph_memory_restores_compilation_counters(monkeypatch):

188 + _patch_module(monkeypatch)

189 + runner = _make_profiling_runner(CUDAGraphMode.FULL)

190 +

191 + def _capture_model() -> int:

192 + compilation_counter.num_cudagraph_captured += 5

193 + compilation_counter.num_gpu_runner_capture_triggers += 1

194 + return 1 << 30

195 +

196 + runner.capture_model = _capture_model

197 + captured_before = compilation_counter.num_cudagraph_captured

198 + triggers_before = compilation_counter.num_gpu_runner_capture_triggers

199 +

200 + cgu.profile_cudagraph_memory(runner)

201 +

202 + # Profiling captures are discarded, so they must not inflate the

203 + # compilation counters; the real capture_model() runs later.

204 + assert compilation_counter.num_cudagraph_captured == captured_before

205 + assert compilation_counter.num_gpu_runner_capture_triggers == triggers_before

206 +

207 +

208 +def test_model_runner_delegates_to_cudagraph_utils(monkeypatch):

209 + runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner)

210 + monkeypatch.setattr(mrv2, "_profile_cudagraph_memory", lambda r: 42)

211 + assert runner.profile_cudagraph_memory() == 42

212 +

213 +

214 +def test_extrapolate_full_graph_memory():

215 + mib = 1 << 20

216 + # No samples (e.g. no FULL graphs): nothing to add.

217 + assert cgu._extrapolate_full_graph_memory([], 0) == 0

218 + # A single graph costs exactly its sample.

219 + assert cgu._extrapolate_full_graph_memory([100 * mib], 1) == 100 * mib

220 + # First capture + per-graph cost for the rest.

221 + assert (

222 + cgu._extrapolate_full_graph_memory([100 * mib, 20 * mib], 5)

223 + == (100 + 4 * 20) * mib

224 + )

225 + # Per-graph cost is floored to account for driver overhead.

226 + assert cgu._extrapolate_full_graph_memory([100 * mib, 0], 3) == (100 + 2 * 1) * mib

227 +

228 +

229 +def test_profile_cudagraph_memory_clears_captured_graphs(monkeypatch):

230 + _patch_module(monkeypatch)

231 + runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)

232 +

233 + cleared: list[str] = []

234 + monkeypatch.setattr(

235 + cgu.CUDAGraphWrapper,

236 + "clear_all_graphs",

237 + classmethod(lambda cls: cleared.append("piecewise")),

238 + )

239 + monkeypatch.setattr(

240 + cgu.BreakableCUDAGraphWrapper,

241 + "clear_all_graphs",

242 + classmethod(lambda cls: cleared.append("breakable")),

243 + )

244 +

245 + cgu.profile_cudagraph_memory(runner)

246 +

247 + # Profiling captures are discarded so the real capture re-captures them

248 + # against the KV cache.

249 + assert cleared == ["piecewise", "breakable"]

250 +

251 +

252 +def test_profile_cudagraph_memory_redirects_wrapper_pools(monkeypatch):

253 + """Piecewise wrappers must capture into the throwaway pool too.

254 +

255 + Profiling graphs captured into the persistent global pool and then

256 + discarded drop the pool's use_count to 0, tripping the c10 allocator's

257 + create_or_incref_pool assert when the real capture reuses that pool.

258 + """

259 + _patch_module(monkeypatch)

260 + runner = _make_profiling_runner(CUDAGraphMode.FULL_AND_PIECEWISE)

261 +

262 + class _FakeWrapper:

263 + def __init__(self) -> None:

264 + self.graph_pool: Any = GLOBAL_POOL

265 + self.pool_during_capture: Any = None

266 +

267 + def clear_graphs(self) -> None:

268 + pass

269 +

270 + wrapper = _FakeWrapper()

271 + cgu.CUDAGraphWrapper._all_instances.add(wrapper)

272 + try:

273 + capture_model = runner.capture_model

274 +

275 + def _capture_model() -> int:

276 + wrapper.pool_during_capture = wrapper.graph_pool

277 + return capture_model()

278 +

279 + runner.capture_model = _capture_model

280 +

281 + cgu.profile_cudagraph_memory(runner)

282 +

283 + assert wrapper.pool_during_capture == THROWAWAY_POOL

284 + assert wrapper.graph_pool == GLOBAL_POOL

285 + finally:

286 + cgu.CUDAGraphWrapper._all_instances.discard(wrapper)

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

+4 -1

@@ -416,7 +416,10 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):

416 )

417

418 assert isinstance(attn_metadata_raw, dict)

419 - attn_metadata_narrowed = attn_metadata_raw[self.prefix]

419 + attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)

420 + if attn_metadata_narrowed is None:

421 + # Profile/warmup dummy runs skip mamba-family metadata.

422 + return

423 assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)

424 m = attn_metadata_narrowed

425 has_initial_state = m.has_initial_state

vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py

+4 -1

@@ -305,7 +305,10 @@ class OlmoHybridGatedDeltaNetAttention(GatedDeltaNetAttention):

305 return

306

307 assert isinstance(attn_metadata, dict)

308 - attn_metadata = attn_metadata[self.prefix] # type: ignore[assignment]

308 + attn_metadata = attn_metadata.get(self.prefix) # type: ignore[assignment]

309 + if attn_metadata is None:

310 + # Profile/warmup dummy runs skip mamba-family metadata.

311 + return

312 assert isinstance(attn_metadata, GDNAttentionMetadata)

313 has_initial_state = attn_metadata.has_initial_state

314 spec_query_start_loc = attn_metadata.spec_query_start_loc

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

+16 -12

@@ -1202,13 +1202,14 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1202 forward_context = get_forward_context()

1203 attn_metadata_raw = forward_context.attn_metadata

1204

1205 - if attn_metadata_raw is None:

1205 + attn_metadata = None

1206 + if isinstance(attn_metadata_raw, dict):

1207 + attn_metadata = attn_metadata_raw.get(self.prefix)

1208 + if attn_metadata is None:

1209 v_dim = core_attn_out.shape[-1] * core_attn_out.shape[-2]

1210 self._warmup_prefill_kernels(qkvz, v_dim)

1211 return

1212

1210 - assert isinstance(attn_metadata_raw, dict)

1211 - attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]

1213 assert isinstance(attn_metadata, GDNAttentionMetadata)

1214

1215 # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved

@@ -1259,12 +1260,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1260 forward_context = get_forward_context()

1261 attn_metadata_raw = forward_context.attn_metadata

1262

1262 - if attn_metadata_raw is None:

1263 + attn_metadata = None

1264 + if isinstance(attn_metadata_raw, dict):

1265 + attn_metadata = attn_metadata_raw.get(self.prefix)

1266 + if attn_metadata is None:

1267 self._warmup_prefill_kernels(mixed_qkv, 0)

1268 return

1269

1266 - assert isinstance(attn_metadata_raw, dict)

1267 - attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]

1270 assert isinstance(attn_metadata, GDNAttentionMetadata)

1271

1272 if (

@@ -1773,12 +1775,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1775 forward_context = get_forward_context()

1776 attn_metadata_raw = forward_context.attn_metadata

1777 qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size

1776 - if attn_metadata_raw is None:

1778 + attn_metadata = None

1779 + if isinstance(attn_metadata_raw, dict):

1780 + attn_metadata = attn_metadata_raw.get(self.prefix)

1781 + if attn_metadata is None:

1782 self._warmup_prefill_kernels(mixed_qkvz[:, :qkv_size], 0)

1783 return

1784

1780 - assert isinstance(attn_metadata_raw, dict)

1781 - attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]

1785 assert isinstance(attn_metadata, GDNAttentionMetadata)

1786 mixed_qkv, output_gate_flat = mixed_qkvz.split(

1787 [qkv_size, self.value_dim // self.tp_size], dim=-1

@@ -1856,12 +1859,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1859 ) -> None:

1860 forward_context = get_forward_context()

1861 attn_metadata_raw = forward_context.attn_metadata

1859 - if attn_metadata_raw is None:

1862 + attn_metadata = None

1863 + if isinstance(attn_metadata_raw, dict):

1864 + attn_metadata = attn_metadata_raw.get(self.prefix)

1865 + if attn_metadata is None:

1866 self._warmup_prefill_kernels(mixed_qkv, 0)

1867 return

1868

1863 - assert isinstance(attn_metadata_raw, dict)

1864 - attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]

1869 assert isinstance(attn_metadata, GDNAttentionMetadata)

1870 if (

1871 self._can_use_fused_gdn_mtp_decode(attn_metadata)

vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py

+2 -1

@@ -542,7 +542,8 @@ class BailingMoELinearAttention(LinearAttention):

542 attn_metadata = forward_context.attn_metadata

543 if attn_metadata is not None:

544 assert isinstance(attn_metadata, dict)

545 - attn_metadata = attn_metadata[self.prefix] # type: ignore

545 + attn_metadata = attn_metadata.get(self.prefix) # type: ignore

546 + if attn_metadata is not None:

547 assert isinstance(attn_metadata, LinearAttentionMetadata)

548 num_actual_tokens = (

549 attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens

vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py

+2 -1

@@ -282,7 +282,8 @@ class MiniMaxText01LinearAttention(LinearAttention):

282 attn_metadata: AttentionMetadata | None = None

283 if attn_metadata_raw is not None:

284 assert isinstance(attn_metadata_raw, dict)

285 - attn_metadata = attn_metadata_raw[self.prefix]

285 + attn_metadata = attn_metadata_raw.get(self.prefix)

286 + if attn_metadata is not None:

287 assert isinstance(attn_metadata, LinearAttentionMetadata)

288 num_actual_tokens = (

289 attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens

vllm/model_executor/layers/mamba/mamba_mixer.py

+2 -1

@@ -269,7 +269,8 @@ class MambaMixer(MambaBase, PluggableLayer):

269 attn_metadata: AttentionMetadata | None = None

270 if attn_metadata_raw is not None:

271 assert isinstance(attn_metadata_raw, dict)

272 - attn_metadata = attn_metadata_raw[self.prefix]

272 + attn_metadata = attn_metadata_raw.get(self.prefix)

273 + if attn_metadata is not None:

274 assert isinstance(attn_metadata, Mamba1AttentionMetadata)

275 query_start_loc_p = attn_metadata.query_start_loc_p

276 state_indices_tensor_p = attn_metadata.state_indices_tensor_p

vllm/model_executor/layers/mamba/mamba_mixer2.py

+2 -1

@@ -710,7 +710,8 @@ class MambaMixer2(MambaBase, PluggableLayer):

710 attn_metadata: AttentionMetadata | None = None

711 if attn_metadata_raw is not None:

712 assert isinstance(attn_metadata_raw, dict)

713 - attn_metadata = attn_metadata_raw[self.prefix]

713 + attn_metadata = attn_metadata_raw.get(self.prefix)

714 + if attn_metadata is not None:

715 assert isinstance(attn_metadata, Mamba2AttentionMetadata)

716 # conv_state must be (..., dim, width-1) for the conv kernels.

717 # DS layout stores it that way directly; SD layout needs a

vllm/model_executor/layers/mamba/short_conv.py

+6 -3

@@ -113,8 +113,10 @@ class ShortConv(MambaBase, PluggableLayer):

113 attn_metadata: AttentionMetadata | None = None

114 if attn_metadata_raw is not None:

115 assert isinstance(attn_metadata_raw, dict)

116 - attn_metadata = attn_metadata_raw[self.prefix]

117 - assert isinstance(attn_metadata, ShortConvAttentionMetadata)

116 + attn_metadata = attn_metadata_raw.get(self.prefix)

117 + assert attn_metadata is None or isinstance(

118 + attn_metadata, ShortConvAttentionMetadata

119 + )

120

121 BCx, _ = self.in_proj(hidden_states)

122 B, C, x = BCx.chunk(3, dim=-1)

@@ -225,7 +227,8 @@ class ShortConv(MambaBase, PluggableLayer):

227 attn_metadata: AttentionMetadata | None = None

228 if attn_metadata_raw is not None:

229 assert isinstance(attn_metadata_raw, dict)

228 - attn_metadata = attn_metadata_raw[self.prefix]

230 + attn_metadata = attn_metadata_raw.get(self.prefix)

231 + if attn_metadata is not None:

232 assert isinstance(attn_metadata, ShortConvAttentionMetadata)

233 conv_state = (

234 self.kv_cache[0]

vllm/v1/worker/gpu/cudagraph_utils.py

+187 -2

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +import gc

4 from collections import defaultdict

5 from collections.abc import Callable

6 from dataclasses import dataclass

7 from itertools import product

7 -from typing import Any, NamedTuple, Protocol

8 +from typing import TYPE_CHECKING, Any, NamedTuple, Protocol

10 import torch

11 import torch.nn as nn

@@ -15,7 +16,8 @@ from vllm.compilation.breakable_cudagraph import (

16 is_breakable_cudagraph_enabled,

17 )

18 from vllm.compilation.counter import compilation_counter

18 -from vllm.config import VllmConfig

19 +from vllm.compilation.cuda_graph import CUDAGraphWrapper

20 +from vllm.config import VllmConfig, set_current_vllm_config

21 from vllm.config.compilation import CUDAGraphMode

22 from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id

23 from vllm.distributed.parallel_state import (

@@ -37,6 +39,9 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

39 from vllm.v1.worker.gpu.model_states.interface import ModelState

40 from vllm.v1.worker.utils import AttentionGroup

42 +if TYPE_CHECKING:

43 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner

44 +

45 logger = init_logger(__name__)

@@ -132,6 +137,12 @@ class CudaGraphManager:

137

138 self._graphs_captured = False

139

140 + # Profiling hooks, set only by profile_cudagraph_memory() below: cap

141 + # FULL-mode capture at the N largest descriptors and record each

142 + # captured FULL graph's memory delta for extrapolation.

143 + self._max_full_descs_to_capture: int | None = None

144 + self._capture_mem_samples: list[int] | None = None

145 +

146 self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}

147 self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}

148

@@ -327,6 +338,14 @@ class CudaGraphManager:

338 continue

339

340 descs = self._capture_descs[mode]

341 + if (

342 + mode == CUDAGraphMode.FULL

343 + and self._max_full_descs_to_capture is not None

344 + ):

345 + # Profiling only: capture a sample of the largest FULL

346 + # graphs; the total cost is extrapolated from their

347 + # per-graph memory deltas.

348 + descs = descs[: self._max_full_descs_to_capture]

349 if is_global_first_rank():

350 descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})")

351 for desc in descs:

@@ -362,6 +381,9 @@ class CudaGraphManager:

381 set_graph_pool_id(self.pool)

382 else:

383 set_graph_pool_id(current_platform.graph_pool_handle())

384 + if self._capture_mem_samples is not None:

385 + torch.accelerator.synchronize()

386 + free_before = torch.accelerator.get_memory_info()[0]

387 with torch.cuda.graph(graph, self.pool):

388 forward_fn(CUDAGraphMode.NONE)

389 # Join offloader's copy stream after forward to avoid

@@ -369,6 +391,10 @@ class CudaGraphManager:

391 # forks copy_stream, but wait_prefetch only happens in

392 # the next forward pass.

393 get_offloader().join_after_forward()

394 + if self._capture_mem_samples is not None:

395 + torch.accelerator.synchronize()

396 + free_after = torch.accelerator.get_memory_info()[0]

397 + self._capture_mem_samples.append(free_before - free_after)

398 self.graphs[desc] = graph

399 compilation_counter.num_cudagraph_captured += 1

400 self._graphs_captured = True

@@ -668,3 +694,162 @@ def prepare_inputs_to_capture(

694 for_capture=full_cudagraph,

695 )

696 return AttentionState(attn_metadata, slot_mappings_by_layer)

697 +

698 +

699 +# ---------------------------------------------------------------------------

700 +# CUDA graph memory profiling

701 +# ---------------------------------------------------------------------------

702 +

703 +# Number of FULL graphs captured during profiling; the total FULL capture

704 +# cost is extrapolated from this sample to avoid a second full capture.

705 +_FULL_GRAPH_PROFILING_SAMPLES = 2

706 +# Floor for the extrapolated per-graph cost (driver overhead per graph).

707 +_MIN_PER_GRAPH_BYTES = 1 << 20

708 +

709 +

710 +@torch.inference_mode()

711 +def profile_cudagraph_memory(runner: "GPUModelRunner") -> int:

712 + """Estimate the GPU memory needed for CUDA graph capture.

713 +

714 + Called during memory profiling, *before* the real KV cache is allocated,

715 + so that ``Worker.determine_available_memory`` can reserve headroom for

716 + graph capture. Bootstraps a minimal KV cache, runs ``capture_model()``

717 + once, then releases everything so the real init/capture path starts clean.

718 +

719 + FULL graphs bake in KV cache pointers, so only the largest few are

720 + captured (into a throwaway pool) and their total cost is extrapolated.

721 + PIECEWISE, encoder and speculator graphs are measured in full. All

722 + profiling captures are discarded afterwards: replaying graphs recorded

723 + against the throwaway profiling state is unsafe (e.g. inductor graph

724 + partition reclaims the storages of earlier cudagraph recordings once the

725 + real capture records new ones, leading to use-after-free crashes).

726 + """

727 + if runner.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:

728 + return 0

729 +

730 + gc.collect()

731 + torch.accelerator.empty_cache()

732 +

733 + with set_current_vllm_config(runner.vllm_config):

734 + _init_minimal_kv_cache_for_profiling(runner)

735 +

736 + manager = runner.cudagraph_manager

737 + assert manager is not None

738 +

739 + # Don't count profiling captures; the real capture_model() runs later.

740 + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured

741 + saved_capture_triggers = compilation_counter.num_gpu_runner_capture_triggers

742 + all_wrappers: list[Any] = []

743 + original_pools: dict[int, Any] = {}

744 + try:

745 + if not manager.needs_capture():

746 + return 0

747 + # Capture all profiling graphs into a throwaway pool so their memory

748 + # is reclaimed on teardown rather than retained by the persistent

749 + # global pool (which the real capture reuses). This must include the

750 + # piecewise wrappers, not just the FULL-graph manager pool: graphs

751 + # captured into the global pool and then discarded drop its use_count

752 + # to 0, and the real capture on the same pool trips the c10 allocator's

753 + # create_or_incref_pool assert ("use_count > 0 INTERNAL ASSERT FAILED").

754 + manager.pool = current_platform.graph_pool_handle()

755 + if manager.use_breakable_cg:

756 + # The breakable runner is otherwise created lazily during capture,

757 + # after the pool swap below, and would capture into the global

758 + # pool. Create it now so its pool gets swapped too.

759 + manager.init_breakable_cg_runner(runner.model)

760 + all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

761 + BreakableCUDAGraphWrapper._all_instances

762 + )

763 + for wrapper in all_wrappers:

764 + original_pools[id(wrapper)] = wrapper.graph_pool

765 + wrapper.graph_pool = manager.pool

766 + manager._max_full_descs_to_capture = _FULL_GRAPH_PROFILING_SAMPLES

767 + mem_samples: list[int] = []

768 + manager._capture_mem_samples = mem_samples

769 +

770 + measured = int(runner.capture_model())

771 +

772 + # The measured delta covers PIECEWISE, encoder and speculator graphs

773 + # plus the sampled FULL graphs; swap the sampled FULL cost for the

774 + # extrapolated total. FULL and PIECEWISE share one pool here just as

775 + # they share the global pool at runtime, so the overlap is not

776 + # double-counted.

777 + num_full_graphs = len(manager._capture_descs.get(CUDAGraphMode.FULL, []))

778 + full_estimate = _extrapolate_full_graph_memory(mem_samples, num_full_graphs)

779 + return max(measured - sum(mem_samples) + full_estimate, 0)

780 + finally:

781 + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

782 + compilation_counter.num_gpu_runner_capture_triggers = saved_capture_triggers

783 + CUDAGraphWrapper.clear_all_graphs()

784 + BreakableCUDAGraphWrapper.clear_all_graphs()

785 + for wrapper in all_wrappers:

786 + if id(wrapper) in original_pools:

787 + wrapper.graph_pool = original_pools[id(wrapper)]

788 + _teardown_profiling_state(runner)

789 +

790 +

791 +def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> int:

792 + """Extrapolate the total FULL capture cost from samples of the largest

793 + graphs. The first capture allocates the pool baseline; later graphs mostly

794 + reuse it, so the second sample is taken as the per-graph cost."""

795 + if not mem_samples:

796 + return 0

797 + first_capture = mem_samples[0]

798 + per_graph = max(mem_samples[1], _MIN_PER_GRAPH_BYTES) if len(mem_samples) > 1 else 0

799 + return first_capture + (total_graphs - 1) * per_graph

800 +

801 +

802 +def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None:

803 + """Allocate the smallest KV cache that still lets every graph be captured."""

804 + from vllm.v1.core.kv_cache_utils import (

805 + get_kv_cache_config_from_groups,

806 + get_kv_cache_groups,

807 + )

808 +

809 + kv_cache_spec = runner.get_kv_cache_spec()

810 + kv_cache_groups = get_kv_cache_groups(runner.vllm_config, kv_cache_spec)

811 + # At least one block per sequence is required to capture the graphs.

812 + min_blocks = (

813 + min(runner.max_num_reqs, runner.compilation_config.max_cudagraph_capture_size)

814 + or 1

815 + )

816 + saved_override = runner.cache_config.num_gpu_blocks_override

817 + runner.cache_config.num_gpu_blocks_override = min_blocks

818 + try:

819 + minimal_config = get_kv_cache_config_from_groups(

820 + runner.vllm_config, kv_cache_groups, available_memory=0

821 + )

822 + finally:

823 + runner.cache_config.num_gpu_blocks_override = saved_override

824 +

825 + runner.initialize_kv_cache(minimal_config, is_profiling=True)

826 + runner.cache_config.num_gpu_blocks = minimal_config.num_blocks

827 +

828 +

829 +def _teardown_profiling_state(runner: "GPUModelRunner") -> None:

830 + """Release the profiling KV cache and captured graphs while keeping model

831 + weights, so the real ``initialize_kv_cache`` starts from a clean slate."""

832 + torch.accelerator.synchronize()

833 + if hasattr(runner, "kv_caches"):

834 + runner.kv_caches.clear()

835 + if hasattr(runner, "attn_groups"):

836 + runner.attn_groups.clear()

837 + if hasattr(runner, "kv_cache_config"):

838 + del runner.kv_cache_config

839 + # Dropping the manager releases the profiling graphs and throwaway pool.

840 + runner.cudagraph_manager = None

841 + # Release encoder graphs captured during profiling; the real

842 + # capture_model() re-captures them.

843 + if runner.model_state.supports_mm_inputs:

844 + runner.model_state.encoder_runner.clear()

845 + # Detach profiling KV tensors held by attention layers.

846 + for layer in runner.compilation_config.static_forward_context.values():

847 + if hasattr(layer, "kv_cache"):

848 + kv_cache = layer.kv_cache

849 + layer.kv_cache = (

850 + torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []

851 + )

852 + runner.cache_config.num_gpu_blocks = None

853 + runner.maybe_remove_all_loras(runner.lora_config)

854 + gc.collect()

855 + torch.accelerator.empty_cache()

vllm/v1/worker/gpu/model_runner.py

+23 -5

@@ -94,6 +94,9 @@ from vllm.v1.worker.gpu.cudagraph_utils import (

94 BatchExecutionDescriptor,

95 ModelCudaGraphManager,

96 )

97 +from vllm.v1.worker.gpu.cudagraph_utils import (

98 + profile_cudagraph_memory as _profile_cudagraph_memory,

99 +)

100 from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp

101 from vllm.v1.worker.gpu.ec_connector import get_ec_connector

102 from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after

@@ -499,7 +502,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

502 def get_kv_cache_spec(self):

503 return get_kv_cache_spec(self.vllm_config)

504

502 - def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

505 + def initialize_kv_cache(

506 + self, kv_cache_config: KVCacheConfig, is_profiling: bool = False

507 + ) -> None:

508 kv_cache_config = deepcopy(kv_cache_config)

509 self.kv_cache_config = kv_cache_config

510

@@ -585,6 +590,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

590 tensor_parallel_size=self.parallel_config.tensor_parallel_size,

591 kv_cache_config=self.kv_cache_config,

592 max_num_reqs=self.max_num_reqs,

593 + is_profiling=is_profiling,

594 )

595 self.cudagraph_manager = ModelCudaGraphManager(

596 self.vllm_config,

@@ -618,7 +624,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

624 self.kernel_block_sizes,

625 self.vllm_config,

626 )

621 - self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

627 + if is_profiling:

628 + self.kv_connector = NO_OP_KV_CONNECTOR

629 + else:

630 + self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict)

631

632 def _init_kv_zero_meta(self) -> None:

633 """Build KV-block zeroing metadata; invoked from gpu_worker."""

@@ -832,9 +841,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

841 if self.pooling_runner is not None:

842 self.pooling_runner.clear()

843

844 + @torch.inference_mode()

845 def profile_cudagraph_memory(self) -> int:

836 - # NOTE(woosuk): It is TBD whether we keep this API or not.

837 - return 0

846 + """Estimate the GPU memory required to capture CUDA graphs."""

847 + return _profile_cudagraph_memory(self)

848

849 @torch.inference_mode()

850 def capture_model(self) -> int:

@@ -1521,12 +1531,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1531 slot_mappings, self.kv_cache_config

1532 )

1533 assert block_tables is not None

1534 + attn_groups = self.attn_groups

1535 + if dummy_run and is_profile:

1536 + # Mamba layers take a cheap warmup path with no metadata;

1537 + # attention metadata is still built so those kernels tune.

1538 + attn_groups = [

1539 + [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)]

1540 + for groups in attn_groups

1541 + ]

1542 attn_metadata = self.model_state.prepare_attn(

1543 input_batch,

1544 batch_desc.cg_mode,

1545 block_tables,

1546 slot_mappings,

1529 - self.attn_groups,

1547 + attn_groups,

1548 self.kv_cache_config,

1549 # FULL replay reads capture-time metadata buffers. Re-stage them

1550 # from the zeroed dummy block tables instead of retaining state