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 14, 05:30 UTC
[Core] Enable decode of context length equal to max model length (#26168)
Signed-off-by: Yannick Schnider <yannick.schnider1@ibm.com>
Yannick Schnider committed
Oct 4, 2025 at 11:59 UTC
f05fea1f5e91007bc7c53a9e3178ffa9becaa403
4 files changed
+32
-15
tests/entrypoints/llm/test_generate.py
+3
-2
@@ -82,10 +82,11 @@ def test_max_model_len():
82
for output in outputs:
83
num_total_tokens = len(output.prompt_token_ids) + len(
84
output.outputs[0].token_ids)
85
- # Total tokens must not exceed max_model_len.
85
+ # Total tokens must not exceed max_model_len + 1 (the last token can be
86
+ # generated with the context length equal to the max model length)
87
# It can be less if generation finishes due to other reasons (e.g., EOS)
88
# before reaching the absolute model length limit.
88
- assert num_total_tokens <= max_model_len
89
+ assert num_total_tokens <= max_model_len + 1
90
91
92
def test_log_stats():
tests/v1/e2e/test_context_length.py
+27
-11
@@ -4,15 +4,22 @@
4
end-to-end tests for context length corner cases of vLLM v1 model runner
5
versus HuggingFace's transformers.
6
7
-This test verifies the following behavior: allow a prefill that fills the
8
-model's maximum context length and then request a single new token.
7
+This test verifies the following behavior: allow prefill and decodes on the
8
+model's maximum context length ``max_model_len`` and get one more token.
9
10
Test strategy
11
-- Build a textual prompt that tokenizes to exactly ``max_model_len`` tokens.
12
-- Run vLLM generation requesting a single new token (max_tokens=1).
13
-- Run HF generation on the same prompt requesting a single token too.
11
+- Build a prompt consisting of exactly ``prompt_len`` tokens.
12
+- Run vLLM generation requesting ``max_tokens`` new tokens.
13
+- Run HF generation on the same prompt requesting the same number of tokens.
14
- Assert both return the same number of generated tokens and the same ids.
15
16
+Test cases
17
+- Prefill a prompt of ``max_model_len`` (2048) and request a single token which
18
+will be sampled after the prefill (context length ``max_model_len``).
19
+- Prefill a prompt of ``max_model_len`` - 1 (2047) and request two tokens where
20
+the 1st will be sampled after the prefill and the 2nd after the first decode
21
+(context length ``max_model_len``).
22
+
23
"""
24
25
import pytest
@@ -27,11 +34,16 @@ from vllm.inputs import TokensPrompt
34
35
@create_new_process_for_each_test()
36
@pytest.mark.parametrize("model", ["JackFram/llama-160m"])
30
-@pytest.mark.parametrize("max_model_len", [2048])
31
-@pytest.mark.parametrize("max_tokens", [1])
32
-def test_prefill_max_context_length(
37
+@pytest.mark.parametrize(
38
+ "prompt_len, max_tokens",
39
+ [
40
+ (2048, 1), # prompt_len = max_model_len
41
+ (2047, 2), # prompt_len = max_model_len - 1
42
+ ],
43
+)
44
+def test_max_context_length(
45
model: str,
34
- max_model_len: int,
46
+ prompt_len: int,
47
max_tokens: int,
48
) -> None:
49
"""Compare vLLM and HuggingFace when the prompt already fills the
@@ -42,8 +54,8 @@ def test_prefill_max_context_length(
54
single token when given the same inputs.
55
"""
56
45
- # Construct a prompt of size max_model_len
46
- prompt_ids = [[43] * max_model_len]
57
+ # Construct a prompt of size prompt_len
58
+ prompt_ids = [[43] * prompt_len]
59
60
# Generate max_tokens new tokens deterministically.
61
sampling_params = [
@@ -54,6 +66,7 @@ def test_prefill_max_context_length(
66
llm = LLM(
67
model=model,
68
tokenizer=model,
69
+ max_model_len=2048,
70
max_num_seqs=1,
71
tensor_parallel_size=1,
72
)
@@ -81,6 +94,9 @@ def test_prefill_max_context_length(
94
# HF returns the prompt + generated tokens. Slice off the prompt.
95
hf_output_ids = hf_generated.cpu().tolist()[0][len(prompt_ids[0]):]
96
97
+ # check that exactly max_tokens tokens were generated with vLLM and HF
98
+ assert len(vllm_output_ids) == len(hf_output_ids) == max_tokens
99
+
100
# check that vLLM outputs (token ids) match HF outputs
101
# Note: for simplicity don't pass detokenized string
102
check_outputs_equal(
vllm/v1/core/sched/scheduler.py
+1
-1
@@ -224,7 +224,7 @@ class Scheduler(SchedulerInterface):
224
# This is necessary when using spec decoding.
225
num_new_tokens = min(
226
num_new_tokens,
227
- self.max_model_len - 1 - request.num_computed_tokens)
227
+ self.max_model_len - request.num_computed_tokens)
228
229
# Schedule encoder inputs.
230
encoder_inputs_to_schedule = None
vllm/v1/core/sched/utils.py
+1
-1
@@ -43,7 +43,7 @@ def remove_all(lst: list, items_to_remove: set) -> list:
43
def check_stop(request: Request,
44
max_model_len: int,
45
pooler_output: Optional[torch.Tensor] = None) -> bool:
46
- if (request.num_tokens >= max_model_len
46
+ if (request.num_tokens > max_model_len
47
or request.num_output_tokens >= request.max_tokens):
48
request.status = RequestStatus.FINISHED_LENGTH_CAPPED
49
return True