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 12, 18:30 UTC
[BugFix] Fix Eagle `IndexError: list index out of range` for even `num_speculative_tokens` (#29102)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Nov 21, 2025 at 19:47 UTC
30d64662387aaa74abcee294f27b83043f2d1ae6
3 files changed
+37
-20
tests/conftest.py
+8
@@ -748,6 +748,14 @@ class VllmRunner:
748
# being captured which can trigger edge cases that we don't handle yet.
749
kwargs["compilation_config"] = {"cudagraph_capture_sizes": [4]}
750
751
+ # Make sure we have atleast one cudagraph large enough for a single decode.
752
+ if (speculative_config := kwargs.get("speculative_config")) and (
753
+ num_speculative_tokens := speculative_config["num_speculative_tokens"]
754
+ ):
755
+ kwargs["compilation_config"]["cudagraph_capture_sizes"].append(
756
+ num_speculative_tokens + 1
757
+ )
758
+
759
with init_ctx:
760
self.llm = LLM(
761
model=model_name,
vllm/config/compilation.py
+10
-6
@@ -950,14 +950,18 @@ class CompilationConfig:
950
)
951
)
952
953
+ if len(rounded_sizes) == 0 and multiple_of <= self.max_cudagraph_capture_size:
954
+ # if one valid but would be round_down use that
955
+ rounded_sizes = [multiple_of]
956
+
957
if len(rounded_sizes) == 0:
954
- logger.warning(
955
- "No valid cudagraph sizes after rounding to multiple of "
956
- " num_speculative_tokens + 1 (%d); please adjust num_speculative_tokens"
957
- " or max_cudagraph_capture_size (or cudagraph_capture_sizes)",
958
- multiple_of,
958
+ raise ValueError(
959
+ f"No valid cudagraph sizes after rounding to multiple of {multiple_of} "
960
+ f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)"
961
+ f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}"
962
+ f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})"
963
+ f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})"
964
)
960
- return
965
966
self.max_cudagraph_capture_size = rounded_sizes[-1]
967
self.cudagraph_capture_sizes = rounded_sizes
vllm/v1/spec_decode/eagle.py
+19
-14
@@ -86,9 +86,9 @@ class EagleProposer:
86
87
self.use_cuda_graph = False
88
89
- compilation_config = self.vllm_config.compilation_config
90
- if compilation_config.mode == CompilationMode.VLLM_COMPILE:
91
- cudagraph_mode = compilation_config.cudagraph_mode
89
+ self.compilation_config = self.vllm_config.compilation_config
90
+ if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:
91
+ cudagraph_mode = self.compilation_config.cudagraph_mode
92
if cudagraph_mode != CUDAGraphMode.NONE and not cudagraph_mode.has_mode(
93
CUDAGraphMode.PIECEWISE
94
):
@@ -103,13 +103,6 @@ class EagleProposer:
103
and not self.speculative_config.enforce_eager
104
)
105
106
- self.cudagraph_batch_sizes = (
107
- (sorted(self.vllm_config.compilation_config.cudagraph_capture_sizes))
108
- if self.use_cuda_graph
109
- else []
110
- )
111
-
112
- self.use_cuda_graph = self.use_cuda_graph and bool(self.cudagraph_batch_sizes)
106
# persistent buffers for cuda graph
107
self.input_ids = torch.zeros(
108
self.max_num_tokens, dtype=torch.int32, device=device
@@ -276,7 +269,10 @@ class EagleProposer:
269
per_layer_attn_metadata[layer_name] = draft_indexer_metadata
270
271
cudagraph_runtime_mode = CUDAGraphMode.NONE
279
- if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
272
+ if (
273
+ self.use_cuda_graph
274
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
275
+ ):
276
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
277
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
278
else:
@@ -366,7 +362,10 @@ class EagleProposer:
362
# Generate the remaining draft tokens.
363
draft_token_ids_list = [draft_token_ids]
364
369
- if self.use_cuda_graph and batch_size <= self.cudagraph_batch_sizes[-1]:
365
+ if (
366
+ self.use_cuda_graph
367
+ and batch_size <= self.compilation_config.max_cudagraph_capture_size
368
+ ):
369
input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size)
370
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
371
else:
@@ -777,7 +776,10 @@ class EagleProposer:
776
self.positions[:num_tokens] = tree_positions.view(-1)
777
self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)
778
780
- if self.use_cuda_graph and num_tokens <= self.cudagraph_batch_sizes[-1]:
779
+ if (
780
+ self.use_cuda_graph
781
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
782
+ ):
783
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
784
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
785
else:
@@ -1114,7 +1116,10 @@ class EagleProposer:
1116
) -> None:
1117
# Determine if CUDA graphs should be used for this run.
1118
cudagraphs_enabled = use_cudagraphs and self.use_cuda_graph
1117
- if cudagraphs_enabled and num_tokens <= self.cudagraph_batch_sizes[-1]:
1119
+ if (
1120
+ cudagraphs_enabled
1121
+ and num_tokens <= self.compilation_config.max_cudagraph_capture_size
1122
+ ):
1123
num_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)
1124
1125
with set_forward_context(