@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 166 lines 5.64 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Truncation must bound the prompt *after* multimodal placeholders expand.
4
5 ``TokenizeParams.apply_post_tokenization`` truncates the unexpanded prompt, so
6 one placeholder token can expand into many and push the result back over
7 ``truncate_prompt_tokens`` (vllm-project/vllm#57038).
8 """
9
10 import numpy as np
11 import pytest
12 import torch
13
14 from vllm.config import ModelConfig, VllmConfig
15 from vllm.exceptions import VLLMValidationError
16 from vllm.multimodal.inputs import PlaceholderRange
17 from vllm.renderers import renderer_from_config
18 from vllm.renderers.base import BaseRenderer
19 from vllm.renderers.hf import _ensure_prompt_embeds_placeholder_token
20 from vllm.renderers.params import TokenizeParams
21
22
23 def _renderer():
24 """A BaseRenderer stub: only the tokenizer attribute is needed here."""
25
26 class _StubRenderer(BaseRenderer):
27 def __init__(self):
28 self.tokenizer = None
29
30 def render_messages(self, *args, **kwargs):
31 raise NotImplementedError
32
33 return _StubRenderer()
34
35
36 def _mm_input(num_tokens: int, offset: int, length: int, **extra):
37 return {
38 "type": "multimodal",
39 "prompt_token_ids": list(range(num_tokens)),
40 "mm_kwargs": {},
41 "mm_hashes": {},
42 "mm_placeholders": {"audio": [PlaceholderRange(offset=offset, length=length)]},
43 **extra,
44 }
45
46
47 def _params(limit: int | None, side: str = "right"):
48 return TokenizeParams(
49 max_total_tokens=4096,
50 truncate_prompt_tokens=limit,
51 truncation_side=side,
52 )
53
54
55 def test_expanded_prompt_is_bounded():
56 """The reported case: audio fits, trailing text is cut to the limit."""
57 engine_input = _mm_input(num_tokens=147, offset=1, length=52)
58
59 _renderer()._truncate_expanded_prompt(engine_input, _params(96))
60
61 assert engine_input["prompt_token_ids"] == list(range(96))
62 assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
63 offset=1, length=52
64 )
65
66
67 def test_prompt_within_limit_is_untouched():
68 engine_input = _mm_input(num_tokens=64, offset=1, length=52)
69
70 _renderer()._truncate_expanded_prompt(engine_input, _params(96))
71
72 assert engine_input["prompt_token_ids"] == list(range(64))
73
74
75 def test_split_multimodal_span_is_rejected():
76 """Cutting mid-span would desync features from token ids, so reject."""
77 engine_input = _mm_input(num_tokens=147, offset=80, length=50)
78
79 with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
80 _renderer()._truncate_expanded_prompt(engine_input, _params(96))
81
82
83 def test_dropped_multimodal_span_is_rejected():
84 """A span entirely past the cut would be discarded silently."""
85 engine_input = _mm_input(num_tokens=147, offset=100, length=20)
86
87 with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
88 _renderer()._truncate_expanded_prompt(engine_input, _params(96))
89
90
91 def test_left_truncation_shifts_placeholder_offsets():
92 engine_input = _mm_input(num_tokens=147, offset=100, length=20)
93
94 _renderer()._truncate_expanded_prompt(engine_input, _params(96, side="left"))
95
96 assert engine_input["prompt_token_ids"] == list(range(51, 147))
97 # The span started at 100 and 51 leading tokens were dropped.
98 assert engine_input["mm_placeholders"]["audio"][0] == PlaceholderRange(
99 offset=49, length=20
100 )
101
102
103 def test_text_only_input_is_untouched():
104 """Only multimodal prompts grow after tokenization, so nothing else runs."""
105 engine_input = {
106 "type": "token",
107 "prompt_token_ids": list(range(147)),
108 "prompt_token_offsets": [(i, i + 1) for i in range(147)],
109 }
110
111 _renderer()._truncate_expanded_prompt(engine_input, _params(96))
112
113 assert engine_input["prompt_token_ids"] == list(range(147))
114 assert len(engine_input["prompt_token_offsets"]) == 147
115
116
117 def test_no_truncation_requested_is_noop():
118 engine_input = _mm_input(num_tokens=147, offset=1, length=52)
119
120 _renderer()._truncate_expanded_prompt(engine_input, _params(None))
121
122 assert engine_input["prompt_token_ids"] == list(range(147))
123
124
125 def test_mixed_prompt_embeds_truncation_covers_embeds_spans():
126 """Bound in _process_singleton, where prompt_embeds spans are already attached."""
127 model_config = ModelConfig(
128 "Qwen/Qwen3-ASR-0.6B-hf",
129 revision="7f1569a48a89f3e3f4dc3a5c9d28bddd903bc76c",
130 dtype="float32",
131 max_model_len=4096,
132 mm_processor_cache_gb=0,
133 enable_prompt_embeds=True,
134 )
135 renderer = renderer_from_config(VllmConfig(model_config=model_config))
136 tok = renderer.tokenizer
137 pe_id = _ensure_prompt_embeds_placeholder_token(tok)
138 hidden = model_config.get_hidden_size()
139 audio = np.zeros(4 * 16000, dtype=np.float32)
140
141 pe_len = 20
142 tensors = [torch.zeros(pe_len, hidden, dtype=torch.float32)]
143
144 def _build_prompt():
145 text = "<|audio_start|><|audio_pad|><|audio_end|>"
146 ids = tok.encode(text) + [pe_id] + tok.encode(" tail" * 40)
147 return {
148 "prompt_token_ids": ids,
149 "multi_modal_data": {"audio": (audio, 16000)},
150 "_prompt_embeds": (tensors, pe_id),
151 }
152
153 # A limit past every span (audio ends at 52, prompt_embeds at 74):
154 # truncation should succeed and bound the prompt.
155 out = renderer._process_singleton(
156 _build_prompt(),
157 tok_params=_params(79),
158 )
159 assert len(out["prompt_token_ids"]) == 79
160
161 # A limit inside the prompt_embeds span (54..74): must be rejected.
162 with pytest.raises(VLLMValidationError, match="truncate_prompt_tokens"):
163 renderer._process_singleton(
164 _build_prompt(),
165 tok_params=_params(64),
166 )