@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 47604137a2
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 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: