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 4, 15:00 UTC
[Model Runner V2] Skip attention metadata rebuild before draft prefill (#40410)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 27, 2026 at 15:38 UTC
4c7c69b4e0aa7062b8a48268abb06c041bcec53d
4 files changed
+88
-40
vllm/v1/worker/gpu/cudagraph_utils.py
+25
-11
@@ -3,7 +3,7 @@
3
from collections import defaultdict
4
from collections.abc import Callable
5
from dataclasses import dataclass
6
-from typing import Any
6
+from typing import Any, NamedTuple
7
8
import torch
9
import torch.nn as nn
@@ -32,6 +32,11 @@ from vllm.v1.worker.utils import AttentionGroup
32
logger = init_logger(__name__)
33
34
35
+class CapturedAttentionState(NamedTuple):
36
+ attn_metadata: dict[str, Any]
37
+ slot_mappings: dict[str, torch.Tensor]
38
+
39
+
40
@dataclass(frozen=True)
41
class BatchExecutionDescriptor:
42
"""Describes the shape of the batch and CG mode to run; this is used to make shape
@@ -177,16 +182,20 @@ class CudaGraphManager:
182
def capture(
183
self,
184
create_forward_fn: Callable[
180
- [BatchExecutionDescriptor], Callable[[CUDAGraphMode], None]
185
+ [BatchExecutionDescriptor],
186
+ tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState],
187
],
188
progress_bar_desc: str = "Capturing CUDA graphs",
183
- ) -> None:
189
+ ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:
190
"""Capture CUDA graphs.
191
192
Args:
193
create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
188
- returns a function that runs forward with a given CUDAGraphMode.
194
+ returns a tuple of (forward_fn, captured_attn_state).
195
"""
196
+ captured_attn_states: dict[
197
+ BatchExecutionDescriptor, CapturedAttentionState
198
+ ] = {}
199
with graph_capture(device=self.device):
200
# Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
201
# activations so FULL activations should fit in already allocated
@@ -200,7 +209,8 @@ class CudaGraphManager:
209
descs = tqdm(descs, desc=f"{progress_bar_desc} ({mode.name})")
210
for desc in descs:
211
# Prepare inputs and get forward function
203
- forward_fn = create_forward_fn(desc)
212
+ forward_fn, attn_state = create_forward_fn(desc)
213
+ captured_attn_states[desc] = attn_state
214
215
# Warmup
216
forward_fn(CUDAGraphMode.NONE)
@@ -228,6 +238,7 @@ class CudaGraphManager:
238
get_offloader().join_after_forward()
239
self.graphs[desc] = graph
240
self._graphs_captured = True
241
+ return captured_attn_states
242
243
def dispatch(
244
self,
@@ -289,13 +300,16 @@ class ModelCudaGraphManager(CudaGraphManager):
300
has_lora: bool = False,
301
use_aux_hidden_state_outputs: bool = False,
302
progress_bar_desc: str = "Capturing CUDA graphs",
292
- ) -> None:
303
+ ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:
304
"""Capture CUDA graphs for model forward pass."""
305
self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs
306
307
def create_forward_fn(
308
desc: BatchExecutionDescriptor,
298
- ) -> Callable[[CUDAGraphMode], None]:
309
+ ) -> tuple[
310
+ Callable[[CUDAGraphMode], None],
311
+ CapturedAttentionState,
312
+ ]:
313
num_tokens = desc.num_tokens
314
num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
315
num_tokens_across_dp = (
@@ -375,9 +389,9 @@ class ModelCudaGraphManager(CudaGraphManager):
389
for k, v in intermediate_tensors.tensors.items():
390
self.intermediate_tensors[k][:num_tokens] = v
391
378
- return forward_fn
392
+ return forward_fn, CapturedAttentionState(attn_metadata, slot_mappings)
393
380
- super().capture(create_forward_fn, progress_bar_desc)
394
+ return super().capture(create_forward_fn, progress_bar_desc)
395
396
def run_fullgraph(
397
self, desc: BatchExecutionDescriptor
@@ -403,7 +417,7 @@ def prepare_inputs_to_capture(
417
block_tables: BlockTables,
418
attn_groups: list[list[AttentionGroup]],
419
kv_cache_config: KVCacheConfig,
406
-) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
420
+) -> CapturedAttentionState:
421
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
422
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
423
slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
@@ -432,4 +446,4 @@ def prepare_inputs_to_capture(
446
kv_cache_config,
447
for_capture=True,
448
)
435
- return attn_metadata, slot_mappings_by_layer
449
+ return CapturedAttentionState(attn_metadata, slot_mappings_by_layer)
vllm/v1/worker/gpu/model_runner.py
+2
-2
@@ -590,7 +590,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
590
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
591
592
with self.maybe_setup_dummy_loras(self.lora_config):
593
- self.cudagraph_manager.capture(
593
+ captured_attn_states = self.cudagraph_manager.capture(
594
self.model,
595
self.model_state,
596
self.input_buffers,
@@ -602,7 +602,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
602
use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
603
)
604
if self.speculator is not None:
605
- self.speculator.capture_model()
605
+ self.speculator.capture(captured_attn_states)
606
607
end_time = time.perf_counter()
608
end_free_gpu_memory = torch.cuda.mem_get_info()[0]
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+48
-7
@@ -10,6 +10,7 @@ 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 (
12
BatchExecutionDescriptor,
13
+ CapturedAttentionState,
14
CudaGraphManager,
15
prepare_inputs_to_capture,
16
)
@@ -18,8 +19,8 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
19
from vllm.v1.worker.utils import AttentionGroup
20
21
21
-class EagleCudaGraphManager(CudaGraphManager):
22
- """CudaGraphManager for Eagle speculative decoding."""
22
+class EagleCudaGraphManagerBase(CudaGraphManager):
23
+ """Base CudaGraphManager for Eagle with a dedicated graph pool."""
24
25
def __init__(
26
self,
@@ -37,6 +38,46 @@ class EagleCudaGraphManager(CudaGraphManager):
38
if cudagraph_mode:
39
self.pool = torch.cuda.graph_pool_handle()
40
41
+
42
+class PrefillEagleCudaGraphManager(EagleCudaGraphManagerBase):
43
+ """Eagle CudaGraphManager for prefill, using pre-built attention states
44
+ from the target model's capture."""
45
+
46
+ def capture(
47
+ self,
48
+ forward_fn: Callable,
49
+ full_cg_attn_states: dict[BatchExecutionDescriptor, CapturedAttentionState],
50
+ progress_bar_desc: str = "Capturing CUDA graphs",
51
+ ) -> None:
52
+ def create_forward_fn(
53
+ desc: BatchExecutionDescriptor,
54
+ ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]:
55
+ num_tokens = desc.num_tokens
56
+ num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
57
+ num_tokens_across_dp = (
58
+ torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
59
+ if self.dp_size > 1
60
+ else None
61
+ )
62
+ attn_state = full_cg_attn_states[desc]
63
+ attn_metadata, slot_mappings = attn_state
64
+ fwd = lambda cg_mode: forward_fn(
65
+ num_reqs,
66
+ num_tokens,
67
+ attn_metadata,
68
+ slot_mappings,
69
+ num_tokens_across_dp,
70
+ cg_mode,
71
+ )
72
+ return fwd, attn_state
73
+
74
+ super().capture(create_forward_fn, progress_bar_desc)
75
+
76
+
77
+class DecodeEagleCudaGraphManager(EagleCudaGraphManagerBase):
78
+ """Eagle CudaGraphManager for decode draft generation, building its own
79
+ attention metadata from scratch."""
80
+
81
def capture(
82
self,
83
forward_fn: Callable,
@@ -47,11 +88,9 @@ class EagleCudaGraphManager(CudaGraphManager):
88
kv_cache_config: KVCacheConfig,
89
progress_bar_desc: str = "Capturing CUDA graphs",
90
) -> None:
50
- """Capture CUDA graphs for Eagle."""
51
-
91
def create_forward_fn(
92
desc: BatchExecutionDescriptor,
54
- ) -> Callable[[CUDAGraphMode], None]:
93
+ ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]:
94
num_tokens = desc.num_tokens
95
num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
96
num_tokens_across_dp = (
@@ -59,7 +98,7 @@ class EagleCudaGraphManager(CudaGraphManager):
98
if self.dp_size > 1
99
else None
100
)
62
- attn_metadata, slot_mappings = prepare_inputs_to_capture(
101
+ attn_state = prepare_inputs_to_capture(
102
num_reqs,
103
num_tokens,
104
model_state,
@@ -68,8 +107,9 @@ class EagleCudaGraphManager(CudaGraphManager):
107
attn_groups,
108
kv_cache_config,
109
)
110
+ attn_metadata, slot_mappings = attn_state
111
72
- return lambda cg_mode: forward_fn(
112
+ fwd = lambda cg_mode: forward_fn(
113
num_reqs,
114
num_tokens,
115
attn_metadata,
@@ -77,5 +117,6 @@ class EagleCudaGraphManager(CudaGraphManager):
117
num_tokens_across_dp,
118
cg_mode,
119
)
120
+ return fwd, attn_state
121
122
super().capture(create_forward_fn, progress_bar_desc)
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+13
-20
@@ -20,6 +20,8 @@ from vllm.v1.worker.gpu.attn_utils import (
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
22
from vllm.v1.worker.gpu.cudagraph_utils import (
23
+ BatchExecutionDescriptor,
24
+ CapturedAttentionState,
25
get_uniform_token_count,
26
)
27
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
@@ -27,7 +29,8 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
29
from vllm.v1.worker.gpu.model_states.interface import ModelState
30
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
31
from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import (
30
- EagleCudaGraphManager,
32
+ DecodeEagleCudaGraphManager,
33
+ PrefillEagleCudaGraphManager,
34
)
35
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
36
@@ -108,13 +111,13 @@ class EagleSpeculator:
111
device=device,
112
)
113
111
- self.prefill_cudagraph_manager: EagleCudaGraphManager | None = None
112
- self.decode_cudagraph_manager: EagleCudaGraphManager | None = None
114
+ self.prefill_cudagraph_manager: PrefillEagleCudaGraphManager | None = None
115
+ self.decode_cudagraph_manager: DecodeEagleCudaGraphManager | None = None
116
117
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
118
cudagraph_mode = self.vllm_config.compilation_config.cudagraph_mode
119
# Initialize cudagraph manager for draft prefill (draft position 0).
117
- self.prefill_cudagraph_manager = EagleCudaGraphManager(
120
+ self.prefill_cudagraph_manager = PrefillEagleCudaGraphManager(
121
self.vllm_config,
122
self.device,
123
cudagraph_mode,
@@ -131,7 +134,7 @@ class EagleSpeculator:
134
cudagraph_mode = CUDAGraphMode.NONE
135
136
# Initialize cudagraph manager for draft decodes (draft positions > 0).
134
- self.decode_cudagraph_manager = EagleCudaGraphManager(
137
+ self.decode_cudagraph_manager = DecodeEagleCudaGraphManager(
138
self.vllm_config,
139
self.device,
140
cudagraph_mode,
@@ -360,7 +363,10 @@ class EagleSpeculator:
363
)
364
return attn_metadata
365
363
- def capture_model(self) -> None:
366
+ def capture(
367
+ self,
368
+ attn_states: dict[BatchExecutionDescriptor, CapturedAttentionState],
369
+ ) -> None:
370
logger.info("Capturing model for Eagle speculator...")
371
# Reset indices to zeros to prevent stale values from prior
372
# dummy runs to cause out-of-bounds indexing during capture.
@@ -374,11 +380,7 @@ class EagleSpeculator:
380
assert self.prefill_cudagraph_manager is not None
381
self.prefill_cudagraph_manager.capture(
382
self.prefill,
377
- self.model_state,
378
- self.input_buffers,
379
- self.block_tables,
380
- self.attn_groups,
381
- self.kv_cache_config,
383
+ attn_states,
384
progress_bar_desc="Capturing eagle prefill CUDA graphs",
385
)
386
@@ -483,15 +485,6 @@ class EagleSpeculator:
485
)
486
487
if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:
486
- # It is necessary to rebuild the attention metadata when
487
- # replaying the FULL graph so that any attention metadata
488
- # builder state is updated.
489
- self._build_draft_attn_metadata(
490
- num_reqs=num_reqs,
491
- num_reqs_padded=prefill_batch_desc.num_reqs or num_reqs,
492
- num_tokens_padded=prefill_batch_desc.num_tokens,
493
- max_query_len=self.num_speculative_steps + 1,
494
- )
488
# Replay the full graph for draft prefill.
489
assert self.prefill_cudagraph_manager is not None
490
self.prefill_cudagraph_manager.run_fullgraph(prefill_batch_desc)