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 5, 04:45 UTC
[Model Runner V2] Disable piecewise cudagraph mode fallback for eagle draft decodes (#39773)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Apr 14, 2026 at 17:47 UTC
3bfe55a03758d57d4dde7db975842dc281b740a0
1 file changed
+19
-9
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+19
-9
@@ -115,13 +115,21 @@ class EagleSpeculator:
115
cudagraph_mode,
116
self.num_speculative_steps + 1,
117
)
118
- # Initialize cudagraph manager for draft generation (draft positions > 0).
118
+
119
+ # PIECEWISE cudagraphs are not supported for eagle draft decodes.
120
+ # PIECEWISE pads num_tokens to the next capture size without padding
121
+ # num_reqs, which can cause attention backends to read past the
122
+ # valid per-request metadata (e.g. FlashInfer's kv_indptr buffer).
123
+ if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
124
+ cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
125
+ else:
126
+ cudagraph_mode = CUDAGraphMode.NONE
127
+
128
+ # Initialize cudagraph manager for draft decodes (draft positions > 0).
129
self.decode_cudagraph_manager = EagleCudaGraphManager(
130
self.vllm_config,
131
self.device,
122
- # Only use FULL graph mode, if available, because draft decodes
123
- # only consist of a single token.
124
- cudagraph_mode.decode_mode(),
132
+ cudagraph_mode,
133
decode_query_len=1,
134
)
135
# Share a single pool between prefill and decode since they never
@@ -366,11 +374,8 @@ class EagleSpeculator:
374
375
# Capture the decode draft generation loop (model forward +
376
# compute_logits + gumbel_sample + update_eagle_inputs, for
369
- # each step).
370
- # For FULL graphs, the entire multi-step loop is recorded as
371
- # one graph. For PIECEWISE, only the model's compiled regions
372
- # are captured, and the rest (compute_logits, gumbel_sample,
373
- # update_eagle_inputs) runs eagerly.
377
+ # each step). For FULL graphs, the entire multi-step loop is
378
+ # recorded as one graph.
379
assert self.decode_cudagraph_manager is not None
380
self.decode_cudagraph_manager.capture(
381
self.generate_draft,
@@ -629,6 +634,11 @@ def _prepare_eagle_inputs_kernel(
634
block = i + tl.arange(0, BLOCK_SIZE)
635
mask = block < max_num_reqs
636
tl.store(eagle_seq_lens_ptr + block, 0, mask=mask)
637
+ # Pad last_token_indices for CUDA graphs.
638
+ for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
639
+ block = i + tl.arange(0, BLOCK_SIZE)
640
+ mask = block < max_num_reqs
641
+ tl.store(last_token_indices_ptr + block, 0, mask=mask)
642
643
644
def prepare_eagle_inputs(