@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b64a2708b0

Bound generation inputs before expensive work (#51447)

Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Clinton Thomas committed Aug 11, 2026 at 10:03 UTC b64a2708b06b6329420702bffce994543c1ec6d2
11 files changed +421 -37
rust/src/server/src/routes/openai/utils/types.rs
+47 -2
@@ -54,9 +54,23 @@ impl StringOrArray {
54 }
55 }
56
57 -/// Validates stop sequences (non-empty strings)
57 +fn max_stop_strings() -> usize {
58 + std::env::var("VLLM_MAX_STOP_STRINGS")
59 + .ok()
60 + .and_then(|value| value.parse().ok())
61 + .unwrap_or(4)
62 +}
63 +
64 +/// Validates stop sequences (at most the configured number of non-empty strings).
65 pub fn validate_stop(stop: &StringOrArray) -> Result<(), validator::ValidationError> {
59 - if stop.as_slice().iter().any(|s| s.is_empty()) {
66 + let stop = stop.as_slice();
67 + let max_stop_strings = max_stop_strings();
68 + if stop.len() > max_stop_strings {
69 + let mut error = validator::ValidationError::new("too_many_stop_strings");
70 + error.code = format!("stop strings must contain at most {max_stop_strings} items").into();
71 + return Err(error);
72 + }
73 + if stop.iter().any(|s| s.is_empty()) {
74 return Err(validator::ValidationError::new(
75 "stop strings cannot be empty",
76 ));
@@ -64,6 +78,37 @@ pub fn validate_stop(stop: &StringOrArray) -> Result<(), validator::ValidationEr
78 Ok(())
79 }
80
81 +#[cfg(test)]
82 +mod tests {
83 + use super::{StringOrArray, validate_stop};
84 +
85 + #[test]
86 + fn validate_stop_accepts_four_stop_strings() {
87 + let stop = StringOrArray::Array(
88 + ["one", "two", "three", "four"].into_iter().map(str::to_string).collect(),
89 + );
90 +
91 + validate_stop(&stop).expect("four stop strings should be accepted");
92 + }
93 +
94 + #[test]
95 + fn validate_stop_rejects_more_than_four_stop_strings() {
96 + let stop = StringOrArray::Array(
97 + ["one", "two", "three", "four", "five"]
98 + .into_iter()
99 + .map(str::to_string)
100 + .collect(),
101 + );
102 +
103 + let error = validate_stop(&stop).expect_err("five stop strings should be rejected");
104 +
105 + assert_eq!(
106 + error.code.as_ref(),
107 + "stop strings must contain at most 4 items"
108 + );
109 + }
110 +}
111 +
112 // ============================================================================
113 // Validation helpers
114 // ============================================================================
tests/test_request_input_bounds.py new
+273
@@ -0,0 +1,273 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Bounds on request-controlled inputs that would otherwise amplify work.
4 +
5 +Consolidates the regression tests for the request-input amplification fixes:
6 +stop-string caps, bad-words dedup/tokenization limit, stop-token-id dedup,
7 +beam-width/sequence caps, and the DeepSeek history-scan bound.
8 +"""
9 +
10 +import os
11 +import subprocess
12 +import sys
13 +from collections.abc import Callable
14 +from typing import Protocol
15 +
16 +import pytest
17 +from pydantic import ValidationError
18 +
19 +import vllm.envs as envs
20 +from vllm import SamplingParams
21 +from vllm.entrypoints.openai.chat_completion.protocol import (
22 + BatchChatCompletionRequest,
23 + ChatCompletionRequest,
24 +)
25 +from vllm.entrypoints.openai.completion.protocol import CompletionRequest
26 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
27 +from vllm.exceptions import VLLMValidationError
28 +from vllm.sampling_params import BeamSearchParams
29 +from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
30 +
31 +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
32 +
33 +
34 +# --- Stop strings: public requests cap the number of stop strings ---------
35 +
36 +
37 +class _StopRequest(Protocol):
38 + stop: str | list[str] | None
39 +
40 +
41 +def _completion_request(stop: list[str]) -> _StopRequest:
42 + return CompletionRequest(model="test-model", prompt="hello", stop=stop)
43 +
44 +
45 +def _chat_request(stop: list[str]) -> _StopRequest:
46 + return ChatCompletionRequest(
47 + model="test-model",
48 + messages=[{"role": "user", "content": "hello"}],
49 + stop=stop,
50 + )
51 +
52 +
53 +def _batch_chat_request(stop: list[str]) -> _StopRequest:
54 + return BatchChatCompletionRequest(
55 + model="test-model",
56 + messages=[[{"role": "user", "content": "hello"}]],
57 + stop=stop,
58 + )
59 +
60 +
61 +def _responses_request(stop: list[str]) -> _StopRequest:
62 + return ResponsesRequest(model="test-model", input="hello", stop=stop)
63 +
64 +
65 +REQUEST_BUILDERS: list[Callable[[list[str]], _StopRequest]] = [
66 + _completion_request,
67 + _chat_request,
68 + _batch_chat_request,
69 + _responses_request,
70 +]
71 +
72 +
73 +@pytest.mark.parametrize("build_request", REQUEST_BUILDERS)
74 +def test_public_requests_accept_four_stop_strings(
75 + build_request: Callable[[list[str]], _StopRequest],
76 +):
77 + stop = ["one", "two", "three", "four"]
78 +
79 + request = build_request(stop)
80 +
81 + assert request.stop == stop
82 +
83 +
84 +@pytest.mark.parametrize("build_request", REQUEST_BUILDERS)
85 +def test_public_requests_reject_more_than_four_stop_strings(
86 + build_request: Callable[[list[str]], _StopRequest],
87 +):
88 + with pytest.raises(ValidationError, match="at most 4"):
89 + build_request(["one", "two", "three", "four", "five"])
90 +
91 +
92 +def test_stop_string_limit_can_be_overridden():
93 + env = os.environ.copy()
94 + env["VLLM_MAX_STOP_STRINGS"] = "1"
95 + code = """
96 +from pydantic import ValidationError
97 +from vllm.entrypoints.openai.completion.protocol import CompletionRequest
98 +
99 +try:
100 + CompletionRequest(
101 + model="test-model",
102 + prompt="hello",
103 + stop=["one", "two"],
104 + )
105 +except ValidationError as error:
106 + assert "at most 1" in str(error)
107 +else:
108 + raise AssertionError("configured stop-string limit was not enforced")
109 +"""
110 +
111 + subprocess.run([sys.executable, "-c", code], check=True, env=env)
112 +
113 +
114 +# --- Stop token ids: duplicates are deduplicated in order ------------------
115 +
116 +
117 +def test_duplicate_stop_token_ids_are_deduplicated_in_order():
118 + params = SamplingParams(stop_token_ids=[42, 7, 42, 9, 7])
119 +
120 + assert params.stop_token_ids == [42, 7, 9]
121 + assert params.all_stop_token_ids == {7, 9, 42}
122 +
123 +
124 +# --- Bad words: dedup, and the tokenization pass is bounded ----------------
125 +
126 +
127 +class MockTokenizer:
128 + max_token_id = 1024
129 +
130 + def __init__(self):
131 + self.calls = 0
132 +
133 + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
134 + self.calls += 1
135 + return [2] if text.startswith(" ") else [1]
136 +
137 +
138 +def test_duplicate_bad_words_are_deduplicated_in_order():
139 + params = SamplingParams(bad_words=["bad", "worse", "bad", "worst"])
140 +
141 + assert params.bad_words == ["bad", "worse", "worst"]
142 +
143 +
144 +def test_bad_word_tokenization_stops_at_worker_limit():
145 + params = SamplingParams(bad_words=[f"word-{i}" for i in range(65)])
146 + tokenizer = MockTokenizer()
147 +
148 + with pytest.raises(VLLMValidationError, match="Too many bad words"):
149 + params.update_from_tokenizer(tokenizer)
150 +
151 + assert tokenizer.calls == 129
152 +
153 +
154 +def test_bad_word_tokenization_limit_can_be_overridden(monkeypatch):
155 + monkeypatch.setenv("VLLM_MAX_NUM_BAD_WORDS", "2")
156 + params = SamplingParams(bad_words=["bad", "worse"])
157 + tokenizer = MockTokenizer()
158 +
159 + with pytest.raises(VLLMValidationError, match="The max number is 2"):
160 + params.update_from_tokenizer(tokenizer)
161 +
162 + assert tokenizer.calls == 3
163 +
164 +
165 +# --- Beam search: beam width / n honor the sequence cap --------------------
166 +
167 +
168 +def _set_max_n(monkeypatch: pytest.MonkeyPatch, value: int) -> None:
169 + monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", str(value))
170 + if hasattr(envs.__getattr__, "cache_clear"):
171 + envs.__getattr__.cache_clear()
172 +
173 +
174 +def test_direct_beam_width_rejects_values_over_sequence_cap(
175 + monkeypatch: pytest.MonkeyPatch,
176 +):
177 + _set_max_n(monkeypatch, 4)
178 +
179 + with pytest.raises(VLLMValidationError, match="beam_width must be at most 4"):
180 + BeamSearchParams(beam_width=5, max_tokens=1)
181 +
182 +
183 +def test_chat_beam_conversion_rejects_n_before_stream_state_allocation(
184 + monkeypatch: pytest.MonkeyPatch,
185 +):
186 + _set_max_n(monkeypatch, 4)
187 + request = ChatCompletionRequest(
188 + model="test-model",
189 + messages=[{"role": "user", "content": "test"}],
190 + n=5,
191 + stream=True,
192 + use_beam_search=True,
193 + max_tokens=1,
194 + )
195 +
196 + with pytest.raises(VLLMValidationError, match="beam_width must be at most 4"):
197 + request.to_beam_search_params(max_tokens=1, default_sampling_params={})
198 +
199 +
200 +def test_chat_beam_conversion_accepts_n_at_sequence_cap(
201 + monkeypatch: pytest.MonkeyPatch,
202 +):
203 + _set_max_n(monkeypatch, 4)
204 + request = ChatCompletionRequest(
205 + model="test-model",
206 + messages=[{"role": "user", "content": "test"}],
207 + n=4,
208 + stream=True,
209 + use_beam_search=True,
210 + max_tokens=1,
211 + )
212 +
213 + params = request.to_beam_search_params(max_tokens=1, default_sampling_params={})
214 +
215 + assert params.beam_width == 4
216 +
217 +
218 +# --- DeepSeek encoders: the last-user scan runs once per conversation ------
219 +
220 +
221 +ENCODING_MODULES = [deepseek_v32_encoding, deepseek_v4_encoding]
222 +
223 +
224 +@pytest.mark.parametrize(
225 + "encoding_module",
226 + ENCODING_MODULES,
227 + ids=["deepseek_v32", "deepseek_v4"],
228 +)
229 +def test_encode_messages_scans_last_user_once_per_conversation(
230 + monkeypatch: pytest.MonkeyPatch,
231 + encoding_module,
232 +):
233 + calls = 0
234 + original_find_last_user_index = encoding_module.find_last_user_index
235 +
236 + def counted_find_last_user_index(messages):
237 + nonlocal calls
238 + calls += 1
239 + return original_find_last_user_index(messages)
240 +
241 + monkeypatch.setattr(
242 + encoding_module,
243 + "find_last_user_index",
244 + counted_find_last_user_index,
245 + )
246 +
247 + messages = [{"role": "user", "content": "Hello"}]
248 + messages.extend({"role": "assistant", "content": "Hi"} for _ in range(8))
249 +
250 + encoding_module.encode_messages(messages, thinking_mode="chat")
251 +
252 + assert calls == 1
253 +
254 +
255 +@pytest.mark.parametrize(
256 + "encoding_module",
257 + ENCODING_MODULES,
258 + ids=["deepseek_v32", "deepseek_v4"],
259 +)
260 +def test_encode_messages_preserves_small_chat_prompt(encoding_module):
261 + prompt = encoding_module.encode_messages(
262 + [
263 + {"role": "user", "content": "Hello"},
264 + {"role": "assistant", "content": "Hi"},
265 + {"role": "assistant", "content": "Again"},
266 + ],
267 + thinking_mode="chat",
268 + )
269 +
270 + assert prompt == (
271 + "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜></think>"
272 + "Hi<｜end▁of▁sentence｜>Again<｜end▁of▁sentence｜>"
273 + )
vllm/entrypoints/openai/chat_completion/protocol.py
+3 -2
@@ -30,6 +30,7 @@ from vllm.entrypoints.openai.engine.protocol import (
30 FunctionDefinition,
31 OpenAIBaseModel,
32 PerRequestTimingMetrics,
33 + StopParam,
34 StreamOptions,
35 ToolCall,
36 UsageInfo,
@@ -228,7 +229,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
229 presence_penalty: float | None = 0.0
230 response_format: AnyResponseFormat | None = None
231 seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
231 - stop: str | list[str] | None = []
232 + stop: StopParam = []
233 stream: bool | None = False
234 stream_options: StreamOptions | None = None
235 temperature: float | None = None
@@ -1077,7 +1078,7 @@ class BatchChatCompletionRequest(OpenAIBaseModel):
1078 presence_penalty: float | None = 0.0
1079 response_format: Any | None = None
1080 seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
1080 - stop: str | list[str] | None = Field(default_factory=list)
1081 + stop: StopParam = Field(default_factory=list)
1082 temperature: float | None = None
1083 top_p: float | None = None
1084 user: str | None = None
vllm/entrypoints/openai/completion/protocol.py
+2 -1
@@ -14,6 +14,7 @@ from vllm.entrypoints.openai.engine.protocol import (
14 AnyResponseFormat,
15 OpenAIBaseModel,
16 PerRequestTimingMetrics,
17 + StopParam,
18 StreamOptions,
19 UsageInfo,
20 structured_outputs_from_response_format,
@@ -61,7 +62,7 @@ class CompletionRequest(OpenAIBaseModel):
62 n: int = 1
63 presence_penalty: float | None = 0.0
64 seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
64 - stop: str | list[str] | None = []
65 + stop: StopParam = []
66 stream: bool | None = False
67 stream_options: StreamOptions | None = None
68 suffix: str | None = None
vllm/entrypoints/openai/engine/protocol.py
+6 -1
@@ -6,7 +6,7 @@
6 import json
7 import time
8 from http import HTTPStatus
9 -from typing import Any, ClassVar, Literal, TypeAlias
9 +from typing import Annotated, Any, ClassVar, Literal, TypeAlias
10
11 import regex as re
12 from pydantic import (
@@ -17,6 +17,7 @@ from pydantic import (
17 model_validator,
18 )
19
20 +import vllm.envs as envs
21 from vllm.config.utils import replace
22 from vllm.entrypoints.chat_utils import make_tool_call_id
23 from vllm.exceptions import VLLMServerError, VLLMValidationError
@@ -27,6 +28,10 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
28
29 logger = init_logger(__name__)
30
31 +StopParam: TypeAlias = (
32 + str | Annotated[list[str], Field(max_length=envs.VLLM_MAX_STOP_STRINGS)] | None
33 +)
34 +
35
36 class OpenAIBaseModel(BaseModel):
37 # OpenAI API does allow extra fields
vllm/entrypoints/openai/responses/protocol.py
+2 -2
@@ -60,7 +60,7 @@ from vllm.entrypoints.chat_utils import (
60 ChatCompletionMessageParam,
61 ChatTemplateContentFormatOption,
62 )
63 -from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
63 +from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, StopParam
64 from vllm.exceptions import VLLMValidationError
65 from vllm.logger import init_logger
66 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
@@ -281,7 +281,7 @@ class ResponsesRequest(OpenAIBaseModel):
281
282 repetition_penalty: float | None = None
283 seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
284 - stop: str | list[str] | None = []
284 + stop: StopParam = []
285 ignore_eos: bool = False
286 vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(
287 default=None,
vllm/envs.py
+14
@@ -108,6 +108,9 @@ if TYPE_CHECKING:
108 VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds
109 VLLM_MAX_N_SEQUENCES: int = 16384
110 VLLM_MAX_COMPLETION_PROMPTS: int = 1024
111 + VLLM_MAX_STOP_STRINGS: int = 4
112 + VLLM_MAX_NUM_BAD_WORDS: int = 128
113 + VLLM_MAX_BAD_WORDS_TOTAL_TOKENS: int = 1024
114 VLLM_PLUGINS: list[str] | None = None
115 VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None
116 VLLM_LORA_RESOLVER_HF_REPO_LIST: str | None = None
@@ -1114,6 +1117,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
1117 "VLLM_MAX_COMPLETION_PROMPTS": lambda: int(
1118 os.environ.get("VLLM_MAX_COMPLETION_PROMPTS", "1024")
1119 ),
1120 + # Maximum number of stop strings allowed in a single request.
1121 + "VLLM_MAX_STOP_STRINGS": lambda: int(os.environ.get("VLLM_MAX_STOP_STRINGS", "4")),
1122 + # Maximum number of bad-word token sequences generated per request.
1123 + "VLLM_MAX_NUM_BAD_WORDS": lambda: int(
1124 + os.environ.get("VLLM_MAX_NUM_BAD_WORDS", "128")
1125 + ),
1126 + # Maximum total number of bad-word tokens (summed across all bad words)
1127 + # allowed per request. Bounds the per-request GPU buffer width.
1128 + "VLLM_MAX_BAD_WORDS_TOTAL_TOKENS": lambda: int(
1129 + os.environ.get("VLLM_MAX_BAD_WORDS_TOTAL_TOKENS", "1024")
1130 + ),
1131 # a list of plugin names to load, separated by commas.
1132 # if this is not set, it means all plugins will be loaded
1133 # if this is set to an empty string, no plugins will be loaded
vllm/sampling_params.py
+33 -13
@@ -32,6 +32,22 @@ MAX_LOGPROB_TOKEN_IDS = 128
32 the per-request row width allocated by the sampler's `LogprobTokenIdsState`."""
33
34
35 +def _verify_num_sequences(value: int, parameter_name: str) -> None:
36 + if not isinstance(value, int):
37 + raise VLLMValidationError(
38 + f"{parameter_name} must be an int, but is of type {type(value)}"
39 + )
40 + if value < 1:
41 + raise VLLMValidationError(f"{parameter_name} must be at least 1, got {value}.")
42 + max_n = envs.VLLM_MAX_N_SEQUENCES
43 + if value > max_n:
44 + raise VLLMValidationError(
45 + f"{parameter_name} must be at most {max_n}, got {value}. "
46 + "To increase this limit, set the VLLM_MAX_N_SEQUENCES "
47 + "environment variable."
48 + )
49 +
50 +
51 def validate_thinking_token_budget(value: int | float | bool | None) -> int | None:
52 """Validate ``thinking_token_budget``; return ``None`` if unset."""
53 if value is None:
@@ -479,9 +495,13 @@ class SamplingParams(
495
496 if self.stop_token_ids is None:
497 self.stop_token_ids = []
498 + else:
499 + self.stop_token_ids = list(dict.fromkeys(self.stop_token_ids))
500
501 if self.bad_words is None:
502 self.bad_words = []
503 + else:
504 + self.bad_words = list(dict.fromkeys(self.bad_words))
505
506 if self.logprobs is True:
507 self.logprobs = 1
@@ -513,19 +533,7 @@ class SamplingParams(
533 self.skip_reading_prefix_cache = self.prompt_logprobs is not None
534
535 def _verify_args(self) -> None:
516 - if not isinstance(self.n, int):
517 - raise VLLMValidationError(
518 - f"n must be an int, but is of type {type(self.n)}"
519 - )
520 - if self.n < 1:
521 - raise VLLMValidationError(f"n must be at least 1, got {self.n}.")
522 - max_n = envs.VLLM_MAX_N_SEQUENCES
523 - if self.n > max_n:
524 - raise VLLMValidationError(
525 - f"n must be at most {max_n}, got {self.n}. "
526 - "To increase this limit, set the VLLM_MAX_N_SEQUENCES "
527 - "environment variable."
528 - )
536 + _verify_num_sequences(self.n, "n")
537 if not -2.0 <= self.presence_penalty <= 2.0:
538 raise VLLMValidationError(
539 f"presence_penalty must be in [-2, 2], got {self.presence_penalty}."
@@ -677,6 +685,7 @@ class SamplingParams(
685 if not self.bad_words:
686 return
687 self._bad_words_token_ids = []
688 + max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
689 for bad_word in self.bad_words:
690 # To prohibit words both at the beginning
691 # and in the middle of text
@@ -696,6 +705,14 @@ class SamplingParams(
705 and len(prompt_token_ids) == len(self._bad_words_token_ids[-1])
706 ):
707 self._bad_words_token_ids.append(prompt_token_ids)
708 + if len(self._bad_words_token_ids) > max_num_bad_words:
709 + raise VLLMValidationError(
710 + f"Too many bad words after tokenization: "
711 + f"{len(self._bad_words_token_ids)}. "
712 + f"The max number is {max_num_bad_words}.",
713 + parameter="bad_words",
714 + value=self.bad_words,
715 + )
716
717 invalid_token_ids = [
718 token_id
@@ -1147,3 +1164,6 @@ class BeamSearchParams(
1164 length_penalty: float = 1.0
1165 include_stop_str_in_output: bool = False
1166 structured_outputs: StructuredOutputsParams | None = None
1167 +
1168 + def __post_init__(self) -> None:
1169 + _verify_num_sequences(self.beam_width, "beam_width")
vllm/tokenizers/deepseek_v32_encoding.py
+13 -3
@@ -119,7 +119,10 @@ def find_last_user_index(messages: list[dict[str, Any]]) -> int:
119
120
121 def render_message(
122 - index: int, messages: list[dict[str, Any]], thinking_mode: str
122 + index: int,
123 + messages: list[dict[str, Any]],
124 + thinking_mode: str,
125 + last_user_idx: int | None = None,
126 ) -> str:
127 if not (0 <= index < len(messages)):
128 raise ValueError(
@@ -130,7 +133,9 @@ def render_message(
133
134 prompt = ""
135 msg = messages[index]
133 - last_user_idx = find_last_user_index(messages)
136 + last_user_idx = (
137 + find_last_user_index(messages) if last_user_idx is None else last_user_idx
138 + )
139
140 role = msg.get("role")
141 content = msg.get("content")
@@ -294,9 +299,14 @@ def encode_messages(
299 if thinking_mode == "thinking" and drop_thinking:
300 full_messages = drop_thinking_messages(full_messages)
301
302 + last_user_idx = find_last_user_index(full_messages)
303 +
304 for idx in range(len(messages)):
305 prompt += render_message(
299 - idx + len(context), full_messages, thinking_mode=thinking_mode
306 + idx + len(context),
307 + full_messages,
308 + thinking_mode=thinking_mode,
309 + last_user_idx=last_user_idx,
310 )
311
312 return prompt
vllm/tokenizers/deepseek_v4_encoding.py
+15 -2
@@ -199,7 +199,14 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
199 # Message Rendering
200 # ============================================================
201
202 -def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
202 +def render_message(
203 + index: int,
204 + messages: List[Dict[str, Any]],
205 + thinking_mode: str,
206 + drop_thinking: bool = True,
207 + reasoning_effort: Optional[str] = None,
208 + last_user_idx: Optional[int] = None,
209 +) -> str:
210 """
211 Render a single message at the given index into its encoded string form.
212
@@ -213,6 +220,7 @@ def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: st
220 drop_thinking: Whether to drop reasoning content from earlier turns.
221 reasoning_effort: Reasoning effort level, one of "low", "high", "max".
222 None is treated as "low".
223 + last_user_idx: Cached index of the last user/developer message.
224
225 Returns:
226 Encoded string for this message.
@@ -222,7 +230,9 @@ def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: st
230
231 prompt = ""
232 msg = messages[index]
225 - last_user_idx = find_last_user_index(messages)
233 + last_user_idx = (
234 + find_last_user_index(messages) if last_user_idx is None else last_user_idx
235 + )
236
237 role = msg.get("role")
238 content = msg.get("content")
@@ -544,6 +554,8 @@ def encode_messages(
554 num_to_render = len(messages)
555 context_len = len(context)
556
557 + last_user_idx = find_last_user_index(full_messages)
558 +
559 for idx in range(num_to_render):
560 prompt += render_message(
561 idx + context_len,
@@ -551,6 +563,7 @@ def encode_messages(
563 thinking_mode=thinking_mode,
564 drop_thinking=effective_drop_thinking,
565 reasoning_effort=reasoning_effort,
566 + last_user_idx=last_user_idx,
567 )
568
569 return prompt
vllm/v1/worker/gpu/sample/bad_words.py
+13 -11
@@ -3,14 +3,12 @@
3 import numpy as np
4 import torch
5
6 +import vllm.envs as envs
7 from vllm.sampling_params import SamplingParams
8 from vllm.triton_utils import tl, triton
9 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
10 from vllm.v1.worker.gpu.states import RequestState
11
11 -MAX_BAD_WORDS_TOTAL_TOKENS = 1024 # Max total tokens for all bad words per request
12 -MAX_NUM_BAD_WORDS = 128 # Max number of bad words per request
13 -
12
13 class BadWordsState:
14 def __init__(self, req_states: RequestState):
@@ -18,15 +16,17 @@ class BadWordsState:
16 self.max_num_reqs = req_states.max_num_reqs
17 self.device = req_states.device
18
21 - # flattened bad word tokens: [max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS]
19 + max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS
20 + max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
21 + # flattened bad word tokens: [max_num_reqs, VLLM_MAX_BAD_WORDS_TOTAL_TOKENS]
22 self.bad_word_token_ids = StagedWriteTensor(
23 - (self.max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS),
23 + (self.max_num_reqs, max_total_tokens),
24 dtype=torch.int32,
25 device=self.device,
26 )
27 - # cumulative offsets of bad words: [max_num_reqs, MAX_NUM_BAD_WORDS + 1]
27 + # cumulative offsets of bad words: [max_num_reqs, VLLM_MAX_NUM_BAD_WORDS + 1]
28 self.bad_word_offsets = StagedWriteTensor(
29 - (self.max_num_reqs, MAX_NUM_BAD_WORDS + 1),
29 + (self.max_num_reqs, max_num_bad_words + 1),
30 dtype=torch.int32,
31 device=self.device,
32 )
@@ -40,10 +40,11 @@ class BadWordsState:
40 return
41
42 num_bad_words = len(bad_words_token_ids)
43 - if num_bad_words > MAX_NUM_BAD_WORDS:
43 + max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS
44 + if num_bad_words > max_num_bad_words:
45 raise ValueError(
46 f"Too many bad words: {num_bad_words}. "
46 - f"The max number is {MAX_NUM_BAD_WORDS}."
47 + f"The max number is {max_num_bad_words}."
48 )
49
50 # Flatten bad words and compute offsets
@@ -53,10 +54,11 @@ class BadWordsState:
54 flattened_tokens.extend(bad_word)
55 offsets.append(len(flattened_tokens))
56
56 - if len(flattened_tokens) > MAX_BAD_WORDS_TOTAL_TOKENS:
57 + max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS
58 + if len(flattened_tokens) > max_total_tokens:
59 raise ValueError(
60 f"Too many total bad word tokens: {len(flattened_tokens)}. "
59 - f"The max is {MAX_BAD_WORDS_TOTAL_TOKENS}."
61 + f"The max is {max_total_tokens}."
62 )
63
64 # Stage writes