@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2d721ab5d8

[Rust Frontend] Align Rust allowed_token_ids validation with Python (#46348)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: reidliu41 <reid201711@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Reid committed Jun 23, 2026 at 16:32 UTC 2d721ab5d82a2a2505480924c615187bd5b793c1
6 files changed +138 -27
rust/src/server/src/error.rs
+12 -2
@@ -102,7 +102,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
102 vllm_text::Error::PromptTooLong { .. }
103 | vllm_text::Error::EmptyPromptTokenIds { .. }
104 | vllm_text::Error::Logprobs(_)
105 - | vllm_text::Error::OutOfVocab(_)
105 + | vllm_text::Error::TokenIds(_)
106 | vllm_text::Error::InvalidThinkingTokenBudget
107 // An empty tokenized prompt detected later, at request prepare
108 // time, surfaces through the transparent Llm wrapper.
@@ -188,7 +188,7 @@ mod tests {
188
189 #[test]
190 fn out_of_vocab_validation_maps_to_invalid_request() {
191 - let error = vllm_text::Error::OutOfVocab(vllm_text::OutOfVocabError {
191 + let error = vllm_text::Error::TokenIds(vllm_text::TokenIdsError::OutOfVocab {
192 parameter: "logprob_token_ids",
193 token_ids: vec![1000],
194 vocab_size: 1000,
@@ -197,6 +197,16 @@ mod tests {
197 assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
198 }
199
200 + #[test]
201 + fn empty_allowed_token_ids_maps_to_invalid_request() {
202 + let error = vllm_text::Error::TokenIds(vllm_text::TokenIdsError::EmptyAllowedTokenIds);
203 + let api_error = text_submit_error("failed to submit completion request", error);
204 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
205 + let response = api_error.to_error_response();
206 + assert_eq!(response.error.error_type, "invalid_request_error");
207 + assert!(response.error.message.contains("allowed_token_ids"));
208 + }
209 +
210 #[test]
211 fn other_submit_errors_stay_internal() {
212 let error = vllm_text::Error::Tokenizer("backend exploded".to_string());
rust/src/server/src/routes/tests.rs
+72
@@ -2237,6 +2237,42 @@ async fn invalid_request_returns_openai_error() {
2237 assert_eq!(json["error"]["type"], "invalid_request_error");
2238 }
2239
2240 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2241 +#[serial]
2242 +async fn chat_completions_empty_allowed_token_ids_returns_openai_error() {
2243 + let mut app = test_app().await;
2244 + let response = app
2245 + .call(
2246 + Request::builder()
2247 + .method("POST")
2248 + .uri("/v1/chat/completions")
2249 + .header("content-type", "application/json")
2250 + .body(Body::from(
2251 + json!({
2252 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2253 + "stream": false,
2254 + "messages": [{"role": "user", "content": "hello"}],
2255 + "allowed_token_ids": []
2256 + })
2257 + .to_string(),
2258 + ))
2259 + .expect("build request"),
2260 + )
2261 + .await
2262 + .expect("call app");
2263 +
2264 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
2265 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2266 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2267 + assert_eq!(json["error"]["type"], "invalid_request_error");
2268 + assert!(
2269 + json["error"]["message"]
2270 + .as_str()
2271 + .expect("message string")
2272 + .contains("allowed_token_ids should not be empty")
2273 + );
2274 +}
2275 +
2276 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2277 #[serial]
2278 async fn non_stream_chat_returns_json_response() {
@@ -2972,6 +3008,42 @@ async fn completions_invalid_request_returns_openai_error() {
3008 assert_eq!(json["error"]["type"], "invalid_request_error");
3009 }
3010
3011 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3012 +#[serial]
3013 +async fn completions_empty_allowed_token_ids_returns_openai_error() {
3014 + let mut app = test_app().await;
3015 + let response = app
3016 + .call(
3017 + Request::builder()
3018 + .method("POST")
3019 + .uri("/v1/completions")
3020 + .header("content-type", "application/json")
3021 + .body(Body::from(
3022 + json!({
3023 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3024 + "prompt": "hello",
3025 + "stream": false,
3026 + "allowed_token_ids": []
3027 + })
3028 + .to_string(),
3029 + ))
3030 + .expect("build request"),
3031 + )
3032 + .await
3033 + .expect("call app");
3034 +
3035 + assert_eq!(response.status(), StatusCode::BAD_REQUEST);
3036 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3037 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
3038 + assert_eq!(json["error"]["type"], "invalid_request_error");
3039 + assert!(
3040 + json["error"]["message"]
3041 + .as_str()
3042 + .expect("message string")
3043 + .contains("allowed_token_ids should not be empty")
3044 + );
3045 +}
3046 +
3047 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3048 #[serial]
3049 async fn non_stream_completions_return_json_response() {
rust/src/text/src/error.rs
+2 -2
@@ -3,7 +3,7 @@ use vllm_engine_core_client::Error as EngineCoreError;
3 use vllm_llm::Error as LlmError;
4
5 pub use crate::lower::logprobs::LogprobsError;
6 -pub use crate::lower::token_ids::OutOfVocabError;
6 +pub use crate::lower::token_ids::TokenIdsError;
7
8 #[derive(Debug, Error)]
9 pub enum Error {
@@ -19,7 +19,7 @@ pub enum Error {
19 #[error(transparent)]
20 Logprobs(#[from] LogprobsError),
21 #[error(transparent)]
22 - OutOfVocab(#[from] OutOfVocabError),
22 + TokenIds(#[from] TokenIdsError),
23 #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]
24 InvalidThinkingTokenBudget,
25 #[error("text request stream `{request_id}` closed before terminal output")]
rust/src/text/src/lib.rs
+1 -1
@@ -7,7 +7,7 @@
7 use std::mem::take;
8
9 pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend};
10 -pub use error::{Error, LogprobsError, OutOfVocabError, Result};
10 +pub use error::{Error, LogprobsError, Result, TokenIdsError};
11 use futures::Stream;
12 pub use lower::{
13 PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens,
rust/src/text/src/lower.rs
+24 -7
@@ -274,7 +274,7 @@ mod tests {
274 use super::*;
275 use crate::backend::hf::HfTextBackend;
276 use crate::backend::{SamplingHints, TextBackend as _};
277 - use crate::error::{LogprobsError, OutOfVocabError};
277 + use crate::error::{LogprobsError, TokenIdsError};
278 use crate::request::{Prompt, TextRequest};
279
280 /// Stub tokenizer that returns empty token IDs — sufficient for tests that
@@ -555,7 +555,7 @@ mod tests {
555
556 assert!(matches!(
557 error,
558 - Error::OutOfVocab(OutOfVocabError {
558 + Error::TokenIds(TokenIdsError::OutOfVocab {
559 parameter: "prompt",
560 token_ids,
561 vocab_size: 2000,
@@ -857,7 +857,7 @@ mod tests {
857
858 assert!(matches!(
859 error,
860 - Error::OutOfVocab(OutOfVocabError {
860 + Error::TokenIds(TokenIdsError::OutOfVocab {
861 parameter: "logprob_token_ids",
862 token_ids,
863 vocab_size: 1000,
@@ -878,7 +878,7 @@ mod tests {
878
879 assert!(matches!(
880 error,
881 - Error::OutOfVocab(OutOfVocabError {
881 + Error::TokenIds(TokenIdsError::OutOfVocab {
882 parameter: "stop_token_ids",
883 token_ids,
884 vocab_size: 1000,
@@ -899,7 +899,7 @@ mod tests {
899
900 assert!(matches!(
901 error,
902 - Error::OutOfVocab(OutOfVocabError {
902 + Error::TokenIds(TokenIdsError::OutOfVocab {
903 parameter: "allowed_token_ids",
904 token_ids,
905 vocab_size: 2000,
@@ -907,6 +907,23 @@ mod tests {
907 ));
908 }
909
910 + #[test]
911 + fn lower_sampling_params_rejects_empty_allowed_token_ids() {
912 + let error = lower_sampling_params_with_limits(
913 + SamplingParams {
914 + allowed_token_ids: Some(vec![]),
915 + ..Default::default()
916 + },
917 + sample_sampling_limits(),
918 + )
919 + .unwrap_err();
920 +
921 + assert!(matches!(
922 + error,
923 + Error::TokenIds(TokenIdsError::EmptyAllowedTokenIds)
924 + ));
925 + }
926 +
927 #[test]
928 fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
929 let tokenizer = FixedTokenizer {
@@ -926,7 +943,7 @@ mod tests {
943
944 assert!(matches!(
945 error,
929 - Error::OutOfVocab(OutOfVocabError {
946 + Error::TokenIds(TokenIdsError::OutOfVocab {
947 parameter: "bad_words",
948 token_ids,
949 vocab_size: 2000,
@@ -947,7 +964,7 @@ mod tests {
964
965 assert!(matches!(
966 error,
950 - Error::OutOfVocab(OutOfVocabError {
967 + Error::TokenIds(TokenIdsError::OutOfVocab {
968 parameter: "logit_bias",
969 token_ids,
970 vocab_size: 1000,
rust/src/text/src/lower/token_ids.rs
+27 -15
@@ -6,21 +6,25 @@ use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
6 use crate::SamplingLimits;
7
8 #[derive(Debug, Error)]
9 -#[error(
10 - "token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \
11 - Vocabulary size: {vocab_size}"
12 -)]
13 -pub struct OutOfVocabError {
14 - pub parameter: &'static str,
15 - pub token_ids: Vec<u32>,
16 - pub vocab_size: usize,
9 +pub enum TokenIdsError {
10 + #[error("allowed_token_ids should not be empty")]
11 + EmptyAllowedTokenIds,
12 + #[error(
13 + "token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \
14 + Vocabulary size: {vocab_size}"
15 + )]
16 + OutOfVocab {
17 + parameter: &'static str,
18 + token_ids: Vec<u32>,
19 + vocab_size: usize,
20 + },
21 }
22
23 fn validate_param(
24 parameter: &'static str,
25 token_ids: impl IntoIterator<Item = u32>,
26 vocab_size: usize,
23 -) -> Result<(), OutOfVocabError> {
27 +) -> Result<(), TokenIdsError> {
28 let invalid_token_ids: Vec<_> = token_ids
29 .into_iter()
30 .filter(|&token_id| token_id as usize >= vocab_size)
@@ -29,7 +33,7 @@ fn validate_param(
33 return Ok(());
34 }
35
32 - Err(OutOfVocabError {
36 + Err(TokenIdsError::OutOfVocab {
37 parameter,
38 token_ids: invalid_token_ids,
39 vocab_size,
@@ -41,7 +45,7 @@ fn validate_param(
45 pub(crate) fn validate_prompt_token_ids(
46 prompt_token_ids: &[u32],
47 limits: &SamplingLimits,
44 -) -> Result<(), OutOfVocabError> {
48 +) -> Result<(), TokenIdsError> {
49 validate_param(
50 "prompt",
51 prompt_token_ids.iter().copied(),
@@ -54,7 +58,7 @@ pub(crate) fn validate_prompt_token_ids(
58 pub(crate) fn validate_vocab_range(
59 params: &EngineCoreSamplingParams,
60 limits: &SamplingLimits,
57 -) -> Result<(), OutOfVocabError> {
61 +) -> Result<(), TokenIdsError> {
62 validate_param(
63 "stop_token_ids",
64 params.stop_token_ids.iter().copied(),
@@ -62,6 +66,9 @@ pub(crate) fn validate_vocab_range(
66 )?;
67
68 if let Some(token_ids) = params.allowed_token_ids.as_deref() {
69 + if token_ids.is_empty() {
70 + return Err(TokenIdsError::EmptyAllowedTokenIds);
71 + }
72 validate_param(
73 "allowed_token_ids",
74 token_ids.iter().copied(),
@@ -104,8 +111,13 @@ mod tests {
111 fn validate_vocab_range_rejects_out_of_vocab_ids() {
112 let error = validate_param("logprob_token_ids", [5_u32, 1000, 1001], 1000).unwrap_err();
113
107 - assert_eq!(error.parameter, "logprob_token_ids");
108 - assert_eq!(error.token_ids, vec![1000, 1001]);
109 - assert_eq!(error.vocab_size, 1000);
114 + assert!(matches!(
115 + error,
116 + TokenIdsError::OutOfVocab {
117 + parameter: "logprob_token_ids",
118 + token_ids,
119 + vocab_size: 1000,
120 + } if token_ids == vec![1000, 1001]
121 + ));
122 }
123 }