@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 80abe0de7d

[Rust Frontend] Support thinking_token_budget for chat and completions (#46137)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Chao-Ju Chen committed Jun 22, 2026 at 16:00 UTC 80abe0de7d20523e465597d823a40ab4a29df20a

13 files changed +178 -7

rust/src/engine-core-client/src/protocol/mod.rs

@@ -277,6 +277,12 @@ pub struct EngineCoreSamplingParams {

277 pub max_tokens: u32,

278 /// Minimum number of tokens to generate before EOS or stop-token handling.

279 pub min_tokens: u32,

280 + /// Maximum number of reasoning ("thinking") tokens to emit before the

281 + /// reasoning section is force-closed. `None` means unlimited; the

282 + /// user-facing `-1` sentinel is normalized to `None` by the frontend before

283 + /// reaching this DTO, so only non-negative values are sent. Enforced

284 + /// engine-side (and only when a reasoning parser is configured).

285 + pub thinking_token_budget: Option<u64>,

286 /// Number of log probabilities to return per generated token.

287 ///

288 /// `None` disables sample logprobs. `-1` requests the full vocabulary.

@@ -345,6 +351,7 @@ impl EngineCoreSamplingParams {

351 seed: None,

352 max_tokens: 65536,

353 min_tokens: 0,

354 + thinking_token_budget: None,

355 logprobs: None,

356 prompt_logprobs: None,

357 min_p: 0.0,

rust/src/engine-core-client/src/tests/client.rs

@@ -150,6 +150,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {

150 top_k: 8,

151 max_tokens: 32,

152 min_tokens: 1,

153 + thinking_token_budget: Some(256),

154 stop_token_ids: vec![151643],

155 eos_token_id: Some(151645),

156 all_stop_token_ids: BTreeSet::from([151643, 151645]),

@@ -2502,6 +2503,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {

2503 seed: None,

2504 max_tokens: 16,

2505 min_tokens: 0,

2506 + thinking_token_budget: None,

2507 logprobs: None,

2508 prompt_logprobs: None,

2509 min_p: 0.0,

rust/src/engine-core-client/src/tests/python_compat.py

@@ -39,6 +39,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):

39 seed: int | None = None

40 max_tokens: int = 16

41 min_tokens: int = 0

42 + thinking_token_budget: int | None = None

43 min_p: float = 0.0

44 frequency_penalty: float = 0.0

45 presence_penalty: float = 0.0

@@ -122,6 +123,7 @@ request = EngineCoreRequest(

123 seed=None,

124 max_tokens=32,

125 min_tokens=1,

126 + thinking_token_budget=256,

127 min_p=0.0,

128 frequency_penalty=0.0,

129 presence_penalty=0.0,

rust/src/server/src/error.rs

+13

@@ -103,6 +103,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {

103 | vllm_text::Error::EmptyPromptTokenIds { .. }

104 | vllm_text::Error::Logprobs(_)

105 | vllm_text::Error::OutOfVocab(_)

106 + | vllm_text::Error::InvalidThinkingTokenBudget

107 // An empty tokenized prompt detected later, at request prepare

108 // time, surfaces through the transparent Llm wrapper.

109 | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })

@@ -127,6 +128,18 @@ mod tests {

128 assert!(response.error.message.contains("9000"));

129 }

130

131 + #[test]

132 + fn invalid_thinking_token_budget_maps_to_invalid_request() {

133 + let api_error = text_submit_error(

134 + "failed to submit completion request",

135 + vllm_text::Error::InvalidThinkingTokenBudget,

136 + );

137 + assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);

138 + let response = api_error.to_error_response();

139 + assert_eq!(response.error.error_type, "invalid_request_error");

140 + assert!(response.error.message.contains("thinking_token_budget"));

141 + }

142 +

143 #[test]

144 fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {

145 let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {

rust/src/server/src/routes/inference/generate/convert.rs

+27

@@ -150,6 +150,33 @@ mod tests {

150 );

151 }

152

153 + #[test]

154 + fn prepare_generate_request_forwards_thinking_token_budget() {

155 + let request: GenerateRequest = serde_json::from_value(json!({

156 + "model": "Qwen/Qwen1.5-0.5B-Chat",

157 + "token_ids": [11, 22, 33],

158 + "sampling_params": {

159 + "thinking_token_budget": 64

160 + }

161 + }))

162 + .expect("parse request");

163 +

164 + let prepared = prepare_generate_request(

165 + request,

166 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

167 + ResolvedRequestContext::default(),

168 + )

169 + .expect("prepare");

170 +

171 + // The raw inference route shares `vllm_text::SamplingParams`, so the

172 + // field is carried through to lowering exactly like the OpenAI routes

173 + // (normalization/validation then happens in `lower_sampling_params`).

174 + assert_eq!(

175 + prepared.text_request.sampling_params.thinking_token_budget,

176 + Some(64)

177 + );

178 + }

179 +

180 #[test]

181 fn prepare_generate_request_gates_continuous_usage_on_include_usage() {

182 let request: GenerateRequest = serde_json::from_value(json!({

rust/src/server/src/routes/openai/chat_completions/convert.rs

+26

@@ -115,6 +115,7 @@ pub(super) fn prepare_chat_request(

115 seed: request.seed,

116 max_tokens: request.max_completion_tokens,

117 min_tokens: request.min_tokens,

118 + thinking_token_budget: request.thinking_token_budget,

119 logprobs: request.logprobs.then_some(top_logprobs),

120 prompt_logprobs,

121 min_p: request.min_p,

@@ -613,6 +614,31 @@ mod tests {

614 assert_eq!(prepared.chat_request.sampling_params, expected);

615 }

616

617 + #[test]

618 + fn prepare_chat_request_passes_through_thinking_token_budget() {

619 + let prepare = |budget: Option<i64>| {

620 + prepare_chat_request(

621 + ChatCompletionRequest {

622 + thinking_token_budget: budget,

623 + ..base_request()

624 + },

625 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

626 + ResolvedRequestContext::default(),

627 + )

628 + .expect("request is valid")

629 + .chat_request

630 + .sampling_params

631 + .thinking_token_budget

632 + };

633 +

634 + // The convert layer forwards the raw value verbatim (including the `-1`

635 + // "unlimited" sentinel); normalization/validation happens during

636 + // lowering (see `vllm_text::lower`).

637 + assert_eq!(prepare(Some(64)), Some(64));

638 + assert_eq!(prepare(Some(-1)), Some(-1));

639 + assert_eq!(prepare(None), None);

640 + }

641 +

642 #[test]

643 fn prepare_chat_request_accepts_developer_messages() {

644 let request = ChatCompletionRequest {

rust/src/server/src/routes/openai/chat_completions/types.rs

+4 -2

@@ -165,8 +165,10 @@ pub struct ChatCompletionRequest {

165 pub bad_words: Option<Vec<String>>,

166

167 // -------- Extra vLLM Parameters --------

168 - /// Token budget for reasoning/thinking

169 - pub thinking_token_budget: Option<u32>,

168 + /// Token budget for reasoning/thinking. Accepts a non-negative integer, or

169 + /// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`

170 + /// to "no budget").

171 + pub thinking_token_budget: Option<i64>,

172

173 /// Whether to include reasoning content in the response

174 #[serde(default = "default_true")]

rust/src/server/src/routes/openai/chat_completions/validate.rs

-5

@@ -108,11 +108,6 @@ pub(super) fn validate_request_compat(

108 "truncate_prompt_tokens",

109 "truncate_prompt_tokens is not supported.",

110 )?;

111 - reject_non_default(

112 - request.thinking_token_budget.as_ref(),

113 - "thinking_token_budget",

114 - "thinking_token_budget is not supported.",

115 - )?;

111 reject_non_default(

112 request.media_io_kwargs.as_ref(),

113 "media_io_kwargs",

rust/src/server/src/routes/openai/completions/convert.rs

+29

@@ -108,6 +108,7 @@ pub(super) fn prepare_completion_request(

108 seed: request.seed,

109 max_tokens,

110 min_tokens: request.min_tokens,

111 + thinking_token_budget: request.thinking_token_budget,

112 logprobs,

113 prompt_logprobs,

114 min_p: request.min_p,

@@ -266,6 +267,34 @@ mod tests {

267 assert!(!prepared.text_request.decode_options.skip_special_tokens);

268 }

269

270 + #[test]

271 + fn prepare_completion_request_passes_through_thinking_token_budget() {

272 + let prepare = |budget: serde_json::Value| {

273 + let request: CompletionRequest = serde_json::from_value(json!({

274 + "model": "Qwen/Qwen1.5-0.5B-Chat",

275 + "prompt": "hello",

276 + "thinking_token_budget": budget,

277 + }))

278 + .expect("parse request");

279 + prepare_completion_request(

280 + request,

281 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),

282 + ResolvedRequestContext::default(),

283 + )

284 + .expect("prepare")

285 + .text_request

286 + .sampling_params

287 + .thinking_token_budget

288 + };

289 +

290 + // The convert layer forwards the raw value verbatim (including the `-1`

291 + // "unlimited" sentinel); normalization/validation happens during

292 + // lowering (see `vllm_text::lower`).

293 + assert_eq!(prepare(json!(64)), Some(64));

294 + assert_eq!(prepare(json!(-1)), Some(-1));

295 + assert_eq!(prepare(json!(null)), None);

296 + }

297 +

298 #[test]

299 fn prepare_completion_request_maps_stream_usage_and_token_format_options() {

300 let request: CompletionRequest = serde_json::from_value(json!({

rust/src/server/src/routes/openai/completions/types.rs

@@ -146,6 +146,11 @@ pub struct CompletionRequest {

146 /// Additional kwargs for structured outputs

147 pub structured_outputs: Option<Value>,

148

149 + /// Token budget for reasoning/thinking. Accepts a non-negative integer, or

150 + /// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`

151 + /// to "no budget").

152 + pub thinking_token_budget: Option<i64>,

153 +

154 /// Request scheduling priority (lower means earlier; default 0)

155 pub priority: Option<i32>,

156

rust/src/text/src/error.rs

@@ -20,6 +20,8 @@ pub enum Error {

20 Logprobs(#[from] LogprobsError),

21 #[error(transparent)]

22 OutOfVocab(#[from] OutOfVocabError),

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")]

26 StreamClosedBeforeTerminalOutput { request_id: String },

27 #[error(transparent)]

rust/src/text/src/lower.rs

+54

@@ -87,6 +87,7 @@ pub fn lower_sampling_params(

87 seed,

88 max_tokens,

89 min_tokens,

90 + thinking_token_budget,

91 logprobs,

92 prompt_logprobs,

93 min_p,

@@ -128,6 +129,7 @@ pub fn lower_sampling_params(

129 prompt_len,

130 )?;

131 let min_tokens = min_tokens.unwrap_or(0);

132 + let thinking_token_budget = normalize_thinking_token_budget(thinking_token_budget)?;

133 let frequency_penalty = frequency_penalty.unwrap_or(0.0);

134 let presence_penalty = presence_penalty.unwrap_or(0.0);

135

@@ -149,6 +151,7 @@ pub fn lower_sampling_params(

151 seed,

152 max_tokens,

153 min_tokens,

154 + thinking_token_budget,

155 logprobs,

156 prompt_logprobs,

157 min_p,

@@ -170,6 +173,21 @@ pub fn lower_sampling_params(

173 Ok(params)

174 }

175

176 +/// Normalize the user-facing `thinking_token_budget` into the engine value.

177 +///

178 +/// Mirrors Python's `validate_thinking_token_budget`

179 +/// (<https://github.com/vllm-project/vllm/blob/ecf9d83520eb217401b47d8a5451a27c5231b8c2/vllm/sampling_params.py#L35-L55>):

180 +/// `None` and the `-1` "unlimited" sentinel both map to `None`; any other

181 +/// negative value is rejected; non-negative values pass through unchanged. Like

182 +/// Python's `int`, no upper bound is imposed.

183 +fn normalize_thinking_token_budget(value: Option<i64>) -> Result<Option<u64>> {

184 + match value {

185 + None | Some(-1) => Ok(None),

186 + Some(budget) if budget >= 0 => Ok(Some(budget as u64)),

187 + Some(_) => Err(Error::InvalidThinkingTokenBudget),

188 + }

189 +}

190 +

191 /// Convert bad-word strings into token-ID sequences, following the Python vLLM

192 /// logic in `SamplingParams.update_from_tokenizer()`.

193 ///

@@ -366,6 +384,36 @@ mod tests {

384 )

385 }

386

387 + #[test]

388 + fn lower_sampling_params_normalizes_thinking_token_budget() {

389 + let lower = |budget: Option<i64>| {

390 + lower_sampling_params_with_limits(

391 + SamplingParams {

392 + thinking_token_budget: budget,

393 + ..SamplingParams::default()

394 + },

395 + sample_sampling_limits(),

396 + )

397 + };

398 +

399 + // Non-negative budgets (including 0) pass through unchanged.

400 + assert_eq!(lower(Some(256)).unwrap().thinking_token_budget, Some(256));

401 + assert_eq!(lower(Some(0)).unwrap().thinking_token_budget, Some(0));

402 + // `None` and the `-1` "unlimited" sentinel both disable the budget.

403 + assert_eq!(lower(None).unwrap().thinking_token_budget, None);

404 + assert_eq!(lower(Some(-1)).unwrap().thinking_token_budget, None);

405 + // No upper bound is imposed, matching Python's `int`.

406 + assert_eq!(

407 + lower(Some(i64::from(u32::MAX) + 1)).unwrap().thinking_token_budget,

408 + Some(u64::from(u32::MAX) + 1)

409 + );

410 + // Other negatives are rejected.

411 + assert!(matches!(

412 + lower(Some(-2)),

413 + Err(Error::InvalidThinkingTokenBudget)

414 + ));

415 + }

416 +

417 #[test]

418 fn lower_text_request_applies_python_style_eos_hints() {

419 let prepared = lower_text_request(

@@ -386,6 +434,7 @@ mod tests {

434 seed: None,

435 max_tokens: 999997,

436 min_tokens: 0,

437 + thinking_token_budget: None,

438 logprobs: None,

439 prompt_logprobs: None,

440 min_p: 0.0,

@@ -437,6 +486,7 @@ mod tests {

486 seed: None,

487 max_tokens: 999997,

488 min_tokens: 0,

489 + thinking_token_budget: None,

490 logprobs: None,

491 prompt_logprobs: None,

492 min_p: 0.0,

@@ -567,6 +617,7 @@ mod tests {

617 seed: None,

618 max_tokens: 40957,

619 min_tokens: 0,

620 + thinking_token_budget: None,

621 logprobs: None,

622 prompt_logprobs: None,

623 min_p: 0.0,

@@ -628,6 +679,7 @@ mod tests {

679 seed: None,

680 max_tokens: 999997,

681 min_tokens: 0,

682 + thinking_token_budget: None,

683 logprobs: None,

684 prompt_logprobs: None,

685 min_p: 0.0,

@@ -697,6 +749,7 @@ mod tests {

749 seed: None,

750 max_tokens: 32,

751 min_tokens: 2,

752 + thinking_token_budget: None,

753 logprobs: None,

754 prompt_logprobs: None,

755 min_p: 0.1,

@@ -929,6 +982,7 @@ mod tests {

982 seed: None,

983 max_tokens: 128,

984 min_tokens: 0,

985 + thinking_token_budget: None,

986 logprobs: None,

987 prompt_logprobs: None,

988 min_p: 0.1,

rust/src/text/src/request.rs

@@ -56,6 +56,12 @@ pub struct SamplingParams {

56 pub max_tokens: Option<u32>,

57 /// Minimum number of tokens to generate before EOS or stop-token handling.

58 pub min_tokens: Option<u32>,

59 + /// Maximum number of reasoning ("thinking") tokens to emit before the

60 + /// reasoning section is force-closed. `None` or the user-facing `-1`

61 + /// "unlimited" sentinel both disable the budget. The raw value is carried

62 + /// here; `-1` is normalized to `None` (and other negatives rejected) during

63 + /// lowering (see `lower_sampling_params`).

64 + pub thinking_token_budget: Option<i64>,

65 /// Number of log probabilities to return per generated token.

66 ///

67 /// `None` disables sample logprobs. `-1` requests the full vocabulary.

@@ -116,6 +122,7 @@ impl Default for SamplingParams {

122 seed: None,

123 max_tokens: None,

124 min_tokens: None,

125 + thinking_token_budget: None,

126 logprobs: None,

127 prompt_logprobs: None,

128 min_p: None,