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 5, 07:15 UTC
[Model Runner V2] support auto resolve cudagraph mode/sizes based on attn backend (#32936)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
Apr 10, 2026 at 23:27 UTC
8f121f7879669b5b6aced90f3d47a518ffd7f95d
6 files changed
+294
-213
vllm/config/compilation.py
+148
@@ -26,6 +26,8 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
26
27
if TYPE_CHECKING:
28
from vllm.config import VllmConfig
29
+ from vllm.v1.attention.backend import AttentionCGSupport
30
+ from vllm.v1.kv_cache_interface import KVCacheConfig
31
else:
32
VllmConfig = object
33
@@ -1241,6 +1243,152 @@ class CompilationConfig:
1243
assert "none" in self.custom_ops
1244
return f"+{op}" in self.custom_ops
1245
1246
+ def resolve_cudagraph_mode_and_sizes(
1247
+ self,
1248
+ min_cg_support: "AttentionCGSupport",
1249
+ min_cg_attn_backend: str | None,
1250
+ uniform_decode_query_len: int = 1,
1251
+ tensor_parallel_size: int = 1,
1252
+ kv_cache_config: "KVCacheConfig | None" = None,
1253
+ max_num_reqs: int | None = None,
1254
+ is_profiling: bool = False,
1255
+ ) -> CUDAGraphMode:
1256
+ from vllm.v1.attention.backend import AttentionCGSupport
1257
+
1258
+ cudagraph_mode = self.cudagraph_mode
1259
+ if cudagraph_mode is None or cudagraph_mode == CUDAGraphMode.NONE:
1260
+ self.cudagraph_mode = CUDAGraphMode.NONE
1261
+ return CUDAGraphMode.NONE
1262
+
1263
+ # Check cudagraph for mixed batch is supported
1264
+ if (
1265
+ cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL
1266
+ and min_cg_support != AttentionCGSupport.ALWAYS
1267
+ ):
1268
+ msg = (
1269
+ f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
1270
+ f"with {min_cg_attn_backend} backend (support: "
1271
+ f"{min_cg_support})"
1272
+ )
1273
+ if min_cg_support == AttentionCGSupport.NEVER:
1274
+ # if not supported any full cudagraphs, just raise it.
1275
+ msg += (
1276
+ "; please try cudagraph_mode=PIECEWISE, and "
1277
+ "make sure compilation mode is VLLM_COMPILE"
1278
+ )
1279
+ raise ValueError(msg)
1280
+
1281
+ # attempt to resolve the full cudagraph related mode
1282
+ if self.splitting_ops_contain_attention():
1283
+ msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE"
1284
+ cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
1285
+ else:
1286
+ msg += "; setting cudagraph_mode=FULL_DECODE_ONLY"
1287
+ cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
1288
+ logger.warning(msg)
1289
+
1290
+ # check that if we are doing decode full-cudagraphs it is supported
1291
+ if (
1292
+ cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
1293
+ and min_cg_support == AttentionCGSupport.NEVER
1294
+ ):
1295
+ msg = (
1296
+ f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
1297
+ f"with {min_cg_attn_backend} backend (support: "
1298
+ f"{min_cg_support})"
1299
+ )
1300
+ if self.mode == CompilationMode.VLLM_COMPILE and (
1301
+ self.splitting_ops_contain_attention()
1302
+ or self.use_inductor_graph_partition
1303
+ ):
1304
+ msg += (
1305
+ "; setting cudagraph_mode=PIECEWISE because "
1306
+ "attention is compiled piecewise"
1307
+ )
1308
+ cudagraph_mode = CUDAGraphMode.PIECEWISE
1309
+ else:
1310
+ msg += (
1311
+ "; setting cudagraph_mode=NONE because "
1312
+ "attention is not compiled piecewise"
1313
+ )
1314
+ cudagraph_mode = CUDAGraphMode.NONE
1315
+ logger.warning(msg)
1316
+
1317
+ # check that if we are doing spec-decode + decode full-cudagraphs it is
1318
+ # supported
1319
+ if (
1320
+ cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
1321
+ and uniform_decode_query_len > 1
1322
+ and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value
1323
+ ):
1324
+ msg = (
1325
+ f"CUDAGraphMode.{cudagraph_mode.name} is not supported"
1326
+ f" with spec-decode for attention backend "
1327
+ f"{min_cg_attn_backend} (support: {min_cg_support})"
1328
+ )
1329
+ if self.splitting_ops_contain_attention():
1330
+ msg += "; setting cudagraph_mode=PIECEWISE"
1331
+ cudagraph_mode = CUDAGraphMode.PIECEWISE
1332
+ else:
1333
+ msg += "; setting cudagraph_mode=NONE"
1334
+ cudagraph_mode = CUDAGraphMode.NONE
1335
+ logger.warning(msg)
1336
+
1337
+ # double check that we can support full cudagraph if they are requested
1338
+ # even after automatic downgrades
1339
+ if (
1340
+ cudagraph_mode.has_full_cudagraphs()
1341
+ and min_cg_support == AttentionCGSupport.NEVER
1342
+ ):
1343
+ raise ValueError(
1344
+ f"CUDAGraphMode.{cudagraph_mode.name} is not "
1345
+ f"supported with {min_cg_attn_backend} backend ("
1346
+ f"support:{min_cg_support}) "
1347
+ "; please try cudagraph_mode=PIECEWISE, "
1348
+ "and make sure compilation mode is VLLM_COMPILE"
1349
+ )
1350
+
1351
+ # Adjust cudagraph sizes to be a multiple of uniform_decode_query_len
1352
+ # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix:
1353
+ # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536
1354
+ # Will be removed in the near future when we have separate cudagraph capture
1355
+ # sizes for decode and mixed prefill-decode.
1356
+ if (
1357
+ cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
1358
+ and uniform_decode_query_len > 1
1359
+ ):
1360
+ self.adjust_cudagraph_sizes_for_spec_decode(
1361
+ uniform_decode_query_len,
1362
+ tensor_parallel_size,
1363
+ )
1364
+
1365
+ # For Mamba models with FULL decode cudagraphs, each decode
1366
+ # sequence needs one Mamba cache block. The decode cudagraph
1367
+ # dispatcher already caps batch sizes at max_num_seqs, so we just
1368
+ # need to verify that enough blocks exist. Raising here instead
1369
+ # of silently capping cudagraph_capture_sizes avoids unintended
1370
+ # restrictions on PIECEWISE (prefill) cudagraphs.
1371
+ # See: https://github.com/vllm-project/vllm/issues/34094
1372
+ if (
1373
+ kv_cache_config is not None
1374
+ and max_num_reqs is not None
1375
+ and cudagraph_mode.has_full_cudagraphs()
1376
+ and not is_profiling
1377
+ and kv_cache_config.has_mamba_layers
1378
+ and max_num_reqs > kv_cache_config.num_blocks
1379
+ ):
1380
+ raise ValueError(
1381
+ f"max_num_seqs ({max_num_reqs}) exceeds available Mamba cache "
1382
+ f"blocks ({kv_cache_config.num_blocks}). Each decode sequence "
1383
+ "requires one Mamba cache block, so CUDA graph capture cannot "
1384
+ "proceed. Please lower max_num_seqs to at most "
1385
+ f"{kv_cache_config.num_blocks} or increase "
1386
+ "gpu_memory_utilization."
1387
+ )
1388
+
1389
+ self.cudagraph_mode = cudagraph_mode
1390
+ return cudagraph_mode
1391
+
1392
def adjust_cudagraph_sizes_for_spec_decode(
1393
self, uniform_decode_query_len: int, tensor_parallel_size: int
1394
):
vllm/v1/worker/gpu/attn_utils.py
+39
-5
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Sequence
4
+from dataclasses import dataclass
5
from typing import Any, cast
6
7
import numpy as np
@@ -8,7 +9,11 @@ import torch
9
10
from vllm.config import VllmConfig, get_layers_from_vllm_config
11
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
11
-from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
12
+from vllm.v1.attention.backend import (
13
+ AttentionBackend,
14
+ AttentionCGSupport,
15
+ CommonAttentionMetadata,
16
+)
17
from vllm.v1.kv_cache_interface import (
18
AttentionSpec,
19
KVCacheConfig,
@@ -18,6 +23,12 @@ from vllm.v1.kv_cache_interface import (
23
from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache
24
25
26
+@dataclass(frozen=True)
27
+class AttentionCGSupportInfo:
28
+ min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS
29
+ min_cg_attn_backend: str | None = None
30
+
31
+
32
def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
33
kv_cache_spec: dict[str, KVCacheSpec] = {}
34
layer_type = cast(type[Any], AttentionLayerBase)
@@ -34,10 +45,17 @@ def init_attn_backend(
45
vllm_config: VllmConfig,
46
device: torch.device,
47
active_layer_names: set[str] | None = None,
37
-):
48
+) -> tuple[
49
+ dict[str, type[AttentionBackend]],
50
+ list[list[AttentionGroup]],
51
+ AttentionCGSupportInfo,
52
+]:
53
attn_backends: dict[str, type[AttentionBackend]] = {}
54
attn_groups: list[list[AttentionGroup]] = []
55
attn_backend_workspace: torch.Tensor | None = None
56
+ # Find minimum cudagraph support across all attention backends
57
+ min_cg_support = AttentionCGSupport.ALWAYS
58
+ min_cg_attn_backend = None
59
for kv_cache_group_id, kv_cache_group_spec in enumerate(
60
kv_cache_config.kv_cache_groups
61
):
@@ -86,8 +104,24 @@ def init_attn_backend(
104
else:
105
if hasattr(builder, "set_workspace_buffer"):
106
builder.set_workspace_buffer(attn_backend_workspace)
107
+ # Check cudagraph support for the attention backend
108
+ cg_support = builder.get_cudagraph_support(
109
+ vllm_config,
110
+ cast(AttentionSpec, kv_cache_group_spec.kv_cache_spec),
111
+ )
112
+ if cg_support.value < min_cg_support.value:
113
+ min_cg_support = cg_support
114
+ min_cg_attn_backend = attn_backend.__name__
115
attn_groups.append(groups)
90
- return attn_backends, attn_groups
116
+
117
+ return (
118
+ attn_backends,
119
+ attn_groups,
120
+ AttentionCGSupportInfo(
121
+ min_cg_support=min_cg_support,
122
+ min_cg_attn_backend=min_cg_attn_backend,
123
+ ),
124
+ )
125
126
127
def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
@@ -110,7 +144,7 @@ def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
144
def _reshape_kv_cache(
145
kv_cache_config: KVCacheConfig,
146
kv_cache_raw_tensors: dict[str, torch.Tensor],
113
- attn_backends: dict[str, AttentionBackend],
147
+ attn_backends: dict[str, type[AttentionBackend]],
148
cache_dtype: str,
149
) -> dict[str, torch.Tensor]:
150
kv_caches: dict[str, torch.Tensor] = {}
@@ -158,7 +192,7 @@ def init_kv_cache(
192
runner_kv_caches: list[torch.Tensor],
193
forward_context: dict[str, Any],
194
kv_cache_config: KVCacheConfig,
161
- attn_backends: dict[str, AttentionBackend],
195
+ attn_backends: dict[str, type[AttentionBackend]],
196
device: torch.device,
197
cache_dtype: str,
198
) -> dict[str, torch.Tensor]:
vllm/v1/worker/gpu/dp_utils.py
+43
-1
@@ -20,7 +20,7 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N
20
21
22
def sync_cudagraph_and_dp_padding(
23
- cudagraph_manager: CudaGraphManager,
23
+ cudagraph_manager: CudaGraphManager | None,
24
desired_batch_desc: BatchExecutionDescriptor,
25
num_tokens: int,
26
num_reqs: int,
@@ -61,6 +61,10 @@ def sync_cudagraph_and_dp_padding(
61
num_reqs=num_reqs,
62
), num_tokens_across_dp
63
64
+ assert cudagraph_manager is not None, (
65
+ "cudagraph_manager should only be None during profile run, "
66
+ "where synced_cg_mode must be NONE across all DP ranks"
67
+ )
68
synced_num_tokens = int(num_tokens_across_dp.max().item())
69
synced_uniform_token_count = uniform_token_counts_across_dp[0]
70
# If ranks disagree on the uniform token count, or its 0 (means None) set to None
@@ -79,3 +83,41 @@ def sync_cudagraph_and_dp_padding(
83
num_tokens_across_dp[:] = synced_desc.num_tokens
84
85
return synced_desc, num_tokens_across_dp
86
+
87
+
88
+def dispatch_cg_and_sync_dp(
89
+ cudagraph_manager: CudaGraphManager | None,
90
+ num_reqs: int,
91
+ num_tokens: int,
92
+ uniform_token_count: int | None,
93
+ dp_size: int,
94
+ dp_rank: int,
95
+ need_eager: bool = False,
96
+) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
97
+ if need_eager:
98
+ batch_desc = BatchExecutionDescriptor(
99
+ cg_mode=CUDAGraphMode.NONE,
100
+ num_tokens=num_tokens,
101
+ num_reqs=num_reqs,
102
+ )
103
+ else:
104
+ assert cudagraph_manager is not None, (
105
+ "cudagraph_manager should only be None during profile run, "
106
+ "where need_eager must be True"
107
+ )
108
+ batch_desc = cudagraph_manager.dispatch(
109
+ num_reqs, num_tokens, uniform_token_count
110
+ )
111
+
112
+ if dp_size == 1:
113
+ return batch_desc, None
114
+
115
+ return sync_cudagraph_and_dp_padding(
116
+ cudagraph_manager,
117
+ batch_desc,
118
+ num_tokens,
119
+ num_reqs,
120
+ uniform_token_count,
121
+ dp_size,
122
+ dp_rank,
123
+ )
vllm/v1/worker/gpu/model_runner.py
+36
-30
@@ -61,7 +61,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
61
ModelCudaGraphManager,
62
get_uniform_token_count,
63
)
64
-from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
64
+from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
65
from vllm.v1.worker.gpu.eplb_utils import EPLBController, step_eplb_after
66
from vllm.v1.worker.gpu.input_batch import (
67
InputBatch,
@@ -176,6 +176,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
176
177
# Draft tokens propagation - for spec-dec + struct outputs.
178
self.draft_tokens_handler = DraftTokensHandler(self.device)
179
+ self.uniform_decode_query_len = 1 + self.num_speculative_steps
180
181
# Pooling models.
182
self.is_pooling_model = self.model_config.runner_type == "pooling"
@@ -224,14 +225,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
225
device=self.device,
226
)
227
227
- # CUDA graphs.
228
+ # For CUDA graphs, and will init cudagraph_manager after init_attn_backend.
229
self.decode_query_len = self.num_speculative_steps + 1
229
- self.cudagraph_manager = ModelCudaGraphManager(
230
- self.vllm_config,
231
- self.device,
232
- self.compilation_config.cudagraph_mode,
233
- decode_query_len=self.decode_query_len,
234
- )
230
+ self.cudagraph_manager: ModelCudaGraphManager | None = None
231
# LoRA-related workers.
232
self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
233
# KV Connector if configured.
@@ -361,9 +357,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
357
cp_interleave=self.cp_interleave,
358
)
359
364
- self.attn_backends, self.attn_groups = init_attn_backend(
360
+ self.attn_backends, self.attn_groups, attn_cg_support = init_attn_backend(
361
self.kv_cache_config, self.vllm_config, self.device
362
)
363
+ cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
364
+ attn_cg_support.min_cg_support,
365
+ attn_cg_support.min_cg_attn_backend,
366
+ self.uniform_decode_query_len,
367
+ self.parallel_config.tensor_parallel_size,
368
+ self.kv_cache_config,
369
+ self.max_num_reqs,
370
+ )
371
+ self.cudagraph_manager = ModelCudaGraphManager(
372
+ self.vllm_config,
373
+ self.device,
374
+ cudagraph_mode,
375
+ decode_query_len=self.decode_query_len,
376
+ )
377
+ if self.speculator is not None:
378
+ self.speculator.init_cudagraph_manager(cudagraph_mode)
379
+
380
check_attention_cp_compatibility(self.vllm_config)
381
if self.speculator is not None:
382
# HACK(woosuk)
@@ -437,6 +450,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
450
intermediate_tensors=intermediate_tensors,
451
dummy_run=True,
452
skip_attn_for_dummy_run=skip_attn,
453
+ is_profile=is_profile,
454
)
455
self.kv_connector.set_disabled(False)
456
@@ -486,6 +500,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
500
dummy_run=True,
501
skip_attn_for_dummy_run=skip_attn,
502
mm_inputs=mm_inputs,
503
+ is_profile=is_profile,
504
)
505
506
assert hidden_states is not None # Last PP rank always has hidden_states
@@ -547,6 +562,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
562
563
@torch.inference_mode()
564
def capture_model(self) -> int:
565
+ assert self.cudagraph_manager is not None
566
if not self.cudagraph_manager.needs_capture():
567
logger.warning(
568
"Skipping CUDA graph capture. To turn on CUDA graph capture, "
@@ -915,6 +931,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
931
intermediate_tensors: IntermediateTensors | None = None,
932
dummy_run: bool = False,
933
skip_attn_for_dummy_run: bool = False,
934
+ is_profile: bool = False,
935
) -> ModelRunnerOutput | IntermediateTensors | None:
936
if not dummy_run:
937
# Update the request states.
@@ -934,34 +951,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
951
max_query_len = max(scheduler_output.num_scheduled_tokens.values())
952
uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)
953
937
- batch_desc = self.cudagraph_manager.dispatch(
938
- num_reqs, num_toks, uniform_tok_count
939
- )
940
- num_tokens_across_dp = None
941
-
954
skip_compiled = False
955
if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
956
# Encoder-decoder models such as Whisper should run eager/non-compiled
957
# when encoder inputs are scheduled, because this step updates
958
# cross-attention cache with dynamic encoder outputs.
947
- # Override batch_desc to NONE.
959
skip_compiled = True
949
- batch_desc = BatchExecutionDescriptor(
950
- cg_mode=CUDAGraphMode.NONE,
951
- num_tokens=num_toks,
952
- num_reqs=num_reqs,
953
- )
960
955
- if self.dp_size > 1:
956
- batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
957
- self.cudagraph_manager,
958
- batch_desc,
959
- num_toks,
960
- num_reqs,
961
- uniform_tok_count,
962
- self.dp_size,
963
- self.dp_rank,
964
- )
961
+ batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
962
+ self.cudagraph_manager,
963
+ num_reqs,
964
+ num_toks,
965
+ uniform_tok_count,
966
+ self.dp_size,
967
+ self.dp_rank,
968
+ need_eager=is_profile or skip_compiled,
969
+ )
970
971
if batch_desc.num_tokens == 0:
972
# All DP ranks have zero tokens to run.
@@ -1059,6 +1064,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1064
# Use explicit cudagraph replay for FULL mode.
1065
# NOTE(woosuk): Here, we don't need to pass the input tensors,
1066
# because they are already copied to the CUDA graph input buffers.
1067
+ assert self.cudagraph_manager is not None
1068
self.kv_connector.pre_forward(scheduler_output)
1069
model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
1070
else:
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+17
-32
@@ -19,10 +19,7 @@ from vllm.v1.worker.gpu.attn_utils import (
19
init_attn_backend,
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
-from vllm.v1.worker.gpu.cudagraph_utils import (
23
- BatchExecutionDescriptor,
24
-)
25
-from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
22
+from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
23
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
24
from vllm.v1.worker.gpu.model_states.interface import ModelState
25
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
@@ -98,15 +95,19 @@ class EagleSpeculator:
95
device=device,
96
)
97
101
- # currently we don't support PIECEWISE for Eagle.
102
- cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
98
+ self.cudagraph_manager: EagleCudaGraphManager | None = None
99
+
100
+ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
101
if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
102
cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
103
else:
104
cudagraph_mode = CUDAGraphMode.NONE
105
106
self.cudagraph_manager = EagleCudaGraphManager(
109
- vllm_config, device, cudagraph_mode, self.draft_tokens
107
+ self.vllm_config,
108
+ self.device,
109
+ cudagraph_mode,
110
+ self.draft_tokens,
111
)
112
113
def load_model(self, target_model: nn.Module) -> None:
@@ -133,7 +134,7 @@ class EagleSpeculator:
134
) -> None:
135
self.model_state = model_state
136
self.kv_cache_config = kv_cache_config
136
- _, self.attn_groups = init_attn_backend(
137
+ _, self.attn_groups, _ = init_attn_backend(
138
kv_cache_config,
139
self.vllm_config,
140
self.device,
@@ -242,29 +243,6 @@ class EagleSpeculator:
243
idx_mapping, query_start_loc, pos, num_tokens_padded
244
)
245
245
- def _dispatch_and_sync_dp(
246
- self,
247
- cudagraph_manager: EagleCudaGraphManager,
248
- num_reqs: int,
249
- num_tokens: int,
250
- uniform_token_count: int | None,
251
- ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
252
- batch_desc = cudagraph_manager.dispatch(
253
- num_reqs, num_tokens, uniform_token_count
254
- )
255
- num_tokens_across_dp = None
256
- if self.dp_size > 1:
257
- batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
258
- cudagraph_manager,
259
- batch_desc,
260
- num_tokens,
261
- num_reqs,
262
- uniform_token_count,
263
- self.dp_size,
264
- self.dp_rank,
265
- )
266
- return batch_desc, num_tokens_across_dp
267
-
246
def _build_draft_attn_metadata(
247
self,
248
num_reqs: int,
@@ -303,8 +281,10 @@ class EagleSpeculator:
281
return attn_metadata
282
283
def capture_model(self) -> None:
284
+ assert self.cudagraph_manager is not None
285
if self.num_speculative_steps == 1:
286
return
287
+
288
logger.info("Capturing model for Eagle speculator...")
289
self.cudagraph_manager.capture(
290
self.generate_draft,
@@ -342,6 +322,7 @@ class EagleSpeculator:
322
dummy_run: bool = False,
323
skip_attn_for_dummy_run: bool = False,
324
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
325
+ is_profile: bool = False,
326
) -> torch.Tensor:
327
# NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
328
# number of rejected tokens, we maintain the size of eagle's input_ids and
@@ -430,11 +411,14 @@ class EagleSpeculator:
411
412
# Each request produces exactly 1 token per draft decode step,
413
# enabling FULL cudagraph.
433
- decode_batch_desc, num_tokens_across_dp = self._dispatch_and_sync_dp(
414
+ decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(
415
self.cudagraph_manager,
416
num_reqs,
417
num_reqs,
418
uniform_token_count=1,
419
+ dp_size=self.dp_size,
420
+ dp_rank=self.dp_rank,
421
+ need_eager=is_profile,
422
)
423
424
attn_metadata_updated = None
@@ -461,6 +445,7 @@ class EagleSpeculator:
445
)
446
447
if decode_batch_desc.cg_mode == CUDAGraphMode.FULL:
448
+ assert self.cudagraph_manager is not None
449
self.cudagraph_manager.run_fullgraph(decode_batch_desc)
450
else:
451
self.generate_draft(
vllm/v1/worker/gpu_model_runner.py
+11
-145
@@ -6307,7 +6307,7 @@ class GPUModelRunner(
6307
cudagraph_mode.
6308
"""
6309
min_cg_support = AttentionCGSupport.ALWAYS
6310
- min_cg_backend_name = None
6310
+ min_cg_attn_backend = None
6311
6312
for attn_backend_set, kv_cache_group in zip(
6313
attention_backends, kv_cache_groups
@@ -6320,152 +6320,18 @@ class GPUModelRunner(
6320
)
6321
if cg_support.value < min_cg_support.value:
6322
min_cg_support = cg_support
6323
- min_cg_backend_name = attn_backend.__name__
6324
- # Flexible resolve the cudagraph mode
6325
- cudagraph_mode = self.compilation_config.cudagraph_mode
6326
- assert cudagraph_mode is not None
6327
- # check cudagraph for mixed batch is supported
6328
- if (
6329
- cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL
6330
- and min_cg_support != AttentionCGSupport.ALWAYS
6331
- ):
6332
- msg = (
6333
- f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
6334
- f"with {min_cg_backend_name} backend (support: "
6335
- f"{min_cg_support})"
6336
- )
6337
- if min_cg_support == AttentionCGSupport.NEVER:
6338
- # if not supported any full cudagraphs, just raise it.
6339
- msg += (
6340
- "; please try cudagraph_mode=PIECEWISE, and "
6341
- "make sure compilation mode is VLLM_COMPILE"
6342
- )
6343
- raise ValueError(msg)
6344
-
6345
- # attempt to resolve the full cudagraph related mode
6346
- if self.compilation_config.splitting_ops_contain_attention():
6347
- msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE"
6348
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6349
- CUDAGraphMode.FULL_AND_PIECEWISE
6350
- )
6351
- else:
6352
- msg += "; setting cudagraph_mode=FULL_DECODE_ONLY"
6353
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6354
- CUDAGraphMode.FULL_DECODE_ONLY
6355
- )
6356
- logger.warning(msg)
6357
-
6358
- # check that if we are doing decode full-cudagraphs it is supported
6359
- if (
6360
- cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
6361
- and min_cg_support == AttentionCGSupport.NEVER
6362
- ):
6363
- msg = (
6364
- f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
6365
- f"with {min_cg_backend_name} backend (support: "
6366
- f"{min_cg_support})"
6367
- )
6368
- if self.compilation_config.mode == CompilationMode.VLLM_COMPILE and (
6369
- self.compilation_config.splitting_ops_contain_attention()
6370
- or self.compilation_config.use_inductor_graph_partition
6371
- ):
6372
- msg += (
6373
- "; setting cudagraph_mode=PIECEWISE because "
6374
- "attention is compiled piecewise"
6375
- )
6376
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6377
- CUDAGraphMode.PIECEWISE
6378
- )
6379
- else:
6380
- msg += (
6381
- "; setting cudagraph_mode=NONE because "
6382
- "attention is not compiled piecewise"
6383
- )
6384
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6385
- CUDAGraphMode.NONE
6386
- )
6387
- logger.warning(msg)
6388
-
6389
- # check that if we are doing spec-decode + decode full-cudagraphs it is
6390
- # supported
6391
- if (
6392
- cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
6393
- and self.uniform_decode_query_len > 1
6394
- and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value
6395
- ):
6396
- msg = (
6397
- f"CUDAGraphMode.{cudagraph_mode.name} is not supported"
6398
- f" with spec-decode for attention backend "
6399
- f"{min_cg_backend_name} (support: {min_cg_support})"
6400
- )
6401
- if self.compilation_config.splitting_ops_contain_attention():
6402
- msg += "; setting cudagraph_mode=PIECEWISE"
6403
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6404
- CUDAGraphMode.PIECEWISE
6405
- )
6406
- else:
6407
- msg += "; setting cudagraph_mode=NONE"
6408
- cudagraph_mode = self.compilation_config.cudagraph_mode = (
6409
- CUDAGraphMode.NONE
6410
- )
6411
- logger.warning(msg)
6412
-
6413
- # double check that we can support full cudagraph if they are requested
6414
- # even after automatic downgrades
6415
- if (
6416
- cudagraph_mode.has_full_cudagraphs()
6417
- and min_cg_support == AttentionCGSupport.NEVER
6418
- ):
6419
- raise ValueError(
6420
- f"CUDAGraphMode.{cudagraph_mode.name} is not "
6421
- f"supported with {min_cg_backend_name} backend ("
6422
- f"support:{min_cg_support}) "
6423
- "; please try cudagraph_mode=PIECEWISE, "
6424
- "and make sure compilation mode is VLLM_COMPILE"
6425
- )
6426
-
6427
- # if we have dedicated decode cudagraphs, and spec-decode is enabled,
6428
- # we need to adjust the cudagraph sizes to be a multiple of the uniform
6429
- # decode query length to avoid: https://github.com/vllm-project/vllm/issues/28207
6430
- # temp-fix: https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536
6431
- # Will be removed in the near future when we have separate cudagraph capture
6432
- # sizes for decode and mixed prefill-decode.
6433
- if (
6434
- cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
6435
- and cudagraph_mode.separate_routine()
6436
- and self.uniform_decode_query_len > 1
6437
- ):
6438
- self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(
6439
- self.uniform_decode_query_len, self.parallel_config.tensor_parallel_size
6440
- )
6441
-
6442
- # For Mamba models with FULL decode cudagraphs, each decode
6443
- # sequence needs one Mamba cache block. The decode cudagraph
6444
- # dispatcher already caps batch sizes at max_num_seqs, so we just
6445
- # need to verify that enough blocks exist. Raising here instead
6446
- # of silently capping cudagraph_capture_sizes avoids unintended
6447
- # restrictions on PIECEWISE (prefill) cudagraphs.
6448
- # See: https://github.com/vllm-project/vllm/issues/34094
6449
- if cudagraph_mode.has_full_cudagraphs() and not is_profiling:
6450
- has_mamba = any(
6451
- isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_groups
6452
- )
6453
- if has_mamba and self.kv_cache_config is not None:
6454
- num_blocks = self.kv_cache_config.num_blocks
6455
- if self.max_num_reqs > num_blocks:
6456
- raise ValueError(
6457
- f"max_num_seqs ({self.max_num_reqs}) exceeds "
6458
- f"available Mamba cache blocks ({num_blocks}). "
6459
- f"Each decode sequence requires one Mamba cache "
6460
- f"block, so CUDA graph capture cannot proceed. "
6461
- f"Please lower max_num_seqs to at most "
6462
- f"{num_blocks} or increase "
6463
- f"gpu_memory_utilization."
6464
- )
6465
-
6323
+ min_cg_attn_backend = attn_backend.__name__
6324
+ cudagraph_mode = self.compilation_config.resolve_cudagraph_mode_and_sizes(
6325
+ min_cg_support,
6326
+ min_cg_attn_backend,
6327
+ self.uniform_decode_query_len,
6328
+ self.parallel_config.tensor_parallel_size,
6329
+ self.kv_cache_config,
6330
+ self.max_num_reqs,
6331
+ is_profiling=is_profiling,
6332
+ )
6333
# Trigger cudagraph dispatching keys initialization after
6334
# resolved cudagraph mode.
6468
- self.compilation_config.cudagraph_mode = cudagraph_mode
6335
self.cudagraph_dispatcher.initialize_cudagraph_keys(
6336
cudagraph_mode, self.uniform_decode_query_len
6337
)