@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 11d3976b88
Fork 0

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 8, 02:30 UTC

[Model Runner V2] support piecewise & mixed cudagraph (#32771)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>

zhrrr committed Feb 19, 2026 at 07:03 UTC 11d3976b883a0a785023d279a48c85c3ad777106
5 files changed +343 -108
vllm/v1/worker/gpu/cudagraph_utils.py
+179 -58
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Callable, Iterable
3 +from collections.abc import Callable
4 from typing import Any
5
6 import numpy as np
@@ -11,7 +11,8 @@ from tqdm import tqdm
11 from vllm.config import VllmConfig
12 from vllm.config.compilation import CUDAGraphMode
13 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14 -from vllm.forward_context import set_forward_context
14 +from vllm.forward_context import BatchDescriptor, set_forward_context
15 +from vllm.utils.math_utils import cdiv
16 from vllm.v1.attention.backend import AttentionMetadataBuilder
17 from vllm.v1.kv_cache_interface import KVCacheConfig
18 from vllm.v1.worker.gpu.attn_utils import (
@@ -34,14 +35,27 @@ class CudaGraphManager:
35 self.max_num_reqs = self.scheduler_config.max_num_seqs
36 self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
37 self.dp_size = vllm_config.parallel_config.data_parallel_size
38 +
39 + self.uniform_decode_query_len = 1
40 + spec_config = vllm_config.speculative_config
41 + if spec_config is not None:
42 + self.uniform_decode_query_len += spec_config.num_speculative_tokens
43 +
44 self.compilation_config = vllm_config.compilation_config
45 assert self.compilation_config is not None
46 self.cudagraph_mode = self.compilation_config.cudagraph_mode
40 - self.cudagraph_sizes = get_cudagraph_sizes(
47 +
48 + use_uniform_decode_cudagraph = (
49 + self.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
50 + and self.cudagraph_mode.separate_routine()
51 + )
52 + self.cudagraph_sizes, self.uniform_decode_cudagraph_sizes = get_cudagraph_sizes(
53 self.compilation_config.cudagraph_capture_sizes,
54 self.max_num_reqs,
55 self.max_num_tokens,
56 self.cudagraph_mode,
57 + self.uniform_decode_query_len,
58 + use_uniform_decode_cudagraph,
59 )
60
61 self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
@@ -54,20 +68,16 @@ class CudaGraphManager:
68 return len(self.cudagraph_sizes) > 0
69
70 def get_cudagraph_size(
57 - self,
58 - num_tokens_after_padding: int,
59 - num_tokens_per_request: Iterable[int],
71 + self, num_tokens: int, uniform_decode: bool = False
72 ) -> int | None:
61 - return get_cudagraph_size(
62 - num_tokens_after_padding,
63 - num_tokens_per_request,
64 - self.cudagraph_sizes,
65 - self.cudagraph_mode,
66 - )
73 + if uniform_decode and self.uniform_decode_cudagraph_sizes:
74 + return self.uniform_decode_cudagraph_sizes.get(num_tokens)
75 + return self.cudagraph_sizes.get(num_tokens)
76
77 def capture_graph(
78 self,
79 num_tokens: int,
80 + capture_cg_mode: CUDAGraphMode,
81 model: nn.Module,
82 input_buffers: InputBuffers,
83 mrope_positions: torch.Tensor | None,
@@ -75,8 +85,25 @@ class CudaGraphManager:
85 block_tables: BlockTables,
86 attn_metadata_builders: list[AttentionMetadataBuilder],
87 kv_cache_config: KVCacheConfig,
88 + has_lora: bool = False,
89 + uniform_decode: bool = False,
90 ) -> None:
79 - num_reqs = min(num_tokens, self.max_num_reqs)
91 + # select and check capture function
92 + assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
93 + f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
94 + )
95 + if capture_cg_mode == CUDAGraphMode.PIECEWISE:
96 + capture_fn = self._capture_piecewise_graph
97 + else:
98 + capture_fn = self._capture_full_graph
99 + # prepare inputs
100 + if uniform_decode:
101 + num_reqs = min(
102 + cdiv(num_tokens, self.uniform_decode_query_len),
103 + self.max_num_reqs,
104 + )
105 + else:
106 + num_reqs = min(num_tokens, self.max_num_reqs)
107 input_ids = input_buffers.input_ids[:num_tokens]
108 positions = input_buffers.positions[:num_tokens]
109 if self.uses_mrope:
@@ -92,6 +119,9 @@ class CudaGraphManager:
119 attn_metadata_builders,
120 self.max_model_len,
121 kv_cache_config,
122 + uniform_decode_query_len=(
123 + self.uniform_decode_query_len if uniform_decode else 0
124 + ),
125 )
126 num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
127
@@ -112,13 +142,40 @@ class CudaGraphManager:
142 if self.hidden_states is None:
143 self.hidden_states = torch.empty_like(hidden_states)
144
145 + capture_fn(
146 + num_tokens=num_tokens,
147 + num_reqs=num_reqs,
148 + model=model,
149 + input_ids=input_ids,
150 + positions=positions,
151 + inputs_embeds=inputs_embeds,
152 + num_tokens_across_dp=num_tokens_across_dp,
153 + attn_metadata=attn_metadata,
154 + slot_mappings=slot_mappings,
155 + has_lora=has_lora,
156 + )
157 +
158 + def _capture_full_graph(
159 + self,
160 + num_tokens: int,
161 + num_reqs: int,
162 + model: nn.Module,
163 + input_ids: torch.Tensor,
164 + positions: torch.Tensor,
165 + inputs_embeds: torch.Tensor | None,
166 + num_tokens_across_dp: torch.Tensor,
167 + attn_metadata: dict[str, Any] | None,
168 + slot_mappings: dict[str, torch.Tensor] | None,
169 + has_lora: bool = False,
170 + ) -> None:
171 + assert attn_metadata is not None
172 # Capture the graph.
173 assert num_tokens not in self.graphs
174 graph = torch.cuda.CUDAGraph()
175 with (
176 set_forward_context(
120 - attn_metadata,
121 - self.vllm_config,
177 + attn_metadata=attn_metadata,
178 + vllm_config=self.vllm_config,
179 num_tokens=num_tokens,
180 cudagraph_runtime_mode=CUDAGraphMode.NONE,
181 num_tokens_across_dp=num_tokens_across_dp,
@@ -131,9 +188,44 @@ class CudaGraphManager:
188 positions=positions,
189 inputs_embeds=inputs_embeds,
190 )
191 + assert self.hidden_states is not None
192 self.hidden_states[:num_tokens] = hidden_states
193 self.graphs[num_tokens] = graph
194
195 + def _capture_piecewise_graph(
196 + self,
197 + num_tokens: int,
198 + num_reqs: int,
199 + model: nn.Module,
200 + input_ids: torch.Tensor,
201 + positions: torch.Tensor,
202 + inputs_embeds: torch.Tensor | None,
203 + num_tokens_across_dp: torch.Tensor,
204 + attn_metadata: dict[str, Any] | None,
205 + slot_mappings: dict[str, torch.Tensor] | None,
206 + has_lora: bool = False,
207 + ) -> None:
208 + # create batch descriptor for piecewise cudagraph dispatch key
209 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens, has_lora=has_lora)
210 +
211 + # Capture run - CUDAGraphWrapper inside torch.compile will auto capture.
212 + with set_forward_context(
213 + attn_metadata=None, # piecewise no need attn_metadata
214 + vllm_config=self.vllm_config,
215 + num_tokens=num_tokens,
216 + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
217 + num_tokens_across_dp=num_tokens_across_dp,
218 + batch_descriptor=batch_descriptor,
219 + slot_mapping=slot_mappings,
220 + ):
221 + hidden_states = model(
222 + input_ids=input_ids,
223 + positions=positions,
224 + inputs_embeds=inputs_embeds,
225 + )
226 + assert self.hidden_states is not None
227 + self.hidden_states[:num_tokens] = hidden_states
228 +
229 @torch.inference_mode()
230 def capture(
231 self,
@@ -144,11 +236,11 @@ class CudaGraphManager:
236 block_tables: BlockTables,
237 attn_metadata_builders: list[AttentionMetadataBuilder],
238 kv_cache_config: KVCacheConfig,
239 + has_lora: bool = False,
240 ) -> None:
148 - capture_graphs(
149 - self.cudagraph_sizes,
150 - self.device,
151 - self.capture_graph,
241 + common_kwargs = dict(
242 + device=self.device,
243 + capture_fn=self.capture_graph,
244 model=model,
245 input_buffers=input_buffers,
246 mrope_positions=mrope_positions,
@@ -156,10 +248,50 @@ class CudaGraphManager:
248 block_tables=block_tables,
249 attn_metadata_builders=attn_metadata_builders,
250 kv_cache_config=kv_cache_config,
251 + has_lora=has_lora,
252 )
253
161 - def run(self, num_tokens: int) -> torch.Tensor:
162 - assert num_tokens in self.graphs
254 + # Phase 1: Capture for mixed prefill-decode batches if needed.
255 + mixed_mode = self.cudagraph_mode.mixed_mode()
256 + if mixed_mode != CUDAGraphMode.NONE:
257 + capture_graphs(
258 + cudagraph_sizes=self.cudagraph_sizes,
259 + capture_cudagraph_mode=mixed_mode,
260 + desc=f"Capturing CUDA graphs (mixed, {mixed_mode.name})",
261 + uniform_decode=False,
262 + **common_kwargs,
263 + )
264 +
265 + # Phase 2: Capture FULL graphs for uniform decode batches if needed.
266 + # This is only needed if we use a separate routine for decode batches
267 + # and the decode_mode is FULL.
268 + if self.uniform_decode_cudagraph_sizes:
269 + capture_graphs(
270 + cudagraph_sizes=self.uniform_decode_cudagraph_sizes,
271 + capture_cudagraph_mode=CUDAGraphMode.FULL,
272 + desc="Capturing CUDA graphs (decode, FULL)",
273 + uniform_decode=True,
274 + **common_kwargs,
275 + )
276 +
277 + def get_cudagraph_runtime_mode(
278 + self, num_reqs: int, num_tokens: int, max_query_len: int
279 + ) -> tuple[CUDAGraphMode, int | None]:
280 + is_uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
281 + num_tokens == max_query_len * num_reqs
282 + )
283 +
284 + cudagraph_size = self.get_cudagraph_size(num_tokens, is_uniform_decode)
285 + if cudagraph_size is None:
286 + cudagraph_mode = CUDAGraphMode.NONE
287 + elif is_uniform_decode:
288 + cudagraph_mode = self.cudagraph_mode.decode_mode()
289 + else:
290 + cudagraph_mode = self.cudagraph_mode.mixed_mode()
291 + return cudagraph_mode, cudagraph_size
292 +
293 + def run_fullgraph(self, num_tokens: int) -> torch.Tensor:
294 + assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens"
295 self.graphs[num_tokens].replay()
296 assert self.hidden_states is not None
297 return self.hidden_states[:num_tokens]
@@ -170,22 +302,18 @@ def get_cudagraph_sizes(
302 max_num_reqs: int,
303 max_num_tokens: int,
304 cudagraph_mode: CUDAGraphMode,
173 -) -> dict[int, int]:
174 - if not cudagraph_mode.has_full_cudagraphs():
175 - return {}
305 + uniform_decode_query_len: int = 1,
306 + uniform_decode_cudagraph: bool = False,
307 +) -> tuple[dict[int, int], dict[int, int]]:
308 + # Support both FULL and PIECEWISE cudagraph modes
309 + if cudagraph_mode == CUDAGraphMode.NONE:
310 + return {}, {}
311 if not capture_sizes:
177 - return {}
312 + return {}, {}
313
314 capture_sizes = sorted(capture_sizes)
180 - # Limit the capture sizes to the max number of requests or tokens.
181 - upper_bound = (
182 - max_num_reqs
183 - if cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY
184 - else max_num_tokens
185 - )
186 - capture_sizes = [x for x in capture_sizes if x <= upper_bound]
315 if not capture_sizes:
188 - return {}
316 + return {}, {}
317
318 cudagraph_sizes: dict[int, int] = {}
319 for i in range(1, capture_sizes[-1] + 1):
@@ -193,45 +321,34 @@ def get_cudagraph_sizes(
321 if i <= x:
322 cudagraph_sizes[i] = x
323 break
196 - return cudagraph_sizes
197 -
198 -
199 -def get_cudagraph_size(
200 - num_tokens_after_dp_padding: int,
201 - num_tokens_per_request: Iterable[int],
202 - cudagraph_sizes: dict[int, int],
203 - cudagraph_mode: CUDAGraphMode,
204 -) -> int | None:
205 - if not cudagraph_mode.has_full_cudagraphs():
206 - # No full CUDA graph is used.
207 - return None
208 -
209 - size = cudagraph_sizes.get(num_tokens_after_dp_padding)
210 - if size is None:
211 - # No CUDA graph for this size.
212 - return None
324
214 - is_mixed = any(x > 1 for x in num_tokens_per_request)
215 - if is_mixed and cudagraph_mode.mixed_mode() != CUDAGraphMode.FULL:
216 - # Prefill is included, and this mode doesn't use CUDA graph for it.
217 - return None
218 - return size
325 + uniform_decode_cudagraph_sizes: dict[int, int] = {}
326 + if uniform_decode_cudagraph:
327 + max_num_tokens = max_num_reqs * uniform_decode_query_len
328 + uniform_decode_cudagraph_sizes = {
329 + k: v
330 + for k, v in cudagraph_sizes.items()
331 + if v <= max_num_tokens and v >= uniform_decode_query_len
332 + }
333 + return cudagraph_sizes, uniform_decode_cudagraph_sizes
334
335
336 def capture_graphs(
337 cudagraph_sizes: dict[int, int],
338 device: torch.device,
339 capture_fn: Callable,
340 + capture_cudagraph_mode: CUDAGraphMode,
341 + desc: str = "Capturing CUDA graphs",
342 **capture_kwargs,
343 ) -> None:
344 # Capture larger graphs first.
345 sizes_to_capture = sorted(set(cudagraph_sizes.values()), reverse=True)
346 if is_global_first_rank():
230 - sizes_to_capture = tqdm(sizes_to_capture, desc="Capturing CUDA graphs")
347 + sizes_to_capture = tqdm(sizes_to_capture, desc=desc)
348
349 with graph_capture(device=device):
350 for size in sizes_to_capture:
234 - capture_fn(size, **capture_kwargs)
351 + capture_fn(size, capture_cudagraph_mode, **capture_kwargs)
352
353
354 def prepare_inputs_to_capture(
@@ -242,8 +359,12 @@ def prepare_inputs_to_capture(
359 attn_metadata_builders: list[AttentionMetadataBuilder],
360 max_model_len: int,
361 kv_cache_config: KVCacheConfig,
362 + uniform_decode_query_len: int = 0,
363 ) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
246 - num_tokens_per_req = num_tokens // num_reqs
364 + if uniform_decode_query_len > 0:
365 + num_tokens_per_req = uniform_decode_query_len
366 + else:
367 + num_tokens_per_req = num_tokens // num_reqs
368
369 query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
370 query_start_loc_np[-1] = num_tokens
vllm/v1/worker/gpu/dp_utils.py
+36 -19
@@ -13,48 +13,65 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N
13
14
15 def get_batch_metadata_across_dp(
16 - num_tokens: int, cudagraph_size: int, dp_size: int, dp_rank: int
17 -) -> tuple[torch.Tensor, torch.Tensor]:
16 + num_tokens: int,
17 + cudagraph_size: int,
18 + cudagraph_runtime_mode: int,
19 + dp_size: int,
20 + dp_rank: int,
21 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
22 assert dp_size > 1
23 # Use CPU group to avoid CPU-GPU synchronization.
24 group = get_dp_group().cpu_group
21 - tensor = torch.zeros(2, dp_size, dtype=torch.int32, device="cpu")
25 + tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu")
26 tensor[0][dp_rank] = num_tokens
27 tensor[1][dp_rank] = cudagraph_size
28 + tensor[2][dp_rank] = cudagraph_runtime_mode
29 dist.all_reduce(tensor, group=group)
25 - return tensor[0], tensor[1]
30 + return tensor[0], tensor[1], tensor[2]
31
32
33 def get_cudagraph_and_dp_padding(
29 - num_tokens: int, cudagraph_size: int | None, dp_size: int, dp_rank: int
30 -) -> tuple[bool, int, torch.Tensor | None]:
34 + num_tokens: int,
35 + cudagraph_size: int | None,
36 + cudagraph_runtime_mode: int,
37 + dp_size: int,
38 + dp_rank: int,
39 +) -> tuple[int, torch.Tensor | None, int]:
40 if dp_size == 1:
41 if cudagraph_size is not None:
33 - return True, cudagraph_size, None
42 + return cudagraph_size, None, cudagraph_runtime_mode
43 else:
35 - return False, num_tokens, None
44 + return num_tokens, None, cudagraph_runtime_mode
45
46 + # Convert None to -1 for sync (indicates no cudagraph available)
47 if num_tokens == 0:
48 cudagraph_size = 0
49 elif cudagraph_size is None:
50 cudagraph_size = -1
41 - num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
42 - num_tokens, cudagraph_size, dp_size, dp_rank
51 +
52 + num_tokens_across_dp, cudagraph_size_across_dp, cudagraph_mode_across_dp = (
53 + get_batch_metadata_across_dp(
54 + num_tokens, cudagraph_size, cudagraph_runtime_mode, dp_size, dp_rank
55 + )
56 )
57 if torch.all(num_tokens_across_dp == 0).item():
58 # All ranks have zero tokens to run.
46 - return False, 0, None
59 + return 0, None, 0
60 +
61 + # Synchronize cudagraph_runtime_mode across ranks by taking the minimum.
62 + synced_cudagraph_mode = int(cudagraph_mode_across_dp.min().item())
63 + # Check if all ranks have valid cudagraph_size.
64 + all_have_cudagraph = torch.all(cudagraph_size_across_dp != -1).item()
65
48 - if torch.all(cudagraph_size_across_dp != -1).item():
49 - # All ranks use CUDA graph or have zero tokens.
50 - # Use CUDA graph for all ranks.
51 - # Pad all ranks to the maximum CUDA graph size.
66 + if synced_cudagraph_mode != 0 and all_have_cudagraph:
67 + # All ranks use cudagraph. Pad to max cudagraph_size.
68 max_cudagraph_size = int(cudagraph_size_across_dp.max().item())
69 num_tokens_across_dp[:] = max_cudagraph_size
54 - return True, max_cudagraph_size, num_tokens_across_dp
70 + return max_cudagraph_size, num_tokens_across_dp, synced_cudagraph_mode
71 else:
56 - # Some ranks do not use CUDA graph. Use eager mode for all ranks.
57 - # No padding is needed except for ranks that have no tokens to run.
72 + # Fall back to eager mode (no cudagraph).
73 + # Either some rank doesn't have cudagraph size or mode is NONE.
74 + synced_cudagraph_mode = 0
75 num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
76 num_tokens_after_padding = int(num_tokens_across_dp[dp_rank].item())
60 - return False, num_tokens_after_padding, num_tokens_across_dp
77 + return num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode
vllm/v1/worker/gpu/model_runner.py
+26 -14
@@ -15,7 +15,7 @@ from vllm.distributed.parallel_state import (
15 get_pp_group,
16 prepare_communication_buffer_for_model,
17 )
18 -from vllm.forward_context import set_forward_context
18 +from vllm.forward_context import BatchDescriptor, set_forward_context
19 from vllm.logger import init_logger
20 from vllm.model_executor.model_loader import get_model_loader
21 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -140,7 +140,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
140 self.do_spec_decode = False
141 self.num_speculative_steps = 0
142 self.speculator = None
143 -
143 self.req_states = RequestState(
144 max_num_reqs=self.max_num_reqs,
145 max_model_len=self.max_model_len,
@@ -458,6 +457,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
457 block_tables=self.block_tables,
458 attn_metadata_builders=self.attn_metadata_builders,
459 kv_cache_config=self.kv_cache_config,
460 + has_lora=self.lora_config is not None,
461 )
462 if self.do_spec_decode:
463 self.speculator.capture_model()
@@ -884,19 +884,26 @@ class GPUModelRunner(LoRAModelRunnerMixin):
884 empty_output = self.kv_connector.no_forward(scheduler_output)
885 return empty_output
886
887 - # Get the CUDA graph size. None means no CUDA graph is used.
888 - cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
889 - scheduler_output.total_num_scheduled_tokens,
890 - scheduler_output.num_scheduled_tokens.values(),
887 + # Get local cudagraph mode and size.
888 + local_cudagraph_mode, local_cudagraph_size = (
889 + self.cudagraph_manager.get_cudagraph_runtime_mode(
890 + num_reqs=len(scheduler_output.num_scheduled_tokens),
891 + num_tokens=scheduler_output.total_num_scheduled_tokens,
892 + max_query_len=max(scheduler_output.num_scheduled_tokens.values()),
893 + )
894 )
892 - use_cudagraph, num_tokens_after_padding, num_tokens_across_dp = (
895 +
896 + # DP sync: num_tokens + cudagraph_size + cudagraph_mode
897 + num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode = (
898 get_cudagraph_and_dp_padding(
899 scheduler_output.total_num_scheduled_tokens,
895 - cudagraph_size,
900 + local_cudagraph_size,
901 + local_cudagraph_mode.value,
902 self.parallel_config.data_parallel_size,
903 self.parallel_config.data_parallel_rank,
904 )
905 )
906 + cudagraph_runtime_mode = CUDAGraphMode(synced_cudagraph_mode)
907 if num_tokens_after_padding == 0:
908 # All DP ranks have zero tokens to run.
909 empty_output = self.kv_connector.no_forward(scheduler_output)
@@ -946,16 +953,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
953 # FIXME(woosuk): Fix warmup for LoRA.
954
955 # Run model.
949 - if use_cudagraph:
950 - # Run CUDA graph.
956 + if cudagraph_runtime_mode == CUDAGraphMode.FULL:
957 + # Use explicit cudagraph replay for FULL mode.
958 # NOTE(woosuk): Here, we don't need to pass the input tensors,
959 # because they are already copied to the CUDA graph input buffers.
960 self.kv_connector.pre_forward(scheduler_output)
954 - hidden_states = self.cudagraph_manager.run(
961 + hidden_states = self.cudagraph_manager.run_fullgraph(
962 input_batch.num_tokens_after_padding
963 )
964 else:
958 - # Run PyTorch model in eager mode.
965 + # For piecewise and eager mode, just call model().
966 positions = input_batch.positions
967 if self.uses_mrope:
968 assert input_batch.mrope_positions is not None
@@ -970,13 +977,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
977 inputs_embeds = None
978 assert intermediate_tensors is not None
979
980 + batch_descriptor = BatchDescriptor(
981 + num_tokens=input_batch.num_tokens_after_padding,
982 + has_lora=self.lora_config is not None,
983 + )
984 +
985 with set_forward_context(
986 input_batch.attn_metadata,
987 self.vllm_config,
988 num_tokens=input_batch.num_tokens_after_padding,
977 - # TODO(woosuk): Support piecewise CUDA graph.
978 - cudagraph_runtime_mode=CUDAGraphMode.NONE,
989 + cudagraph_runtime_mode=cudagraph_runtime_mode,
990 num_tokens_across_dp=num_tokens_across_dp,
991 + batch_descriptor=batch_descriptor,
992 slot_mapping=input_batch.slot_mappings,
993 ):
994 self.kv_connector.pre_forward(scheduler_output)
vllm/v1/worker/gpu/spec_decode/eagle.py
+27 -9
@@ -7,7 +7,7 @@ import torch.nn as nn
7
8 from vllm.config import VllmConfig
9 from vllm.config.compilation import CUDAGraphMode
10 -from vllm.forward_context import set_forward_context
10 +from vllm.forward_context import BatchDescriptor, set_forward_context
11 from vllm.logger import init_logger
12 from vllm.model_executor.model_loader import get_model
13 from vllm.triton_utils import tl, triton
@@ -103,14 +103,17 @@ class EagleSpeculator:
103 attn_metadata: dict[str, Any] | None,
104 slot_mappings: dict[str, torch.Tensor] | None,
105 num_tokens_across_dp: torch.Tensor | None,
106 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
107 ) -> tuple[torch.Tensor, torch.Tensor]:
108 + batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
109 with set_forward_context(
110 attn_metadata,
111 self.vllm_config,
112 num_tokens=num_tokens,
111 - cudagraph_runtime_mode=CUDAGraphMode.NONE,
113 + cudagraph_runtime_mode=cudagraph_runtime_mode,
114 num_tokens_across_dp=num_tokens_across_dp,
115 slot_mapping=slot_mappings,
116 + batch_descriptor=batch_descriptor,
117 ):
118 ret_hidden_states = self.model(
119 input_ids=self.input_buffers.input_ids[:num_tokens],
@@ -127,9 +130,11 @@ class EagleSpeculator:
130 def generate_draft(
131 self,
132 num_reqs: int,
133 + num_tokens_padded: int,
134 attn_metadata: dict[str, Any],
135 slot_mappings: dict[str, torch.Tensor],
136 num_tokens_across_dp: torch.Tensor | None,
137 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
138 ) -> None:
139 pos = self.input_buffers.positions[:num_reqs]
140 query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -137,8 +142,14 @@ class EagleSpeculator:
142 for step in range(1, self.num_speculative_steps):
143 # Run the eagle model.
144 last_hidden_states, hidden_states = self.run_model(
140 - num_reqs, attn_metadata, slot_mappings, num_tokens_across_dp
145 + num_tokens_padded,
146 + attn_metadata,
147 + slot_mappings,
148 + num_tokens_across_dp,
149 + cudagraph_runtime_mode,
150 )
151 + last_hidden_states = last_hidden_states[:num_reqs]
152 + hidden_states = hidden_states[:num_reqs]
153 logits = self.model.compute_logits(last_hidden_states)
154
155 # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
@@ -283,12 +294,14 @@ class EagleSpeculator:
294 )
295
296 cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
286 - if cudagraph_size is not None:
287 - # Run CUDA graph.
288 - self.cudagraph_manager.run(cudagraph_size)
297 + cudagraph_mode = self.cudagraph_manager.cudagraph_mode
298 + if cudagraph_size is not None and cudagraph_mode == CUDAGraphMode.FULL:
299 + # Run full CUDA graph.
300 + self.cudagraph_manager.run_fullgraph(cudagraph_size)
301 return self.draft_tokens[:num_reqs]
302
291 - # Run eager mode.
303 + # Run eager or piecewise CUDA graph.
304 + num_tokens_padded = cudagraph_size if cudagraph_size is not None else num_reqs
305 query_start_loc_cpu = torch.arange(
306 num_reqs + 1, dtype=torch.int32, device="cpu"
307 )
@@ -312,8 +325,13 @@ class EagleSpeculator:
325 slot_mappings, self.kv_cache_config
326 )
327 self.generate_draft(
315 - num_reqs, attn_metadata, slot_mappings_by_layer, num_tokens_across_dp=None
316 - ) # FIXME
328 + num_reqs,
329 + num_tokens_padded,
330 + attn_metadata,
331 + slot_mappings_by_layer,
332 + num_tokens_across_dp=None, # FIXME
333 + cudagraph_runtime_mode=cudagraph_mode,
334 + )
335 return self.draft_tokens[:num_reqs]
336
337
vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py
+75 -8
@@ -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 Callable
4 +from typing import Any
5
6 import torch
7
@@ -31,16 +32,17 @@ class EagleCudaGraphManager:
32 self.compilation_config = vllm_config.compilation_config
33 assert self.compilation_config is not None
34
34 - self.cudagraph_mode = self.compilation_config.cudagraph_mode
35 - if self.cudagraph_mode == CUDAGraphMode.FULL:
36 - # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
37 - self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
35 + # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
36 + self.cudagraph_mode = self.compilation_config.cudagraph_mode.decode_mode()
37
39 - self.cudagraph_sizes = get_cudagraph_sizes(
38 + # only need to capture uniform decode cudagraph sizes (the 2nd return value)
39 + _, self.cudagraph_sizes = get_cudagraph_sizes(
40 self.compilation_config.cudagraph_capture_sizes,
41 self.max_num_reqs,
42 self.max_num_tokens,
43 self.cudagraph_mode,
44 + uniform_decode_query_len=1,
45 + uniform_decode_cudagraph=True,
46 )
47
48 self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
@@ -54,12 +56,21 @@ class EagleCudaGraphManager:
56 def capture_graph(
57 self,
58 num_tokens: int,
59 + capture_cg_mode: CUDAGraphMode,
60 generate_fn: Callable,
61 input_buffers: InputBuffers,
62 block_tables: BlockTables,
63 attn_metadata_builders: list[AttentionMetadataBuilder],
64 kv_cache_config: KVCacheConfig,
65 ) -> None:
66 + assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
67 + f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
68 + )
69 + if capture_cg_mode == CUDAGraphMode.PIECEWISE:
70 + capture_fn = self._capture_piecewise_graph
71 + else:
72 + capture_fn = self._capture_full_graph
73 +
74 num_reqs = min(num_tokens, self.max_num_reqs)
75 attn_metadata, slot_mappings = prepare_inputs_to_capture(
76 num_reqs,
@@ -69,19 +80,70 @@ class EagleCudaGraphManager:
80 attn_metadata_builders,
81 self.max_model_len,
82 kv_cache_config,
83 + uniform_decode_query_len=1,
84 )
85 num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
86
87 # Warm up.
76 - generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
88 + generate_fn(
89 + num_reqs,
90 + num_tokens,
91 + attn_metadata,
92 + slot_mappings,
93 + num_tokens_across_dp,
94 + CUDAGraphMode.NONE,
95 + )
96
97 # Capture the graph.
98 + capture_fn(
99 + num_reqs=num_reqs,
100 + num_tokens=num_tokens,
101 + generate_fn=generate_fn,
102 + attn_metadata=attn_metadata,
103 + slot_mappings=slot_mappings,
104 + num_tokens_across_dp=num_tokens_across_dp,
105 + )
106 +
107 + def _capture_full_graph(
108 + self,
109 + num_reqs: int,
110 + num_tokens: int,
111 + generate_fn: Callable,
112 + attn_metadata: dict[str, Any],
113 + slot_mappings: dict[str, torch.Tensor],
114 + num_tokens_across_dp: torch.Tensor,
115 + ) -> None:
116 assert num_tokens not in self.graphs
117 graph = torch.cuda.CUDAGraph()
118 with torch.cuda.graph(graph, self.pool):
82 - generate_fn(num_tokens, attn_metadata, slot_mappings, num_tokens_across_dp)
119 + generate_fn(
120 + num_reqs,
121 + num_tokens,
122 + attn_metadata,
123 + slot_mappings,
124 + num_tokens_across_dp,
125 + CUDAGraphMode.NONE,
126 + )
127 self.graphs[num_tokens] = graph
128
129 + def _capture_piecewise_graph(
130 + self,
131 + num_reqs: int,
132 + num_tokens: int,
133 + generate_fn: Callable,
134 + attn_metadata: dict[str, Any],
135 + slot_mappings: dict[str, torch.Tensor],
136 + num_tokens_across_dp: torch.Tensor,
137 + ) -> None:
138 + generate_fn(
139 + num_reqs,
140 + num_tokens,
141 + attn_metadata,
142 + slot_mappings,
143 + num_tokens_across_dp,
144 + CUDAGraphMode.PIECEWISE,
145 + )
146 +
147 @torch.inference_mode()
148 def capture(
149 self,
@@ -91,10 +153,15 @@ class EagleCudaGraphManager:
153 attn_metadata_builders: list[AttentionMetadataBuilder],
154 kv_cache_config: KVCacheConfig,
155 ) -> None:
156 + if self.cudagraph_mode == CUDAGraphMode.NONE:
157 + return
158 +
159 capture_graphs(
160 self.cudagraph_sizes,
161 self.device,
162 self.capture_graph,
163 + capture_cudagraph_mode=self.cudagraph_mode,
164 + desc=f"Capturing eagle CUDA graphs ({self.cudagraph_mode.name})",
165 generate_fn=generate_fn,
166 input_buffers=input_buffers,
167 block_tables=block_tables,
@@ -102,6 +169,6 @@ class EagleCudaGraphManager:
169 kv_cache_config=kv_cache_config,
170 )
171
105 - def run(self, num_tokens: int) -> None:
172 + def run_fullgraph(self, num_tokens: int) -> None:
173 assert num_tokens in self.graphs
174 self.graphs[num_tokens].replay()