@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 91945b6e4a

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 2, 01:00 UTC

[Bug Fix][Model Runner V2][Spec Decode] Warmup & capture with different attention states for speculator prefill (#44253)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

Giancarlo Delfin committed Jun 3, 2026 at 15:32 UTC 91945b6e4ade361125837228179cee01e6573023

4 files changed +59 -34

vllm/v1/worker/gpu/cudagraph_utils.py

+46 -25

@@ -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, NamedTuple

6 +from typing import Any, NamedTuple, Protocol

8 import torch

9 import torch.nn as nn

@@ -37,11 +37,16 @@ from vllm.v1.worker.utils import AttentionGroup

37 logger = init_logger(__name__)

40 -class CapturedAttentionState(NamedTuple):

40 +class AttentionState(NamedTuple):

41 attn_metadata: dict[str, Any] | None

42 slot_mappings: dict[str, torch.Tensor]

45 +class AttentionStatePair(NamedTuple):

46 + warmup: AttentionState

47 + captured: AttentionState

48 +

49 +

50 @dataclass(frozen=True)

51 class BatchExecutionDescriptor:

52 """Describes the shape of the batch and CG mode to run; this is used to make shape

@@ -53,6 +58,18 @@ class BatchExecutionDescriptor:

58 uniform_token_count: int | None = None

61 +class CreateForwardFn(Protocol):

62 + """Factory that prepares inputs (OUTSIDE the graph) and returns a tuple of

63 + (forward_fn, attn_state). Called with warmup=True for the warmup pass and

64 + warmup=False for the captured pass."""

65 +

66 + def __call__(

67 + self,

68 + desc: BatchExecutionDescriptor,

69 + warmup: bool,

70 + ) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]: ...

71 +

72 +

73 def _is_compatible(

74 desc: BatchExecutionDescriptor,

75 num_reqs: int,

@@ -198,21 +215,21 @@ class CudaGraphManager:

215 @torch.inference_mode()

216 def capture(

217 self,

201 - create_forward_fn: Callable[

202 - [BatchExecutionDescriptor],

203 - tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState],

204 - ],

218 + create_forward_fn: CreateForwardFn,

219 progress_bar_desc: str = "Capturing CUDA graphs",

206 - ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:

220 + ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:

221 """Capture CUDA graphs.

222

223 Args:

224 create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and

211 - returns a tuple of (forward_fn, captured_attn_state).

225 + returns a tuple of (forward_fn, attn_state). For FULL cudagraph

226 + mode, it is invoked once with warmup=True for the warmup pass,

227 + and again with warmup=False for the captured pass. For attention

228 + backends that perform lazy metadata initialization (e.g. FlashMLA),

229 + FULL cudagraph capture requires distinct metadatas for warmup and

230 + capture.

231 """

213 - captured_attn_states: dict[

214 - BatchExecutionDescriptor, CapturedAttentionState

215 - ] = {}

232 + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {}

233 with graph_capture(device=self.device):

234 # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger

235 # activations so FULL activations should fit in already allocated

@@ -226,7 +243,7 @@ class CudaGraphManager:

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

244 for desc in descs:

245 # Prepare inputs and get forward function

229 - forward_fn, attn_state = create_forward_fn(desc)

246 + forward_fn, warmup_attn_state = create_forward_fn(desc, warmup=True)

247

248 # Warmup

249 forward_fn(CUDAGraphMode.NONE)

@@ -236,15 +253,18 @@ class CudaGraphManager:

253 "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc

254 )

255 if desc.cg_mode == CUDAGraphMode.PIECEWISE:

239 - captured_attn_states[desc] = attn_state

256 + attn_states[desc] = AttentionStatePair(

257 + warmup_attn_state, warmup_attn_state

258 + )

259 forward_fn(CUDAGraphMode.PIECEWISE)

260 else:

242 - # Capture with fresh attention state. The warmup

243 - # attention state is discarded because some backends

244 - # (e.g. FlashMLA) perform lazy initializations that

245 - # must be captured in the graph.

246 - forward_fn, attn_state = create_forward_fn(desc)

247 - captured_attn_states[desc] = attn_state

261 + # Capture with fresh attention state.

262 + forward_fn, capture_attn_state = create_forward_fn(

263 + desc, warmup=False

264 + )

265 + attn_states[desc] = AttentionStatePair(

266 + warmup_attn_state, capture_attn_state

267 + )

268 assert desc not in self.graphs, (

269 f"Graph already captured for {desc}"

270 )

@@ -262,7 +282,7 @@ class CudaGraphManager:

282 self.graphs[desc] = graph

283 compilation_counter.num_cudagraph_captured += 1

284 self._graphs_captured = True

265 - return captured_attn_states

285 + return attn_states

286

287 def dispatch(

288 self,

@@ -337,7 +357,7 @@ class ModelCudaGraphManager(CudaGraphManager):

357 has_lora: bool = False,

358 use_aux_hidden_state_outputs: bool = False,

359 progress_bar_desc: str = "Capturing CUDA graphs",

340 - ) -> dict[BatchExecutionDescriptor, CapturedAttentionState]:

360 + ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:

361 """Capture CUDA graphs for model forward pass."""

362 self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs

363 if self.use_breakable_cg:

@@ -345,9 +365,10 @@ class ModelCudaGraphManager(CudaGraphManager):

365

366 def create_forward_fn(

367 desc: BatchExecutionDescriptor,

368 + warmup: bool,

369 ) -> tuple[

370 Callable[[CUDAGraphMode], None],

350 - CapturedAttentionState,

371 + AttentionState,

372 ]:

373 num_tokens = desc.num_tokens

374 num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)

@@ -435,7 +456,7 @@ class ModelCudaGraphManager(CudaGraphManager):

456 for k, v in intermediate_tensors.tensors.items():

457 self.intermediate_tensors[k][:num_tokens] = v

458

438 - return forward_fn, CapturedAttentionState(attn_metadata, slot_mappings)

459 + return forward_fn, AttentionState(attn_metadata, slot_mappings)

460

461 return super().capture(create_forward_fn, progress_bar_desc)

462

@@ -464,7 +485,7 @@ def prepare_inputs_to_capture(

485 attn_groups: list[list[AttentionGroup]],

486 kv_cache_config: KVCacheConfig,

487 skip_attn: bool = False,

467 -) -> CapturedAttentionState:

488 +) -> AttentionState:

489 input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)

490 input_block_tables = block_tables.get_dummy_block_tables(num_reqs)

491 slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)

@@ -495,4 +516,4 @@ def prepare_inputs_to_capture(

516 kv_cache_config,

517 for_capture=True,

518 )

498 - return CapturedAttentionState(attn_metadata, slot_mappings_by_layer)

519 + return AttentionState(attn_metadata, slot_mappings_by_layer)

vllm/v1/worker/gpu/model_runner.py

+2 -2

@@ -678,7 +678,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

678 start_free_gpu_memory = torch.cuda.mem_get_info()[0]

679

680 with self.maybe_setup_dummy_loras(self.lora_config):

681 - captured_attn_states = self.cudagraph_manager.capture(

681 + attn_states = self.cudagraph_manager.capture(

682 self.model,

683 self.model_state,

684 self.input_buffers,

@@ -690,7 +690,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

690 use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,

691 )

692 if self.speculator is not None:

693 - self.speculator.capture(captured_attn_states)

693 + self.speculator.capture(attn_states)

694

695 end_time = time.perf_counter()

696 end_free_gpu_memory = torch.cuda.mem_get_info()[0]

vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py

+9 -5

@@ -8,8 +8,9 @@ from vllm.config.compilation import CUDAGraphMode

8 from vllm.v1.kv_cache_interface import KVCacheConfig

9 from vllm.v1.worker.gpu.block_table import BlockTables

10 from vllm.v1.worker.gpu.cudagraph_utils import (

11 + AttentionState,

12 + AttentionStatePair,

13 BatchExecutionDescriptor,

12 - CapturedAttentionState,

14 CudaGraphManager,

15 prepare_inputs_to_capture,

16 )

@@ -25,12 +26,13 @@ class PrefillEagleCudaGraphManager(CudaGraphManager):

26 def capture(

27 self,

28 forward_fn: Callable,

28 - full_cg_attn_states: dict[BatchExecutionDescriptor, CapturedAttentionState],

29 + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],

30 progress_bar_desc: str = "Capturing CUDA graphs",

31 ) -> None:

32 def create_forward_fn(

33 desc: BatchExecutionDescriptor,

33 - ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]:

34 + warmup: bool,

35 + ) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]:

36 num_tokens = desc.num_tokens

37 num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)

38 num_tokens_across_dp = (

@@ -38,7 +40,8 @@ class PrefillEagleCudaGraphManager(CudaGraphManager):

40 if self.dp_size > 1

41 else None

42 )

41 - attn_state = full_cg_attn_states[desc]

43 + attn_state_pair = attn_states[desc]

44 + attn_state = attn_state_pair.warmup if warmup else attn_state_pair.captured

45 attn_metadata, slot_mappings = attn_state

46 fwd = lambda cg_mode: forward_fn(

47 num_reqs,

@@ -69,7 +72,8 @@ class DecodeEagleCudaGraphManager(CudaGraphManager):

72 ) -> None:

73 def create_forward_fn(

74 desc: BatchExecutionDescriptor,

72 - ) -> tuple[Callable[[CUDAGraphMode], None], CapturedAttentionState]:

75 + warmup: bool,

76 + ) -> tuple[Callable[[CUDAGraphMode], None], AttentionState]:

77 num_tokens = desc.num_tokens

78 num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)

79 num_tokens_across_dp = (

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

+2 -2

@@ -20,8 +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 + AttentionStatePair,

24 BatchExecutionDescriptor,

24 - CapturedAttentionState,

25 get_uniform_token_count,

26 )

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

@@ -426,7 +426,7 @@ class EagleSpeculator:

426

427 def capture(

428 self,

429 - attn_states: dict[BatchExecutionDescriptor, CapturedAttentionState],

429 + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],

430 ) -> None:

431 logger.info("Capturing model for Eagle speculator...")

432 # Reset indices to zeros to prevent stale values from prior