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 13, 14:30 UTC
[Bugfix] Spec decode + structured output + spec model max len edge case (#28298)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Nov 8, 2025 at 19:44 UTC
47604137a26f825295e6e9cc7314dea2bd4f7eef
3 files changed
+36
-8
tests/v1/spec_decode/test_max_len.py
+30
-3
@@ -7,6 +7,7 @@ import pytest
7
from tests.utils import get_attn_backend_list_based_on_platform
8
from vllm import LLM, SamplingParams
9
from vllm.platforms import current_platform
10
+from vllm.sampling_params import StructuredOutputsParams
11
12
_PROMPTS = [
13
"1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1",
@@ -56,8 +57,34 @@ def test_eagle_max_len(
57
"method": "eagle",
58
"model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
59
"num_speculative_tokens": num_speculative_tokens,
60
+ "max_model_len": 80,
61
},
60
- max_model_len=100,
62
+ max_model_len=200,
63
)
62
- sampling_params = SamplingParams(max_tokens=100, ignore_eos=True)
63
- llm.generate(_PROMPTS, sampling_params)
64
+ sampling_params = SamplingParams(max_tokens=200, ignore_eos=True)
65
+ outputs = llm.generate(_PROMPTS, sampling_params)
66
+ for o in outputs:
67
+ assert o.outputs[0].finish_reason == "length", (
68
+ "This test is only meaningful if the output "
69
+ "is truncated due to max length"
70
+ )
71
+
72
+ sampling_params = SamplingParams(
73
+ max_tokens=200,
74
+ structured_outputs=StructuredOutputsParams(
75
+ regex="^" + "a b c d e " * 15 + "$"
76
+ ),
77
+ )
78
+ output = llm.generate(_PROMPTS, sampling_params)
79
+ for o in output:
80
+ assert o.prompt_token_ids is not None
81
+ assert (
82
+ len(o.prompt_token_ids)
83
+ < 80
84
+ < len(o.prompt_token_ids) + len(o.outputs[0].token_ids)
85
+ < 200
86
+ ), (
87
+ "This test is only meaningful if the output "
88
+ "is longer than the eagle max length"
89
+ )
90
+ assert o.outputs[0].text == "a b c d e " * 15
vllm/v1/core/sched/scheduler.py
+4
-4
@@ -325,6 +325,9 @@ class Scheduler(SchedulerInterface):
325
scheduled_spec_decode_tokens[request.request_id] = (
326
request.spec_token_ids
327
)
328
+ # New spec tokens will be set in `update_draft_token_ids` before the
329
+ # next step when applicable.
330
+ request.spec_token_ids = []
331
332
# Encoder-related.
333
if encoder_inputs_to_schedule:
@@ -1149,10 +1152,7 @@ class Scheduler(SchedulerInterface):
1152
continue
1153
1154
# Add newly generated spec token ids to the request.
1152
- if not spec_token_ids:
1153
- # NOTE(woosuk): request.spec_token_ids should be updated.
1154
- request.spec_token_ids.clear()
1155
- elif self.structured_output_manager.should_advance(request):
1155
+ if self.structured_output_manager.should_advance(request):
1156
metadata = request.structured_output_request
1157
request.spec_token_ids = metadata.grammar.validate_tokens( # type: ignore[union-attr]
1158
spec_token_ids
vllm/v1/structured_output/__init__.py
+2
-1
@@ -269,9 +269,10 @@ class StructuredOutputManager:
269
and token is not None
270
and not structured_output_request.grammar.is_terminated()
271
):
272
- assert structured_output_request.grammar.accept_tokens(
272
+ accepted = structured_output_request.grammar.accept_tokens(
273
req_id, [token]
274
)
275
+ assert accepted, (token, req_id, scheduled_spec_decode_tokens)
276
state_advancements += 1
277
cumulative_index += 1
278
if state_advancements > 0: