@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 483463f735
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 7, 00:15 UTC

[MRV2] Extensible CG dispatch rework (#35959)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Mar 9, 2026 at 16:58 UTC 483463f735c41c36a41431044fa537dc4c81fc3c
9 files changed +545 -636
vllm/config/compilation.py
+3
@@ -97,6 +97,9 @@ class CUDAGraphMode(enum.Enum):
97 def __str__(self) -> str:
98 return self.name
99
100 + def __bool__(self) -> bool:
101 + return self != CUDAGraphMode.NONE
102 +
103
104 @config
105 class PassConfig:
vllm/v1/worker/gpu/block_table.py
+24 -8
@@ -104,19 +104,24 @@ class BlockTables:
104 self.num_blocks.copy_to_uva()
105
106 def gather_block_tables(
107 - self, idx_mapping: torch.Tensor
107 + self,
108 + idx_mapping: torch.Tensor,
109 + num_reqs_padded: int,
110 ) -> tuple[torch.Tensor, ...]:
111 num_reqs = idx_mapping.shape[0]
110 - _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs)](
112 + # Launch kernel with num_reqs_padded to fuse zeroing of padded rows.
113 + _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)](
114 idx_mapping,
115 self.block_table_ptrs,
116 self.input_block_table_ptrs,
117 self.block_table_strides,
118 self.num_blocks.gpu,
119 self.num_blocks.gpu.stride(0),
120 + num_reqs,
121 + self.input_block_tables[0].shape[1], # max_num_blocks
122 BLOCK_SIZE=1024, # type: ignore
123 )
119 - return tuple(block_table[:num_reqs] for block_table in self.input_block_tables)
124 + return tuple(bt[:num_reqs_padded] for bt in self.input_block_tables)
125
126 def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
127 # NOTE(woosuk): The output may be used for CUDA graph capture.
@@ -130,6 +135,7 @@ class BlockTables:
135 idx_mapping: torch.Tensor,
136 query_start_loc: torch.Tensor,
137 positions: torch.Tensor,
138 + num_tokens_padded: int,
139 ) -> torch.Tensor:
140 num_reqs = idx_mapping.shape[0]
141 num_tokens = positions.shape[0]
@@ -151,7 +157,7 @@ class BlockTables:
157 PAD_ID=PAD_SLOT_ID,
158 TRITON_BLOCK_SIZE=1024, # type: ignore
159 )
154 - return self.slot_mappings[:, :num_tokens]
160 + return self.slot_mappings[:, :num_tokens_padded]
161
162 def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
163 # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.
@@ -173,21 +179,31 @@ def _gather_block_tables_kernel(
179 block_table_strides, # [num_kv_cache_groups]
180 num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs]
181 num_blocks_stride,
182 + num_reqs, # actual number of requests (for padding)
183 + max_num_blocks, # stride for zeroing padded rows
184 BLOCK_SIZE: tl.constexpr,
185 ):
186 # kv cache group id
187 group_id = tl.program_id(0)
188 batch_idx = tl.program_id(1)
181 - req_idx = tl.load(batch_idx_to_req_idx + batch_idx)
189
190 + stride = tl.load(block_table_strides + group_id)
191 + dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
192 + dst_row_ptr = dst_block_table_ptr + batch_idx * stride
193 +
194 + if batch_idx >= num_reqs:
195 + # Zero out padded rows.
196 + for i in tl.range(0, max_num_blocks, BLOCK_SIZE):
197 + offset = i + tl.arange(0, BLOCK_SIZE)
198 + tl.store(dst_row_ptr + offset, 0, mask=offset < max_num_blocks)
199 + return
200 +
201 + req_idx = tl.load(batch_idx_to_req_idx + batch_idx)
202 group_num_blocks_ptr = num_blocks_ptr + group_id * num_blocks_stride
203 num_blocks = tl.load(group_num_blocks_ptr + req_idx)
204
186 - stride = tl.load(block_table_strides + group_id)
205 src_block_table_ptr = _load_ptr(src_block_table_ptrs + group_id, tl.int32)
206 src_row_ptr = src_block_table_ptr + req_idx * stride
189 - dst_block_table_ptr = _load_ptr(dst_block_table_ptrs + group_id, tl.int32)
190 - dst_row_ptr = dst_block_table_ptr + batch_idx * stride
207
208 for i in tl.range(0, num_blocks, BLOCK_SIZE):
209 offset = i + tl.arange(0, BLOCK_SIZE)
vllm/v1/worker/gpu/cudagraph_utils.py
+285 -316
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections import defaultdict
4 from collections.abc import Callable
5 +from dataclasses import dataclass
6 from typing import Any
7
8 import torch
@@ -11,235 +13,260 @@ from vllm.config import VllmConfig
13 from vllm.config.compilation import CUDAGraphMode
14 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
15 from vllm.forward_context import BatchDescriptor, set_forward_context
16 +from vllm.logger import init_logger
17 from vllm.model_executor.offloader.base import get_offloader
15 -from vllm.utils.math_utils import cdiv
18 +from vllm.platforms import current_platform
19 from vllm.v1.kv_cache_interface import KVCacheConfig
20 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
21 from vllm.v1.worker.gpu.block_table import BlockTables
22 from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
20 -from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_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.utils import AttentionGroup
26
27 +logger = init_logger(__name__)
28 +
29 +
30 +@dataclass(frozen=True)
31 +class BatchExecutionDescriptor:
32 + """Describes the shape of the batch and CG mode to run; this is used to make shape
33 + matches between the capture and runtime."""
34 +
35 + cg_mode: CUDAGraphMode
36 + num_tokens: int
37 + num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs)
38 + uniform_token_count: int | None = None
39 +
40 +
41 +def _is_compatible(
42 + desc: BatchExecutionDescriptor,
43 + num_reqs: int,
44 + num_tokens: int,
45 + uniform_token_count: int | None,
46 +) -> bool:
47 + # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
48 + # desc.num_reqs=None means no request padding needed (PIECEWISE)
49 + return (
50 + (
51 + desc.uniform_token_count is None
52 + or desc.uniform_token_count == uniform_token_count
53 + )
54 + and (desc.num_reqs is None or desc.num_reqs >= num_reqs)
55 + and desc.num_tokens >= num_tokens
56 + )
57 +
58 +
59 +def get_uniform_token_count(
60 + num_reqs: int,
61 + num_tokens: int,
62 + max_query_len: int,
63 +) -> int | None:
64 + """
65 + Return the uniform token count if batch is uniform, else None.
66 + A batch is uniform if all requests have the same number of tokens.
67 + """
68 + if (max_query_len == num_tokens // num_reqs) and (
69 + num_tokens == max_query_len * num_reqs
70 + ):
71 + return max_query_len
72 + return None
73 +
74
75 class CudaGraphManager:
76 def __init__(
77 self,
78 vllm_config: VllmConfig,
30 - use_aux_hidden_state_outputs: bool,
79 device: torch.device,
80 + cudagraph_mode: CUDAGraphMode,
81 + decode_query_len: int,
82 ):
83 self.vllm_config = vllm_config
34 - self.scheduler_config = vllm_config.scheduler_config
35 - self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
84 self.device = device
37 -
38 - self.max_model_len = vllm_config.model_config.max_model_len
39 - self.max_num_reqs = self.scheduler_config.max_num_seqs
40 - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
85 + self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
86 + self.compilation_config = vllm_config.compilation_config
87 + assert self.compilation_config is not None
88 + self.cudagraph_mode = cudagraph_mode
89 + self.decode_query_len = decode_query_len
90 self.dp_size = vllm_config.parallel_config.data_parallel_size
91
43 - self.uniform_decode_query_len = 1
44 - spec_config = vllm_config.speculative_config
45 - if spec_config is not None:
46 - self.uniform_decode_query_len += spec_config.num_speculative_tokens
92 + self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
93 + self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
94
48 - self.compilation_config = vllm_config.compilation_config
49 - assert self.compilation_config is not None
50 - self.cudagraph_mode = self.compilation_config.cudagraph_mode
95 + self._graphs_captured = False
96 + self._candidates: list[list[BatchExecutionDescriptor]] = []
97 + self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
98 + self._init_candidates()
99
52 - use_uniform_decode_cudagraph = (
53 - self.cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
54 - and self.cudagraph_mode.separate_routine()
55 - )
56 - self.cudagraph_sizes, self.uniform_decode_cudagraph_sizes = get_cudagraph_sizes(
57 - self.compilation_config.cudagraph_capture_sizes,
58 - self.max_num_reqs,
59 - self.max_num_tokens,
60 - self.cudagraph_mode,
61 - self.uniform_decode_query_len,
62 - use_uniform_decode_cudagraph,
63 - )
100 + def _init_candidates(self) -> None:
101 + """Build priority-ordered candidate lists for each token count."""
102 + capture_sizes = self.compilation_config.cudagraph_capture_sizes
103 + if not (self.cudagraph_mode and capture_sizes):
104 + return
105
65 - self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
66 - self.pool = None
67 - if self.cudagraph_mode != CUDAGraphMode.NONE:
68 - self.pool = torch.cuda.graph_pool_handle()
69 - self.hidden_states: torch.Tensor | None = None
70 - self.aux_hidden_states: list[torch.Tensor] = []
106 + capture_sizes = sorted(capture_sizes)
107 + max_decode_tokens = self.max_num_reqs * self.decode_query_len
108 + decode_mode = self.cudagraph_mode.decode_mode()
109 + mixed_mode = self.cudagraph_mode.mixed_mode()
110 + separate_decode_routine = self.cudagraph_mode.separate_routine()
111 +
112 + descs_by_token_count = defaultdict(list)
113 + descs_by_mode = defaultdict(list)
114 +
115 + for num_tokens in capture_sizes:
116 + # Capture uniform decode specfifc graphs if required
117 + # (i.e. separate decode routine)
118 + if (
119 + separate_decode_routine
120 + and decode_mode
121 + and self.decode_query_len <= num_tokens <= max_decode_tokens
122 + ):
123 + desc = BatchExecutionDescriptor(
124 + cg_mode=decode_mode,
125 + num_tokens=num_tokens,
126 + num_reqs=num_tokens // self.decode_query_len,
127 + uniform_token_count=self.decode_query_len,
128 + )
129 + descs_by_mode[decode_mode].append(desc)
130 + descs_by_token_count[num_tokens].append(desc)
131 +
132 + if mixed_mode:
133 + # for PIECEWISE graphs there is no limit on requests when replaying
134 + # i.e. no request padding is needed
135 + # so we leave it as None
136 + num_reqs = (
137 + min(num_tokens, self.max_num_reqs)
138 + if mixed_mode == CUDAGraphMode.FULL
139 + else None
140 + )
141 + desc = BatchExecutionDescriptor(
142 + cg_mode=mixed_mode,
143 + num_tokens=num_tokens,
144 + num_reqs=num_reqs,
145 + )
146 + descs_by_mode[mixed_mode].append(desc)
147 + descs_by_token_count[num_tokens].append(desc)
148 +
149 + if not descs_by_token_count:
150 + return
151 +
152 + sorted_padded = sorted(descs_by_token_count.keys())
153 + self._candidates = [[] for _ in range(sorted_padded[-1] + 1)]
154 +
155 + current_range_start = 0
156 + for cg_size in sorted_padded:
157 + for i in range(current_range_start, cg_size + 1):
158 + self._candidates[i] = descs_by_token_count[cg_size]
159 + current_range_start = cg_size + 1
160 +
161 + for mode, descs in descs_by_mode.items():
162 + descs.sort(key=lambda d: d.num_tokens, reverse=True)
163 + self._capture_descs[mode] = descs
164
165 def needs_capture(self) -> bool:
73 - return len(self.cudagraph_sizes) > 0
74 -
75 - def get_cudagraph_size(
76 - self, num_tokens: int, uniform_decode: bool = False
77 - ) -> int | None:
78 - if uniform_decode and self.uniform_decode_cudagraph_sizes:
79 - return self.uniform_decode_cudagraph_sizes.get(num_tokens)
80 - return self.cudagraph_sizes.get(num_tokens)
166 + return len(self._capture_descs) > 0
167
82 - def capture_graph(
168 + @torch.inference_mode()
169 + def capture(
170 self,
84 - num_tokens: int,
85 - capture_cg_mode: CUDAGraphMode,
86 - model: nn.Module,
87 - model_state: ModelState,
88 - input_buffers: InputBuffers,
89 - block_tables: BlockTables,
90 - attn_groups: list[list[AttentionGroup]],
91 - kv_cache_config: KVCacheConfig,
92 - has_lora: bool = False,
93 - uniform_decode: bool = False,
171 + create_forward_fn: Callable[
172 + [BatchExecutionDescriptor], Callable[[CUDAGraphMode], None]
173 + ],
174 + progress_bar_desc: str = "Capturing CUDA graphs",
175 ) -> None:
95 - # select and check capture function
96 - assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
97 - f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
98 - )
99 - if capture_cg_mode == CUDAGraphMode.PIECEWISE:
100 - capture_fn = self._capture_piecewise_graph
101 - else:
102 - capture_fn = self._capture_full_graph
103 - # prepare inputs
104 - if uniform_decode:
105 - num_reqs = min(
106 - cdiv(num_tokens, self.uniform_decode_query_len),
107 - self.max_num_reqs,
108 - )
109 - else:
110 - num_reqs = min(num_tokens, self.max_num_reqs)
111 -
112 - model_inputs = {
113 - "input_ids": input_buffers.input_ids[:num_tokens],
114 - "positions": input_buffers.positions[:num_tokens],
115 - # NOTE: Values returned by `prepare_dummy_inputs` will override the
116 - # default values above.
117 - **model_state.prepare_dummy_inputs(num_reqs, num_tokens),
118 - }
119 -
120 - attn_metadata, slot_mappings = prepare_inputs_to_capture(
121 - num_reqs,
122 - num_tokens,
123 - model_state,
124 - input_buffers,
125 - block_tables,
126 - attn_groups,
127 - kv_cache_config,
128 - )
129 - num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
130 -
131 - # Warm up.
132 - with set_forward_context(
133 - attn_metadata,
134 - self.vllm_config,
135 - num_tokens=num_tokens,
136 - cudagraph_runtime_mode=CUDAGraphMode.NONE,
137 - num_tokens_across_dp=num_tokens_across_dp,
138 - slot_mapping=slot_mappings,
139 - ):
140 - model_output = model(**model_inputs)
141 - if self.use_aux_hidden_state_outputs:
142 - hidden_states, aux_hidden_states = model_output
143 - else:
144 - hidden_states = model_output
145 - aux_hidden_states = None
146 -
147 - # Allocate output buffers if not already done.
148 - if self.hidden_states is None:
149 - self.hidden_states = torch.empty_like(hidden_states)
150 - if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
151 - self.aux_hidden_states = [torch.empty_like(x) for x in aux_hidden_states]
152 -
153 - capture_fn(
154 - num_tokens=num_tokens,
155 - num_reqs=num_reqs,
156 - model=model,
157 - model_inputs=model_inputs,
158 - num_tokens_across_dp=num_tokens_across_dp,
159 - attn_metadata=attn_metadata,
160 - slot_mappings=slot_mappings,
161 - has_lora=has_lora,
162 - )
163 -
164 - def _capture_full_graph(
176 + """Capture CUDA graphs.
177 +
178 + Args:
179 + create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
180 + returns a function that runs forward with a given CUDAGraphMode.
181 + """
182 + with graph_capture(device=self.device):
183 + # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
184 + # activations so FULL activations should fit in already allocated
185 + # buffers in the graph pool.
186 + for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]:
187 + if mode not in self._capture_descs:
188 + continue
189 +
190 + descs = self._capture_descs[mode]
191 + if is_global_first_rank():
192 + descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})")
193 + for desc in descs:
194 + # Prepare inputs and get forward function
195 + forward_fn = create_forward_fn(desc)
196 +
197 + # Warmup
198 + forward_fn(CUDAGraphMode.NONE)
199 +
200 + # Capture
201 + logger.debug(
202 + "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
203 + )
204 + if desc.cg_mode == CUDAGraphMode.PIECEWISE:
205 + forward_fn(CUDAGraphMode.PIECEWISE)
206 + else:
207 + assert desc not in self.graphs, (
208 + f"Graph already captured for {desc}"
209 + )
210 + graph = torch.cuda.CUDAGraph()
211 + # Sync offloader's copy stream before capture.
212 + # Ensure any pre-capture prefetches from offloader are complete.
213 + get_offloader().sync_prev_onload()
214 + with torch.cuda.graph(graph, self.pool):
215 + forward_fn(CUDAGraphMode.NONE)
216 + # Join offloader's copy stream after forward to avoid
217 + # unjoined stream error. The last layer's start_prefetch
218 + # forks copy_stream, but wait_prefetch only happens in
219 + # the next forward pass.
220 + get_offloader().join_after_forward()
221 + self.graphs[desc] = graph
222 + self._graphs_captured = True
223 +
224 + def dispatch(
225 self,
166 - num_tokens: int,
226 num_reqs: int,
168 - model: nn.Module,
169 - model_inputs: dict[str, torch.Tensor | None],
170 - num_tokens_across_dp: torch.Tensor,
171 - attn_metadata: dict[str, Any] | None,
172 - slot_mappings: dict[str, torch.Tensor] | None,
173 - has_lora: bool = False,
174 - ) -> None:
175 - assert attn_metadata is not None
176 - # Capture the graph.
177 - assert num_tokens not in self.graphs
178 - graph = torch.cuda.CUDAGraph()
227 + num_tokens: int,
228 + uniform_token_count: int | None,
229 + ) -> BatchExecutionDescriptor:
230 + """Find matching cudagraph descriptor from priority-ordered candidates."""
231 + if self._graphs_captured and 0 < num_tokens < len(self._candidates):
232 + for desc in self._candidates[num_tokens]:
233 + if _is_compatible(desc, num_reqs, num_tokens, uniform_token_count):
234 + return desc
235 + return BatchExecutionDescriptor(
236 + cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, num_reqs=num_reqs
237 + )
238
180 - # Sync offloader's copy stream before capture.
181 - # Ensure any pre-capture prefetches from offloader are complete.
239 + def run_fullgraph(self, desc: BatchExecutionDescriptor):
240 + """Replay a captured FULL cudagraph."""
241 + assert desc.cg_mode == CUDAGraphMode.FULL, (
242 + f"Expected FULL mode, got {desc.cg_mode}"
243 + )
244 + assert desc in self.graphs, f"No cudagraph for {desc}"
245 + # Sync offloader before replay - needed when transitioning from
246 + # eager/piecewise to full cudagraph (e.g., prefill → decode).
247 + # The previous eager iteration's start_prefetch may have queued
248 + # H2D copies on copy_stream that the graph's captured events
249 + # cannot see. Without this, replay could overwrite static buffers
250 + # while those copies are still in flight.
251 get_offloader().sync_prev_onload()
252 + self.graphs[desc].replay()
253 +
254 +
255 +class ModelCudaGraphManager(CudaGraphManager):
256 + """CudaGraphManager with model-specific capture and hidden state management."""
257
184 - with (
185 - set_forward_context(
186 - attn_metadata=attn_metadata,
187 - vllm_config=self.vllm_config,
188 - num_tokens=num_tokens,
189 - cudagraph_runtime_mode=CUDAGraphMode.NONE,
190 - num_tokens_across_dp=num_tokens_across_dp,
191 - slot_mapping=slot_mappings,
192 - ),
193 - torch.cuda.graph(graph, self.pool),
194 - ):
195 - model_output = model(**model_inputs)
196 -
197 - # Join offloader's copy stream after forward to avoid unjoined
198 - # stream error. The last layer's start_prefetch forks copy_stream,
199 - # but wait_prefetch only happens in the next forward pass.
200 - get_offloader().join_after_forward()
201 -
202 - if self.use_aux_hidden_state_outputs:
203 - hidden_states, aux_hidden_states = model_output
204 - else:
205 - hidden_states = model_output
206 - aux_hidden_states = None
207 -
208 - # Copy outputs to the output buffers.
209 - assert self.hidden_states is not None
210 - self.hidden_states[:num_tokens] = hidden_states
211 - if self.use_aux_hidden_state_outputs:
212 - for i, aux_hidden in enumerate(aux_hidden_states):
213 - self.aux_hidden_states[i][:num_tokens] = aux_hidden
214 - self.graphs[num_tokens] = graph
215 -
216 - def _capture_piecewise_graph(
258 + def __init__(
259 self,
218 - num_tokens: int,
219 - num_reqs: int,
220 - model: nn.Module,
221 - model_inputs: dict[str, torch.Tensor | None],
222 - num_tokens_across_dp: torch.Tensor,
223 - attn_metadata: dict[str, Any] | None,
224 - slot_mappings: dict[str, torch.Tensor] | None,
225 - has_lora: bool = False,
226 - ) -> None:
227 - # create batch descriptor for piecewise cudagraph dispatch key
228 - batch_descriptor = BatchDescriptor(num_tokens=num_tokens, has_lora=has_lora)
229 -
230 - # Capture run - CUDAGraphWrapper inside torch.compile will auto capture.
231 - with set_forward_context(
232 - attn_metadata=None, # piecewise no need attn_metadata
233 - vllm_config=self.vllm_config,
234 - num_tokens=num_tokens,
235 - cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
236 - num_tokens_across_dp=num_tokens_across_dp,
237 - batch_descriptor=batch_descriptor,
238 - slot_mapping=slot_mappings,
239 - ):
240 - model(**model_inputs)
260 + vllm_config: VllmConfig,
261 + device: torch.device,
262 + cudagraph_mode: CUDAGraphMode,
263 + decode_query_len: int,
264 + ):
265 + super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
266 + self.hidden_states: torch.Tensor | None = None
267 + self.aux_hidden_states: list[torch.Tensor] = []
268 + self.use_aux_hidden_state_outputs = False
269
242 - @torch.inference_mode()
270 def capture(
271 self,
272 model: nn.Module,
@@ -249,139 +276,81 @@ class CudaGraphManager:
276 attn_groups: list[list[AttentionGroup]],
277 kv_cache_config: KVCacheConfig,
278 has_lora: bool = False,
279 + use_aux_hidden_state_outputs: bool = False,
280 + progress_bar_desc: str = "Capturing CUDA graphs",
281 ) -> None:
253 - common_kwargs = dict(
254 - device=self.device,
255 - capture_fn=self.capture_graph,
256 - model=model,
257 - model_state=model_state,
258 - input_buffers=input_buffers,
259 - block_tables=block_tables,
260 - attn_groups=attn_groups,
261 - kv_cache_config=kv_cache_config,
262 - has_lora=has_lora,
263 - )
282 + """Capture CUDA graphs for model forward pass."""
283 + self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
284
265 - # Phase 1: Capture for mixed prefill-decode batches if needed.
266 - mixed_mode = self.cudagraph_mode.mixed_mode()
267 - if mixed_mode != CUDAGraphMode.NONE:
268 - capture_graphs(
269 - cudagraph_sizes=self.cudagraph_sizes,
270 - capture_cudagraph_mode=mixed_mode,
271 - desc=f"Capturing CUDA graphs (mixed, {mixed_mode.name})",
272 - uniform_decode=False,
273 - **common_kwargs,
285 + def create_forward_fn(
286 + desc: BatchExecutionDescriptor,
287 + ) -> Callable[[CUDAGraphMode], None]:
288 + num_tokens = desc.num_tokens
289 + num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
290 + num_tokens_across_dp = (
291 + torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
292 + if self.dp_size > 1
293 + else None
294 )
275 -
276 - # Phase 2: Capture FULL graphs for uniform decode batches if needed.
277 - # This is only needed if we use a separate routine for decode batches
278 - # and the decode_mode is FULL.
279 - if self.uniform_decode_cudagraph_sizes:
280 - capture_graphs(
281 - cudagraph_sizes=self.uniform_decode_cudagraph_sizes,
282 - capture_cudagraph_mode=CUDAGraphMode.FULL,
283 - desc="Capturing CUDA graphs (decode, FULL)",
284 - uniform_decode=True,
285 - **common_kwargs,
295 + attn_metadata, slot_mappings = prepare_inputs_to_capture(
296 + num_reqs,
297 + num_tokens,
298 + model_state,
299 + input_buffers,
300 + block_tables,
301 + attn_groups,
302 + kv_cache_config,
303 )
304
288 - def get_cudagraph_runtime_mode(
289 - self, num_reqs: int, num_tokens: int, max_query_len: int
290 - ) -> tuple[CUDAGraphMode, int | None]:
291 - is_uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
292 - num_tokens == max_query_len * num_reqs
293 - )
294 -
295 - cudagraph_size = self.get_cudagraph_size(num_tokens, is_uniform_decode)
296 - if cudagraph_size is None:
297 - cudagraph_mode = CUDAGraphMode.NONE
298 - elif is_uniform_decode:
299 - cudagraph_mode = self.cudagraph_mode.decode_mode()
300 - else:
301 - cudagraph_mode = self.cudagraph_mode.mixed_mode()
302 -
303 - if (
304 - cudagraph_mode == CUDAGraphMode.FULL
305 - and cudagraph_size is not None
306 - and cudagraph_size not in self.graphs
307 - ):
308 - # If graph wasn't captured yet, fall back to eager.
309 - # This might happen when the dummy run is called before capture.
310 - cudagraph_mode = CUDAGraphMode.NONE
311 - cudagraph_size = None
312 - return cudagraph_mode, cudagraph_size
305 + def forward_fn(cg_mode: CUDAGraphMode) -> None:
306 + batch_descriptor = (
307 + BatchDescriptor(num_tokens=num_tokens)
308 + if cg_mode == CUDAGraphMode.PIECEWISE
309 + else None
310 + )
311 + with set_forward_context(
312 + attn_metadata if cg_mode != CUDAGraphMode.PIECEWISE else None,
313 + self.vllm_config,
314 + num_tokens=num_tokens,
315 + cudagraph_runtime_mode=cg_mode,
316 + num_tokens_across_dp=num_tokens_across_dp,
317 + slot_mapping=slot_mappings,
318 + batch_descriptor=batch_descriptor,
319 + ):
320 + model_inputs = {
321 + "input_ids": input_buffers.input_ids[:num_tokens],
322 + "positions": input_buffers.positions[:num_tokens],
323 + }
324 + model_output = model(**model_inputs)
325 + if self.use_aux_hidden_state_outputs:
326 + hidden_states, aux_hidden_states = model_output
327 + else:
328 + hidden_states = model_output
329 + aux_hidden_states = []
330 + if self.hidden_states is None:
331 + self.hidden_states = torch.empty_like(hidden_states)
332 + if self.use_aux_hidden_state_outputs and not self.aux_hidden_states:
333 + self.aux_hidden_states = [
334 + torch.empty_like(x) for x in aux_hidden_states
335 + ]
336 + self.hidden_states[:num_tokens] = hidden_states
337 + for i, aux in enumerate(aux_hidden_states):
338 + self.aux_hidden_states[i][:num_tokens] = aux
339 +
340 + return forward_fn
341 +
342 + super().capture(create_forward_fn, progress_bar_desc)
343
344 def run_fullgraph(
315 - self, num_tokens: int
345 + self, desc: BatchExecutionDescriptor
346 ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
317 - assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens"
318 - # Sync offloader before replay - needed when transitioning from
319 - # eager/piecewise to full cudagraph (e.g., prefill → decode).
320 - # The previous eager iteration's start_prefetch may have queued
321 - # H2D copies on copy_stream that the graph's captured events
322 - # cannot see. Without this, replay could overwrite static buffers
323 - # while those copies are still in flight.
324 - get_offloader().sync_prev_onload()
325 - self.graphs[num_tokens].replay()
347 + """Replay a captured FULL cudagraph and return hidden states."""
348 + super().run_fullgraph(desc)
349 assert self.hidden_states is not None
327 - hidden_states = self.hidden_states[:num_tokens]
350 + hidden_states = self.hidden_states[: desc.num_tokens]
351 if not self.use_aux_hidden_state_outputs:
352 return hidden_states
330 - return hidden_states, [x[:num_tokens] for x in self.aux_hidden_states]
331 -
332 -
333 -def get_cudagraph_sizes(
334 - capture_sizes: list[int] | None,
335 - max_num_reqs: int,
336 - max_num_tokens: int,
337 - cudagraph_mode: CUDAGraphMode,
338 - uniform_decode_query_len: int = 1,
339 - uniform_decode_cudagraph: bool = False,
340 -) -> tuple[dict[int, int], dict[int, int]]:
341 - # Support both FULL and PIECEWISE cudagraph modes
342 - if cudagraph_mode == CUDAGraphMode.NONE:
343 - return {}, {}
344 - if not capture_sizes:
345 - return {}, {}
346 -
347 - capture_sizes = sorted(capture_sizes)
348 - if not capture_sizes:
349 - return {}, {}
350 -
351 - cudagraph_sizes: dict[int, int] = {}
352 - for i in range(1, capture_sizes[-1] + 1):
353 - for x in capture_sizes:
354 - if i <= x:
355 - cudagraph_sizes[i] = x
356 - break
357 -
358 - uniform_decode_cudagraph_sizes: dict[int, int] = {}
359 - if uniform_decode_cudagraph:
360 - max_num_tokens = max_num_reqs * uniform_decode_query_len
361 - uniform_decode_cudagraph_sizes = {
362 - k: v
363 - for k, v in cudagraph_sizes.items()
364 - if v <= max_num_tokens and v >= uniform_decode_query_len
365 - }
366 - return cudagraph_sizes, uniform_decode_cudagraph_sizes
367 -
368 -
369 -def capture_graphs(
370 - cudagraph_sizes: dict[int, int],
371 - device: torch.device,
372 - capture_fn: Callable,
373 - capture_cudagraph_mode: CUDAGraphMode,
374 - desc: str = "Capturing CUDA graphs",
375 - **capture_kwargs,
376 -) -> None:
377 - # Capture larger graphs first.
378 - sizes_to_capture = sorted(set(cudagraph_sizes.values()), reverse=True)
379 - if is_global_first_rank():
380 - sizes_to_capture = tqdm(sizes_to_capture, desc=desc)
381 -
382 - with graph_capture(device=device):
383 - for size in sizes_to_capture:
384 - capture_fn(size, capture_cudagraph_mode, **capture_kwargs)
353 + return hidden_states, [x[: desc.num_tokens] for x in self.aux_hidden_states]
354
355
356 def prepare_inputs_to_capture(
vllm/v1/worker/gpu/dp_utils.py
+53 -49
@@ -1,9 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 import torch
6 import torch.distributed as dist
7
8 +from vllm.config.compilation import CUDAGraphMode
9 from vllm.distributed.parallel_state import get_dp_group
10 +from vllm.v1.worker.gpu.cudagraph_utils import (
11 + BatchExecutionDescriptor,
12 + CudaGraphManager,
13 +)
14
15
16 def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | None:
@@ -12,66 +19,63 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N
19 return torch.full((dp_size,), num_tokens, dtype=torch.int32, device="cpu")
20
21
15 -def get_batch_metadata_across_dp(
22 +def sync_cudagraph_and_dp_padding(
23 + cudagraph_manager: CudaGraphManager,
24 + desired_batch_desc: BatchExecutionDescriptor,
25 num_tokens: int,
17 - cudagraph_size: int,
18 - cudagraph_runtime_mode: int,
26 + num_reqs: int,
27 + uniform_token_count: int | None,
28 dp_size: int,
29 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.
30 +) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
31 + """
32 + Coordinates the batch descriptor and DP padding across all ranks.
33 +
34 + Returns (synced_batch_desc, num_tokens_across_dp).
35 + """
36 + assert dp_size > 1, "DP size must be greater than 1"
37 group = get_dp_group().cpu_group
38 tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu")
39 tensor[0][dp_rank] = num_tokens
27 - tensor[1][dp_rank] = cudagraph_size
28 - tensor[2][dp_rank] = cudagraph_runtime_mode
40 + tensor[1][dp_rank] = desired_batch_desc.cg_mode.value
41 + tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None)
42 dist.all_reduce(tensor, group=group)
30 - return tensor[0], tensor[1], tensor[2]
43
44 + num_tokens_across_dp = tensor[0]
45 + cg_mode_across_dp = tensor[1]
46 + uniform_token_counts_across_dp = tensor[2]
47
33 -def get_cudagraph_and_dp_padding(
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:
42 - return cudagraph_size, None, cudagraph_runtime_mode
43 - else:
44 - return num_tokens, None, cudagraph_runtime_mode
48 + if torch.all(num_tokens_across_dp == 0).item():
49 + synced_desc = BatchExecutionDescriptor(
50 + cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0
51 + )
52 + return synced_desc, None
53
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
54 + synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))
55
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 + # If any rank wants to run eager, all ranks run eager
57 + if synced_cg_mode == CUDAGraphMode.NONE:
58 + return BatchExecutionDescriptor(
59 + cg_mode=CUDAGraphMode.NONE,
60 + num_tokens=num_tokens,
61 + num_reqs=num_reqs,
62 + ), num_tokens_across_dp
63 +
64 + synced_num_tokens = int(num_tokens_across_dp.max().item())
65 + synced_uniform_token_count = uniform_token_counts_across_dp[0]
66 + # If ranks disagree on the uniform token count, or its 0 (means None) set to None
67 + if synced_uniform_token_count == 0 or not torch.all(
68 + uniform_token_counts_across_dp == synced_uniform_token_count
69 + ):
70 + synced_uniform_token_count = None
71 +
72 + # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
73 + # so we don't perform request padding for PIECEWISE graphs
74 + synced_desc = cudagraph_manager.dispatch(
75 + num_reqs, synced_num_tokens, synced_uniform_token_count
76 )
57 - if torch.all(num_tokens_across_dp == 0).item():
58 - # All ranks have zero tokens to run.
59 - return 0, None, 0
77
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()
78 + # Update num_tokens_across_dp to reflect padded size.
79 + num_tokens_across_dp[:] = synced_desc.num_tokens
80
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
70 - return max_cudagraph_size, num_tokens_across_dp, synced_cudagraph_mode
71 - else:
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())
77 - return num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode
81 + return synced_desc, num_tokens_across_dp
vllm/v1/worker/gpu/input_batch.py
+4 -1
@@ -37,6 +37,7 @@ class InputBatch:
37 # batch_idx -> req_id
38 req_ids: list[str]
39 num_reqs: int
40 + num_reqs_after_padding: int
41
42 # batch_idx -> req_state_idx
43 idx_mapping: torch.Tensor
@@ -123,6 +124,7 @@ class InputBatch:
124 return cls(
125 req_ids=req_ids,
126 num_reqs=num_reqs,
127 + num_reqs_after_padding=num_reqs,
128 idx_mapping=idx_mapping,
129 idx_mapping_np=idx_mapping_np,
130 expanded_idx_mapping=expanded_idx_mapping,
@@ -330,7 +332,8 @@ def combine_sampled_and_draft_tokens(
332 cu_num_logits: torch.Tensor,
333 num_logits: int,
334 ) -> torch.Tensor:
333 - num_reqs = seq_lens.shape[0]
335 + # use idx_mapping.shape[0] for actual request count
336 + num_reqs = idx_mapping.shape[0]
337 num_speculative_steps = draft_tokens.shape[-1]
338
339 logits_indices = torch.empty(
vllm/v1/worker/gpu/model_runner.py
+76 -57
@@ -40,7 +40,6 @@ from vllm.model_executor.model_loader import get_model_loader
40 from vllm.multimodal import MULTIMODAL_REGISTRY
41 from vllm.sequence import IntermediateTensors
42 from vllm.tasks import SupportedTask
43 -from vllm.utils.math_utils import cdiv
43 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
44 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
45 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -57,8 +56,12 @@ from vllm.v1.worker.gpu.attn_utils import (
56 from vllm.v1.worker.gpu.block_table import BlockTables
57 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
58 from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
60 -from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
61 -from vllm.v1.worker.gpu.dp_utils import get_cudagraph_and_dp_padding
59 +from vllm.v1.worker.gpu.cudagraph_utils import (
60 + BatchExecutionDescriptor,
61 + ModelCudaGraphManager,
62 + get_uniform_token_count,
63 +)
64 +from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
65 from vllm.v1.worker.gpu.input_batch import (
66 InputBatch,
67 InputBuffers,
@@ -137,6 +140,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
140 self.is_first_pp_rank = True
141 self.is_last_pp_rank = True
142
143 + # Data parallelism.
144 + self.dp_size = self.parallel_config.data_parallel_size
145 + self.dp_rank = self.parallel_config.data_parallel_rank
146 +
147 # Decode context parallelism.
148 self.dcp_size = self.parallel_config.decode_context_parallel_size
149 self.use_dcp = self.dcp_size > 1
@@ -193,10 +200,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
200 self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
201
202 # CUDA graphs.
196 - self.cudagraph_manager = CudaGraphManager(
203 + self.decode_query_len = self.num_speculative_steps + 1
204 + self.cudagraph_manager = ModelCudaGraphManager(
205 self.vllm_config,
198 - self.use_aux_hidden_state_outputs,
206 self.device,
207 + self.compilation_config.cudagraph_mode,
208 + decode_query_len=self.decode_query_len,
209 )
210 # Structured outputs worker.
211 self.structured_outputs_worker = StructuredOutputsWorker(
@@ -331,17 +340,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
340 **kwargs,
341 ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
342 # Create a dummy scheduler output.
343 + num_reqs = min(num_tokens, self.max_num_reqs)
344 if uniform_decode:
335 - # Align tokens to uniform_decode_query_len for cudagraph
336 - # compatibility across DP ranks.
337 - query_len = self.cudagraph_manager.uniform_decode_query_len
338 - num_reqs = min(cdiv(num_tokens, query_len), self.max_num_reqs)
339 - num_tokens = num_reqs * query_len
340 - num_tokens_per_request = [query_len] * num_reqs
341 - else:
342 - num_reqs = min(num_tokens, self.max_num_reqs)
343 - num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
344 - num_tokens_per_request[-1] += num_tokens % num_reqs
345 + # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
346 + # and for spec-decode with MTP we want to make sure the dummy runs use
347 + # 1+num_speculative_tokens we use max here, this will likely be eventually
348 + # changed in the worker: https://github.com/vllm-project/vllm/pull/35243
349 + num_tokens = max(num_tokens, self.decode_query_len)
350 + num_reqs = num_tokens // self.decode_query_len
351 + assert num_tokens % self.decode_query_len == 0
352 + num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
353 + num_tokens_per_request[-1] += num_tokens % num_reqs
354 +
355 assert sum(num_tokens_per_request) == num_tokens
356 num_scheduled_tokens = {
357 f"_dummy_req_{i}": n for i, n in enumerate(num_tokens_per_request)
@@ -498,13 +508,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
508
509 with self.maybe_setup_dummy_loras(self.lora_config):
510 self.cudagraph_manager.capture(
501 - model=self.model,
502 - model_state=self.model_state,
503 - input_buffers=self.input_buffers,
504 - block_tables=self.block_tables,
505 - attn_groups=self.attn_groups,
506 - kv_cache_config=self.kv_cache_config,
511 + self.model,
512 + self.model_state,
513 + self.input_buffers,
514 + self.block_tables,
515 + self.attn_groups,
516 + self.kv_cache_config,
517 has_lora=self.lora_config is not None,
518 + use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
519 )
520 if self.speculator is not None:
521 self.speculator.capture_model()
@@ -592,9 +603,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
603 )
604
605 def prepare_inputs(
595 - self, scheduler_output: SchedulerOutput, num_tokens_after_padding: int
606 + self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor
607 ) -> InputBatch:
608 num_tokens = scheduler_output.total_num_scheduled_tokens
609 + num_tokens_after_padding = batch_desc.num_tokens
610 assert num_tokens > 0
611 num_tokens_per_req = scheduler_output.num_scheduled_tokens
612 num_reqs = len(num_tokens_per_req)
@@ -644,6 +656,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
656 )
657
658 # Get query_start_loc.
659 + # num_reqs_padded is None for PIECEWISE graphs (no request padding needed)
660 + num_reqs_padded = batch_desc.num_reqs or num_reqs
661 query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
662 query_start_loc_np[0] = 0
663 np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1 : num_reqs + 1])
@@ -651,8 +665,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
665 # Some attention backends like FA3 require query_start_loc to be non-decreasing.
666 query_start_loc_np[num_reqs + 1 :] = num_tokens
667 async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
654 - query_start_loc_np = query_start_loc_np[: num_reqs + 1]
655 - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
668 + query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
669 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
670
671 # Get prefill tokens if any.
672 if self.req_states.any_prefills(idx_mapping_np):
@@ -674,7 +688,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
688 self.input_buffers.positions,
689 self.input_buffers.seq_lens,
690 )
677 - seq_lens = self.input_buffers.seq_lens[:num_reqs]
691 + seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
692
693 dcp_local_seq_lens = None
694 if self.use_dcp:
@@ -687,7 +701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
701 self.dcp_rank,
702 self.cp_interleave,
703 )
690 - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
704 + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded]
705
706 # Some input token ids are directly read from the last sampled tokens
707 # and draft tokens. Also, get the logits indices to sample tokens from.
@@ -706,6 +720,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
720 return InputBatch(
721 req_ids=req_ids,
722 num_reqs=num_reqs,
723 + num_reqs_after_padding=num_reqs_padded,
724 idx_mapping=idx_mapping,
725 idx_mapping_np=idx_mapping_np,
726 expanded_idx_mapping=expanded_idx_mapping,
@@ -729,13 +744,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
744 def prepare_attn(
745 self, input_batch: InputBatch
746 ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
732 - # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
733 - block_tables = self.block_tables.gather_block_tables(input_batch.idx_mapping)
734 - # Compute slot mappings: [num_kv_cache_groups, num_tokens]
747 + # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
748 + block_tables = self.block_tables.gather_block_tables(
749 + input_batch.idx_mapping,
750 + num_reqs_padded=input_batch.num_reqs_after_padding,
751 + )
752 + # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
753 + # Kernel pads beyond num_tokens with PAD_SLOT_ID.
754 slot_mappings = self.block_tables.compute_slot_mappings(
755 input_batch.idx_mapping,
756 input_batch.query_start_loc,
757 input_batch.positions,
758 + num_tokens_padded=input_batch.num_tokens_after_padding,
759 )
760 return block_tables, slot_mappings
761
@@ -851,27 +871,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
871 empty_output = self.kv_connector.no_forward(scheduler_output)
872 return empty_output
873
854 - # Get local cudagraph mode and size.
855 - local_cudagraph_mode, local_cudagraph_size = (
856 - self.cudagraph_manager.get_cudagraph_runtime_mode(
857 - num_reqs=len(scheduler_output.num_scheduled_tokens),
858 - num_tokens=scheduler_output.total_num_scheduled_tokens,
859 - max_query_len=max(scheduler_output.num_scheduled_tokens.values()),
860 - )
874 + # Get batch descriptor and sync across DP ranks.
875 + num_reqs = len(scheduler_output.num_scheduled_tokens)
876 + num_toks = scheduler_output.total_num_scheduled_tokens
877 + max_query_len = max(scheduler_output.num_scheduled_tokens.values())
878 + uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)
879 +
880 + batch_desc = self.cudagraph_manager.dispatch(
881 + num_reqs, num_toks, uniform_tok_count
882 )
883 + num_tokens_across_dp = None
884
863 - # DP sync: num_tokens + cudagraph_size + cudagraph_mode
864 - num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode = (
865 - get_cudagraph_and_dp_padding(
866 - scheduler_output.total_num_scheduled_tokens,
867 - local_cudagraph_size,
868 - local_cudagraph_mode.value,
869 - self.parallel_config.data_parallel_size,
870 - self.parallel_config.data_parallel_rank,
885 + if self.dp_size > 1:
886 + batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
887 + self.cudagraph_manager,
888 + batch_desc,
889 + num_toks,
890 + num_reqs,
891 + uniform_tok_count,
892 + self.dp_size,
893 + self.dp_rank,
894 )
872 - )
873 - cudagraph_runtime_mode = CUDAGraphMode(synced_cudagraph_mode)
874 - if num_tokens_after_padding == 0:
895 +
896 + if batch_desc.num_tokens == 0:
897 # All DP ranks have zero tokens to run.
898 empty_output = self.kv_connector.no_forward(scheduler_output)
899 return empty_output
@@ -879,9 +901,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
901 if not dummy_run:
902 # Common case.
903 # Prepare all the inputs and copy to the input buffers.
882 - input_batch = self.prepare_inputs(
883 - scheduler_output, num_tokens_after_padding
884 - )
904 + input_batch = self.prepare_inputs(scheduler_output, batch_desc)
905 block_tables, slot_mappings = self.prepare_attn(input_batch)
906
907 if self.lora_config:
@@ -894,9 +914,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
914 self._set_active_loras(*lora_inputs)
915 else:
916 # No actual tokens to run. A dummy run for DP or memory profiling.
897 - num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
917 input_batch = InputBatch.make_dummy(
899 - num_reqs, num_tokens_after_padding, self.input_buffers
918 + batch_desc.num_reqs or num_reqs,
919 + batch_desc.num_tokens,
920 + self.input_buffers,
921 )
922 if not skip_attn_for_dummy_run:
923 block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
@@ -948,14 +969,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
969 model_inputs["intermediate_tensors"] = intermediate_tensors
970
971 # Run model.
951 - if cudagraph_runtime_mode == CUDAGraphMode.FULL:
972 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
973 # Use explicit cudagraph replay for FULL mode.
974 # NOTE(woosuk): Here, we don't need to pass the input tensors,
975 # because they are already copied to the CUDA graph input buffers.
976 self.kv_connector.pre_forward(scheduler_output)
956 - model_output = self.cudagraph_manager.run_fullgraph(
957 - input_batch.num_tokens_after_padding
958 - )
977 + model_output = self.cudagraph_manager.run_fullgraph(batch_desc)
978 if self.use_aux_hidden_state_outputs:
979 hidden_states, aux_hidden_states = model_output
980 else:
@@ -972,7 +991,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
991 attn_metadata,
992 self.vllm_config,
993 num_tokens=input_batch.num_tokens_after_padding,
975 - cudagraph_runtime_mode=cudagraph_runtime_mode,
994 + cudagraph_runtime_mode=batch_desc.cg_mode,
995 num_tokens_across_dp=num_tokens_across_dp,
996 batch_descriptor=batch_descriptor,
997 slot_mapping=slot_mappings_by_layer,
vllm/v1/worker/gpu/model_states/default.py
+5 -2
@@ -142,12 +142,15 @@ class DefaultModelState(ModelState):
142 attn_groups: list[list[AttentionGroup]],
143 kv_cache_config: KVCacheConfig,
144 ) -> dict[str, Any]:
145 + # Use padded sizes - padding is handled by model_runner.prepare_attn.
146 + num_reqs = input_batch.num_reqs_after_padding
147 + num_tokens = input_batch.num_tokens_after_padding
148 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
149 max_query_len = input_batch.num_scheduled_tokens.max().item()
150 attn_metadata = build_attn_metadata(
151 attn_groups=attn_groups,
149 - num_reqs=input_batch.num_reqs,
150 - num_tokens=input_batch.num_tokens,
152 + num_reqs=num_reqs,
153 + num_tokens=num_tokens,
154 query_start_loc_gpu=input_batch.query_start_loc,
155 query_start_loc_cpu=query_start_loc_cpu,
156 max_query_len=max_query_len,
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+52 -175
@@ -1,214 +1,91 @@
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
4
5 import torch
6
7 from vllm.config import VllmConfig
8 from vllm.config.compilation import CUDAGraphMode
10 -from vllm.model_executor.offloader.base import get_offloader
9 from vllm.v1.kv_cache_interface import KVCacheConfig
10 from vllm.v1.worker.gpu.block_table import BlockTables
11 from vllm.v1.worker.gpu.cudagraph_utils import (
14 - capture_graphs,
15 - get_cudagraph_sizes,
12 + BatchExecutionDescriptor,
13 + CudaGraphManager,
14 prepare_inputs_to_capture,
15 )
18 -from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
16 from vllm.v1.worker.gpu.input_batch import InputBuffers
17 from vllm.v1.worker.gpu.model_states.interface import ModelState
18 from vllm.v1.worker.utils import AttentionGroup
19
20
24 -class EagleCudaGraphManager:
25 - def __init__(self, vllm_config: VllmConfig, device: torch.device):
26 - self.vllm_config = vllm_config
27 - self.scheduler_config = vllm_config.scheduler_config
28 - self.device = device
21 +class EagleCudaGraphManager(CudaGraphManager):
22 + """CudaGraphManager for Eagle speculative decoding (FULL mode only)."""
23
30 - self.max_model_len = vllm_config.model_config.max_model_len
31 - self.max_num_reqs = self.scheduler_config.max_num_seqs
32 - self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
33 - self.dp_size = vllm_config.parallel_config.data_parallel_size
34 - self.compilation_config = vllm_config.compilation_config
35 - assert self.compilation_config is not None
36 -
37 - # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.
38 - self.cudagraph_mode = self.compilation_config.cudagraph_mode.decode_mode()
39 -
40 - # only need to capture uniform decode cudagraph sizes (the 2nd return value)
41 - _, self.cudagraph_sizes = get_cudagraph_sizes(
42 - self.compilation_config.cudagraph_capture_sizes,
43 - self.max_num_reqs,
44 - self.max_num_tokens,
45 - self.cudagraph_mode,
46 - uniform_decode_query_len=1,
47 - uniform_decode_cudagraph=True,
24 + def __init__(
25 + self,
26 + vllm_config: VllmConfig,
27 + device: torch.device,
28 + cudagraph_mode: CUDAGraphMode,
29 + draft_tokens: torch.Tensor,
30 + ):
31 + assert not cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE), (
32 + "EagleCudaGraphManager does not support PIECEWISE mode yet"
33 )
49 -
50 - self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
51 - self.pool = None
52 - if self.cudagraph_mode != CUDAGraphMode.NONE:
34 + # Eagle always uses uniform decode with query_len=1
35 + super().__init__(vllm_config, device, cudagraph_mode, decode_query_len=1)
36 + self.draft_tokens = draft_tokens
37 +
38 + # Use a dedicated pool for Eagle to avoid memory overlap with the main
39 + # model's cudagraph. The base class uses a shared global pool, but Eagle's
40 + # internal allocations (e.g., gumbel_sample temporaries) can conflict with
41 + # the main model's allocations when sharing the same pool.
42 + if cudagraph_mode:
43 self.pool = torch.cuda.graph_pool_handle()
44
55 - def get_cudagraph_size(self, num_tokens: int) -> int | None:
56 - return self.cudagraph_sizes.get(num_tokens)
57 -
58 - def get_cudagraph_runtime_mode(
59 - self, num_tokens: int
60 - ) -> tuple[CUDAGraphMode, int | None]:
61 - cudagraph_size = self.get_cudagraph_size(num_tokens)
62 - if cudagraph_size is None:
63 - cudagraph_mode = CUDAGraphMode.NONE
64 - else:
65 - cudagraph_mode = self.cudagraph_mode
66 -
67 - if (
68 - cudagraph_mode == CUDAGraphMode.FULL
69 - and cudagraph_size is not None
70 - and cudagraph_size not in self.graphs
71 - ):
72 - # If graph wasn't captured yet, fall back to eager.
73 - # This might happen when the dummy run is called before capture.
74 - cudagraph_mode = CUDAGraphMode.NONE
75 - cudagraph_size = None
76 - return cudagraph_mode, cudagraph_size
77 -
78 - def capture_graph(
45 + def capture(
46 self,
80 - num_tokens: int,
81 - capture_cg_mode: CUDAGraphMode,
47 generate_fn: Callable,
48 model_state: ModelState,
49 input_buffers: InputBuffers,
50 block_tables: BlockTables,
51 attn_groups: list[list[AttentionGroup]],
52 kv_cache_config: KVCacheConfig,
53 + progress_bar_desc: str = "Capturing CUDA graphs",
54 ) -> None:
89 - assert capture_cg_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], (
90 - f"Invalid capture_cudagraph_mode for capture: {capture_cg_mode}"
91 - )
92 - if capture_cg_mode == CUDAGraphMode.PIECEWISE:
93 - capture_fn = self._capture_piecewise_graph
94 - else:
95 - capture_fn = self._capture_full_graph
96 -
97 - num_reqs = min(num_tokens, self.max_num_reqs)
98 - attn_metadata, slot_mappings = prepare_inputs_to_capture(
99 - num_reqs,
100 - num_tokens,
101 - model_state,
102 - input_buffers,
103 - block_tables,
104 - attn_groups,
105 - kv_cache_config,
106 - )
107 - num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
108 -
109 - # Warm up.
110 - generate_fn(
111 - num_reqs,
112 - num_tokens,
113 - attn_metadata,
114 - slot_mappings,
115 - num_tokens_across_dp,
116 - CUDAGraphMode.NONE,
117 - )
118 -
119 - # Capture the graph.
120 - capture_fn(
121 - num_reqs=num_reqs,
122 - num_tokens=num_tokens,
123 - generate_fn=generate_fn,
124 - attn_metadata=attn_metadata,
125 - slot_mappings=slot_mappings,
126 - num_tokens_across_dp=num_tokens_across_dp,
127 - )
128 -
129 - def _capture_full_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 - assert num_tokens not in self.graphs
139 - graph = torch.cuda.CUDAGraph()
140 -
141 - # Sync offloader's copy stream before capture.
142 - # Ensure any pre-capture prefetches from offloader are complete.
143 - get_offloader().sync_prev_onload()
55 + """Capture CUDA graphs for Eagle speculative decoding (FULL mode only)."""
56 +
57 + def create_forward_fn(
58 + desc: BatchExecutionDescriptor,
59 + ) -> Callable[[CUDAGraphMode], None]:
60 + num_tokens = desc.num_tokens
61 + num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
62 + num_tokens_across_dp = (
63 + torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
64 + if self.dp_size > 1
65 + else None
66 + )
67 + attn_metadata, slot_mappings = prepare_inputs_to_capture(
68 + num_reqs,
69 + num_tokens,
70 + model_state,
71 + input_buffers,
72 + block_tables,
73 + attn_groups,
74 + kv_cache_config,
75 + )
76
145 - with torch.cuda.graph(graph, self.pool):
146 - generate_fn(
77 + return lambda cg_mode: generate_fn(
78 num_reqs,
79 num_tokens,
80 attn_metadata,
81 slot_mappings,
82 num_tokens_across_dp,
152 - CUDAGraphMode.NONE,
83 + cg_mode,
84 )
154 - # Join offloader's copy stream after forward to avoid unjoined
155 - # stream error. The last layer's start_prefetch forks copy_stream,
156 - # but wait_prefetch only happens in the next forward pass.
157 - get_offloader().join_after_forward()
158 - self.graphs[num_tokens] = graph
159 -
160 - def _capture_piecewise_graph(
161 - self,
162 - num_reqs: int,
163 - num_tokens: int,
164 - generate_fn: Callable,
165 - attn_metadata: dict[str, Any],
166 - slot_mappings: dict[str, torch.Tensor],
167 - num_tokens_across_dp: torch.Tensor,
168 - ) -> None:
169 - generate_fn(
170 - num_reqs,
171 - num_tokens,
172 - attn_metadata,
173 - slot_mappings,
174 - num_tokens_across_dp,
175 - CUDAGraphMode.PIECEWISE,
176 - )
177 -
178 - @torch.inference_mode()
179 - def capture(
180 - self,
181 - generate_fn: Callable,
182 - model_state: ModelState,
183 - input_buffers: InputBuffers,
184 - block_tables: BlockTables,
185 - attn_groups: list[list[AttentionGroup]],
186 - kv_cache_config: KVCacheConfig,
187 - ) -> None:
188 - if self.cudagraph_mode == CUDAGraphMode.NONE:
189 - return
85
191 - capture_graphs(
192 - self.cudagraph_sizes,
193 - self.device,
194 - self.capture_graph,
195 - capture_cudagraph_mode=self.cudagraph_mode,
196 - desc=f"Capturing eagle CUDA graphs ({self.cudagraph_mode.name})",
197 - generate_fn=generate_fn,
198 - model_state=model_state,
199 - input_buffers=input_buffers,
200 - block_tables=block_tables,
201 - attn_groups=attn_groups,
202 - kv_cache_config=kv_cache_config,
203 - )
86 + super().capture(create_forward_fn, progress_bar_desc)
87
205 - def run_fullgraph(self, num_tokens: int) -> None:
206 - assert num_tokens in self.graphs
207 - # Sync offloader before replay - needed when transitioning from
208 - # eager/piecewise to full cudagraph (e.g., prefill → decode).
209 - # The previous eager iteration's start_prefetch may have queued
210 - # H2D copies on copy_stream that the graph's captured events
211 - # cannot see. Without this, replay could overwrite static buffers
212 - # while those copies are still in flight.
213 - get_offloader().sync_prev_onload()
214 - self.graphs[num_tokens].replay()
88 + def run_fullgraph(self, desc: BatchExecutionDescriptor) -> torch.Tensor:
89 + """Replay a captured FULL cudagraph and return draft tokens."""
90 + super().run_fullgraph(desc)
91 + return self.draft_tokens
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+43 -28
@@ -16,7 +16,7 @@ from vllm.v1.worker.gpu.attn_utils import (
16 build_slot_mappings_by_layer,
17 )
18 from vllm.v1.worker.gpu.block_table import BlockTables
19 -from vllm.v1.worker.gpu.dp_utils import get_cudagraph_and_dp_padding
19 +from vllm.v1.worker.gpu.dp_utils import sync_cudagraph_and_dp_padding
20 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
21 from vllm.v1.worker.gpu.model_states.interface import ModelState
22 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
@@ -75,7 +75,16 @@ class EagleSpeculator:
75 device=device,
76 )
77
78 - self.cudagraph_manager = EagleCudaGraphManager(vllm_config, device)
78 + # currently we don't support PIECEWISE for Eagle.
79 + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
80 + if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
81 + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
82 + else:
83 + cudagraph_mode = CUDAGraphMode.NONE
84 +
85 + self.cudagraph_manager = EagleCudaGraphManager(
86 + vllm_config, device, cudagraph_mode, self.draft_tokens
87 + )
88
89 def load_model(self, target_model: nn.Module) -> None:
90 self.model = load_eagle_model(target_model, self.vllm_config)
@@ -171,7 +180,7 @@ class EagleSpeculator:
180 )
181 if attn_metadata is not None:
182 self.block_tables.compute_slot_mappings(
174 - idx_mapping, query_start_loc, pos
183 + idx_mapping, query_start_loc, pos, num_tokens_padded
184 )
185
186 def capture_model(self) -> None:
@@ -185,6 +194,7 @@ class EagleSpeculator:
194 self.block_tables,
195 self.attn_groups,
196 self.kv_cache_config,
197 + progress_bar_desc="Capturing eagle CUDA graphs",
198 )
199
200 @torch.inference_mode()
@@ -251,6 +261,7 @@ class EagleSpeculator:
261 logits = self.model.compute_logits(sample_hidden_states)
262
263 num_reqs = input_batch.num_reqs
264 + num_reqs_padded = input_batch.num_reqs_after_padding
265 # NOTE(woosuk): For draft sampling, we only consider the temperature
266 # and ignore the other sampling parameters such as top_k and top_p,
267 # for simplicity and performance.
@@ -292,48 +303,52 @@ class EagleSpeculator:
303 self.max_num_reqs,
304 )
305
295 - if not (dummy_run and skip_attn_for_dummy_run):
296 - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
297 - slot_mappings = self.block_tables.compute_slot_mappings(
298 - idx_mapping, query_start_loc, pos
299 - )
306 + # Get batch descriptor and sync across DP ranks.
307 + # Eagle uses FULL-only mode, dispatch with uniform_token_count=1 for decode
308
301 - cudagraph_mode, cudagraph_size = (
302 - self.cudagraph_manager.get_cudagraph_runtime_mode(num_reqs)
303 - )
304 - num_tokens_padded, num_tokens_across_dp, synced_cudagraph_mode = (
305 - get_cudagraph_and_dp_padding(
309 + batch_desc = self.cudagraph_manager.dispatch(num_reqs, num_reqs, 1)
310 + num_tokens_across_dp = None
311 +
312 + if self.dp_size > 1:
313 + batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
314 + self.cudagraph_manager,
315 + batch_desc,
316 num_reqs,
307 - cudagraph_size,
308 - cudagraph_mode.value,
317 + num_reqs,
318 + 1, # uniform_token_count
319 self.dp_size,
320 self.dp_rank,
321 )
312 - )
313 - cudagraph_mode = CUDAGraphMode(synced_cudagraph_mode)
314 - if cudagraph_mode == CUDAGraphMode.FULL:
315 - # Run full CUDA graph.
316 - self.cudagraph_manager.run_fullgraph(num_tokens_padded)
317 - return self.draft_tokens[:num_reqs]
322 +
323 + if not (dummy_run and skip_attn_for_dummy_run):
324 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
325 + slot_mappings = self.block_tables.compute_slot_mappings(
326 + idx_mapping, query_start_loc, pos, batch_desc.num_tokens
327 + )
328 +
329 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
330 + return self.cudagraph_manager.run_fullgraph(batch_desc)[:num_reqs]
331
332 # Run eager or piecewise CUDA graph.
333 attn_metadata_updated = None
334 slot_mappings_updated = None
335 if not (dummy_run and skip_attn_for_dummy_run):
336 query_start_loc_cpu = torch.arange(
324 - num_reqs + 1, dtype=torch.int32, device="cpu"
337 + num_reqs_padded + 1, dtype=torch.int32, device="cpu"
338 )
326 - block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
339 + block_tables = [
340 + x[:num_reqs_padded] for x in self.block_tables.input_block_tables
341 + ]
342
343 # FIXME(woosuk): This is UNSAFE!!
344 attn_metadata_updated = build_attn_metadata(
345 attn_groups=self.attn_groups,
331 - num_reqs=num_reqs,
332 - num_tokens=num_reqs,
346 + num_reqs=num_reqs_padded,
347 + num_tokens=num_reqs_padded,
348 query_start_loc_gpu=query_start_loc,
349 query_start_loc_cpu=query_start_loc_cpu,
350 max_query_len=1,
336 - seq_lens=self.input_buffers.seq_lens[:num_reqs],
351 + seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
352 max_seq_len=self.max_model_len,
353 block_tables=block_tables,
354 slot_mappings=slot_mappings,
@@ -345,11 +360,11 @@ class EagleSpeculator:
360
361 self.generate_draft(
362 num_reqs,
348 - num_tokens_padded,
363 + batch_desc.num_tokens,
364 attn_metadata_updated,
365 slot_mappings_updated,
366 num_tokens_across_dp=num_tokens_across_dp,
352 - cudagraph_runtime_mode=cudagraph_mode,
367 + cudagraph_runtime_mode=batch_desc.cg_mode,
368 )
369 return self.draft_tokens[:num_reqs]
370