@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 26587f9519

[BugFix][ModelRunner V2] Fix stale attn metadata in speculator prefill cudagraph capture (#48261)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Jul 13, 2026 at 17:39 UTC 26587f9519e22a5c4549ead7595ad9ca3229c4fd

7 files changed +72 -108

vllm/v1/worker/gpu/cudagraph_utils.py

+17 -35

@@ -44,11 +44,6 @@ class AttentionState(NamedTuple):

44 slot_mappings: dict[str, torch.Tensor]

47 -class AttentionStatePair(NamedTuple):

48 - warmup: AttentionState

49 - captured: AttentionState

50 -

51 -

47 @dataclass(frozen=True)

48 class BatchExecutionDescriptor:

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

@@ -62,15 +57,15 @@ class BatchExecutionDescriptor:

59 class CreateForwardFn(Protocol):

65 - """Factory that prepares inputs (OUTSIDE the graph) and returns a tuple of

66 - (forward_fn, attn_state). Called with warmup=True for the warmup pass and

67 - warmup=False for the captured pass."""

60 + """Factory that prepares inputs (OUTSIDE the graph) and returns a

61 + forward_fn. Called with warmup=True for the warmup pass and warmup=False

62 + for the captured pass."""

64 def __call__(

65 self,

66 desc: BatchExecutionDescriptor,

67 warmup: bool,

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

68 + ) -> Callable[[CUDAGraphMode], None]: ...

71 def _is_compatible(

@@ -301,19 +296,18 @@ class CudaGraphManager:

296 self,

297 create_forward_fn: CreateForwardFn,

298 progress_bar_desc: str = "Capturing CUDA graphs",

304 - ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:

299 + ) -> None:

300 """Capture CUDA graphs.

301

302 Args:

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

309 - returns a tuple of (forward_fn, attn_state). For FULL cudagraph

310 - mode, it is invoked once with warmup=True for the warmup pass,

311 - and again with warmup=False for the captured pass. For attention

312 - backends that perform lazy metadata initialization (e.g. FlashMLA),

313 - FULL cudagraph capture requires distinct metadatas for warmup and

314 - capture.

304 + returns a forward_fn. For FULL cudagraph mode, it is invoked

305 + once with warmup=True for the warmup pass, and again with

306 + warmup=False for the captured pass. For attention backends

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

308 + FULL cudagraph capture requires distinct metadatas for warmup

309 + and capture.

310 """

316 - attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {}

311 with graph_capture(device=self.device):

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

313 # activations so FULL activations should fit in already allocated

@@ -327,7 +321,7 @@ class CudaGraphManager:

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

322 for desc in descs:

323 # Prepare inputs and get forward function

330 - forward_fn, warmup_attn_state = create_forward_fn(desc, warmup=True)

324 + forward_fn = create_forward_fn(desc, warmup=True)

325

326 # Warmup

327 forward_fn(CUDAGraphMode.NONE)

@@ -337,18 +331,10 @@ class CudaGraphManager:

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

332 )

333 if desc.cg_mode == CUDAGraphMode.PIECEWISE:

340 - attn_states[desc] = AttentionStatePair(

341 - warmup_attn_state, warmup_attn_state

342 - )

334 forward_fn(CUDAGraphMode.PIECEWISE)

335 else:

336 # Capture with fresh attention state.

346 - forward_fn, capture_attn_state = create_forward_fn(

347 - desc, warmup=False

348 - )

349 - attn_states[desc] = AttentionStatePair(

350 - warmup_attn_state, capture_attn_state

351 - )

337 + forward_fn = create_forward_fn(desc, warmup=False)

338 assert desc not in self.graphs, (

339 f"Graph already captured for {desc}"

340 )

@@ -366,7 +352,6 @@ class CudaGraphManager:

352 self.graphs[desc] = graph

353 compilation_counter.num_cudagraph_captured += 1

354 self._graphs_captured = True

369 - return attn_states

355

356 def dispatch(

357 self,

@@ -461,7 +446,7 @@ class ModelCudaGraphManager(CudaGraphManager):

446 use_aux_hidden_state_outputs: bool = False,

447 lora_capture_hook: Callable[[int, int, int], None] | None = None,

448 progress_bar_desc: str = "Capturing CUDA graphs",

464 - ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:

449 + ) -> None:

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

451 self.use_aux_hidden_state_outputs = use_aux_hidden_state_outputs

452 if self.use_breakable_cg:

@@ -470,10 +455,7 @@ class ModelCudaGraphManager(CudaGraphManager):

455 def create_forward_fn(

456 desc: BatchExecutionDescriptor,

457 warmup: bool,

473 - ) -> tuple[

474 - Callable[[CUDAGraphMode], None],

475 - AttentionState,

476 - ]:

458 + ) -> Callable[[CUDAGraphMode], None]:

459 num_tokens = desc.num_tokens

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

461

@@ -571,9 +553,9 @@ class ModelCudaGraphManager(CudaGraphManager):

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

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

555

574 - return forward_fn, AttentionState(attn_metadata, slot_mappings)

556 + return forward_fn

557

576 - return super().capture(create_forward_fn, progress_bar_desc)

558 + super().capture(create_forward_fn, progress_bar_desc)

559

560 def run_fullgraph(

561 self, desc: BatchExecutionDescriptor

vllm/v1/worker/gpu/model_runner.py

+7 -3

@@ -479,7 +479,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

479 if isinstance(self.speculator, DraftModelSpeculator):

480 # HACK(woosuk)

481 self.speculator.set_attn(

482 - self.model_state, self.kv_cache_config, self.block_tables

482 + self.model_state,

483 + self.kv_cache_config,

484 + self.block_tables,

485 + self.input_buffers,

486 + self.attn_groups,

487 )

488

489 self.kv_caches: list[torch.Tensor] = []

@@ -706,7 +710,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

710 start_free_gpu_memory = torch.accelerator.get_memory_info()[0]

711

712 with self.maybe_setup_dummy_loras(self.lora_config):

709 - attn_states = self.cudagraph_manager.capture(

713 + self.cudagraph_manager.capture(

714 self.model,

715 self.model_state,

716 self.input_buffers,

@@ -719,7 +723,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

723 lora_capture_hook=create_lora_capture_hook(self.lora_config, self),

724 )

725 if self.speculator is not None:

722 - self.speculator.capture(attn_states)

726 + self.speculator.capture()

727

728 end_time = time.perf_counter()

729 end_free_gpu_memory = torch.accelerator.get_memory_info()[0]

vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py

+11 -45

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

11 BatchExecutionDescriptor,

12 CudaGraphManager,

13 prepare_inputs_to_capture,

@@ -19,45 +17,15 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState

17 from vllm.v1.worker.utils import AttentionGroup

22 -class PrefillSpeculatorCudaGraphManager(CudaGraphManager):

23 - """CudaGraphManager for draft prefill, using pre-built attention states

24 - from the target model's capture."""

20 +class SpeculatorCudaGraphManager(CudaGraphManager):

21 + """CudaGraphManager for draft prefill and decode.

26 - def capture(

27 - self,

28 - forward_fn: Callable,

29 - attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],

30 - progress_bar_desc: str = "Capturing CUDA graphs",

31 - ) -> None:

32 - def create_forward_fn(

33 - desc: BatchExecutionDescriptor,

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 = (

39 - torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")

40 - if self.dp_size > 1

41 - else None

42 - )

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,

48 - num_tokens,

49 - attn_metadata,

50 - slot_mappings,

51 - num_tokens_across_dp,

52 - cg_mode,

53 - )

54 - return fwd, attn_state

55 -

56 - super().capture(create_forward_fn, progress_bar_desc)

57 -

58 -

59 -class DecodeSpeculatorCudaGraphManager(CudaGraphManager):

60 - """CudaGraphManager for draft decode, building its own attention metadata."""

23 + Builds fresh dummy inputs and attention metadata for every warmup and

24 + capture pass so that the contents of the shared persistent buffers

25 + (e.g. query_start_loc, seq_lens, FA3 scheduler metadata) always match

26 + the batch descriptor being captured. Reusing metadata built during an

27 + earlier capture would execute kernels with stale buffer contents.

28 + """

30 def capture(

31 self,

@@ -72,7 +40,7 @@ class DecodeSpeculatorCudaGraphManager(CudaGraphManager):

40 def create_forward_fn(

41 desc: BatchExecutionDescriptor,

42 warmup: bool,

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

43 + ) -> Callable[[CUDAGraphMode], None]:

44 num_tokens = desc.num_tokens

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

46 num_tokens_across_dp = (

@@ -80,7 +48,7 @@ class DecodeSpeculatorCudaGraphManager(CudaGraphManager):

48 if self.dp_size > 1

49 else None

50 )

83 - attn_state = prepare_inputs_to_capture(

51 + attn_metadata, slot_mappings = prepare_inputs_to_capture(

52 num_reqs,

53 num_tokens,

54 model_state,

@@ -90,9 +58,8 @@ class DecodeSpeculatorCudaGraphManager(CudaGraphManager):

58 kv_cache_config,

59 skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),

60 )

93 - attn_metadata, slot_mappings = attn_state

95 - fwd = lambda cg_mode: forward_fn(

62 + return lambda cg_mode: forward_fn(

63 num_reqs,

64 num_tokens,

65 attn_metadata,

@@ -100,6 +67,5 @@ class DecodeSpeculatorCudaGraphManager(CudaGraphManager):

67 num_tokens_across_dp,

68 cg_mode,

69 )

103 - return fwd, attn_state

71 super().capture(create_forward_fn, progress_bar_desc)

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

+14 -12

@@ -12,15 +12,13 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

12 from vllm.triton_utils import tl, triton

13 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer

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

15 - AttentionStatePair,

15 BatchExecutionDescriptor,

16 get_uniform_token_count,

17 )

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

19 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

20 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (

22 - DecodeSpeculatorCudaGraphManager,

23 - PrefillSpeculatorCudaGraphManager,

21 + SpeculatorCudaGraphManager,

22 )

23 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator

@@ -47,8 +45,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

45 self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device

46 )

50 - self.prefill_cudagraph_manager: PrefillSpeculatorCudaGraphManager | None = None

51 - self.decode_cudagraph_manager: DecodeSpeculatorCudaGraphManager | None = None

48 + self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None

49 + self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None

51 @property

52 def advance_draft_positions(self) -> bool:

@@ -62,7 +60,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

61 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:

62 # Initialize cudagraph manager for draft prefill (draft position 0).

65 - self.prefill_cudagraph_manager = PrefillSpeculatorCudaGraphManager(

63 + self.prefill_cudagraph_manager = SpeculatorCudaGraphManager(

64 self.vllm_config,

65 self.device,

66 cudagraph_mode,

@@ -76,17 +74,14 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

74 cudagraph_mode = CUDAGraphMode.NONE

76 # Initialize cudagraph manager for draft decodes (draft positions > 0).

79 - self.decode_cudagraph_manager = DecodeSpeculatorCudaGraphManager(

77 + self.decode_cudagraph_manager = SpeculatorCudaGraphManager(

78 self.vllm_config,

79 self.device,

80 cudagraph_mode,

81 decode_query_len=1,

82 )

86 - def capture(

87 - self,

88 - attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],

89 - ) -> None:

84 + def capture(self) -> None:

85 logger.info("Capturing model for speculator...")

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

87 # dummy runs to cause out-of-bounds indexing during capture.

@@ -97,12 +92,19 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

92 # For FULL graphs, the entire routine is recorded as one graph.

93 # For PIECEWISE, only the model's compiled regions are captured

94 # and the rest (compute_logits, gumbel_sample) runs eagerly.

95 + # Draft prefill reuses the target model's attention metadata at

96 + # runtime, so capture builds its dummy metadata through the target

97 + # model runner's builders and buffers.

98 assert self.prefill_cudagraph_manager is not None

99 if self.prefill_cudagraph_manager.use_breakable_cg:

100 self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)

101 self.prefill_cudagraph_manager.capture(

102 self._prefill,

105 - attn_states,

103 + self.model_state,

104 + self.target_input_buffers,

105 + self.block_tables,

106 + self.target_attn_groups,

107 + self.kv_cache_config,

108 progress_bar_desc="Capturing prefill CUDA graphs",

109 )

110

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

+2 -3

@@ -77,7 +77,7 @@ class DFlashCudaGraphManager(CudaGraphManager):

77 def create_forward_fn(

78 desc: BatchExecutionDescriptor,

79 warmup: bool,

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

80 + ) -> Callable[[CUDAGraphMode], None]:

81 num_tokens = desc.num_tokens

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

83 num_tokens_across_dp = (

@@ -98,7 +98,7 @@ class DFlashCudaGraphManager(CudaGraphManager):

98 )

99 attn_metadata, slot_mappings = attn_state

100

101 - fwd = lambda cg_mode: forward_fn(

101 + return lambda cg_mode: forward_fn(

102 num_reqs,

103 num_tokens,

104 attn_metadata,

@@ -106,6 +106,5 @@ class DFlashCudaGraphManager(CudaGraphManager):

106 num_tokens_across_dp,

107 cg_mode,

108 )

109 - return fwd, attn_state

109

110 super().capture(create_forward_fn, progress_bar_desc)

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

+11 -2

@@ -25,6 +25,7 @@ from vllm.v1.worker.gpu.spec_decode.dflash.utils import (

25 )

26 from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator

27 from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id

28 +from vllm.v1.worker.utils import AttentionGroup

30 logger = init_logger(__name__)

@@ -97,7 +98,7 @@ class DFlashSpeculator(DraftModelSpeculator):

98 decode_query_len=self.num_query_per_req,

99 )

100

100 - def capture(self, attn_states: dict | None = None) -> None:

101 + def capture(self) -> None:

102 logger.info("Capturing model for %s speculator...", self._speculator_name)

103 # Reset sampling indices to zero to prevent stale values from prior

104 # dummy runs from being baked into the captured graph.

@@ -128,8 +129,16 @@ class DFlashSpeculator(DraftModelSpeculator):

129 model_state: ModelState,

130 kv_cache_config: KVCacheConfig,

131 block_tables: BlockTables,

132 + target_input_buffers: InputBuffers,

133 + target_attn_groups: list[list[AttentionGroup]],

134 ) -> None:

132 - super().set_attn(model_state, kv_cache_config, block_tables)

135 + super().set_attn(

136 + model_state,

137 + kv_cache_config,

138 + block_tables,

139 + target_input_buffers,

140 + target_attn_groups,

141 + )

142

143 self.draft_kv_cache_group_ids = [

144 gid for gid, g in enumerate(self.attn_groups) if g

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

+10 -8

@@ -18,13 +18,10 @@ from vllm.v1.worker.gpu.attn_utils import (

18 init_attn_backend,

19 )

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

21 -from vllm.v1.worker.gpu.cudagraph_utils import (

22 - AttentionStatePair,

23 - BatchExecutionDescriptor,

24 -)

21 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

22 from vllm.v1.worker.gpu.model_states.interface import ModelState

23 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

24 +from vllm.v1.worker.utils import AttentionGroup

26 logger = init_logger(__name__)

@@ -35,10 +32,7 @@ class BaseSpeculator(ABC):

32 pass

34 @abstractmethod

38 - def capture(

39 - self,

40 - attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],

41 - ) -> None:

35 + def capture(self) -> None:

36 pass

38 @abstractmethod

@@ -186,6 +180,8 @@ class DraftModelSpeculator(BaseSpeculator):

180 model_state: ModelState,

181 kv_cache_config: KVCacheConfig,

182 block_tables: BlockTables,

183 + target_input_buffers: InputBuffers,

184 + target_attn_groups: list[list[AttentionGroup]],

185 ) -> None:

186 self.model_state = model_state

187 self.kv_cache_config = kv_cache_config

@@ -196,6 +192,12 @@ class DraftModelSpeculator(BaseSpeculator):

192 active_layer_names=self.draft_attn_layer_names,

193 )

194 self.block_tables = block_tables

195 + # The target model runner's buffers and attention groups. Draft

196 + # prefill reuses the target model's attention metadata, so its

197 + # cudagraph capture must build dummy metadata through the same

198 + # builders and buffers.

199 + self.target_input_buffers = target_input_buffers

200 + self.target_attn_groups = target_attn_groups

201

202 def _build_draft_attn_metadata(

203 self,