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

[Rust Frontend] Support echo for token-ID completion prompts (#46219)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: reidliu41 <reid201711@gmail.com>

Reid committed Jun 23, 2026 at 16:04 UTC accaa434f36b37a35b3e68eede167415ecc83c51

4 files changed +275 -23

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

+7 -1

@@ -51,7 +51,13 @@ pub async fn completions(

51 let request_context = resolve_request_context(&headers, body.request_id.as_deref());

52 let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;

54 - let prepared = match prepare_completion_request(body, &lora_resolution, request_context) {

54 + let tokenizer = state.chat.text().tokenizer();

55 + let prepared = match prepare_completion_request(

56 + body,

57 + &lora_resolution,

58 + request_context,

59 + tokenizer.as_ref(),

60 + ) {

61 Ok(prepared) => prepared,

62 Err(error) => return error.into_response(),

63 };

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

+87 -13

@@ -1,4 +1,6 @@

1 -use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest};

1 +use thiserror_ext::AsReport as _;

2 +use vllm_text::tokenizer::Tokenizer;

3 +use vllm_text::{Prompt, SamplingParams, TextDecodeOptions, TextRequest};

5 use super::types::CompletionRequest;

6 use crate::error::ApiError;

@@ -29,8 +31,7 @@ pub(super) struct ResponseOptions {

31 pub include_continuous_usage: bool,

32 /// Whether the caller requested prompt-only echo via `max_tokens=0`.

33 pub prompt_only: bool,

32 - /// Original text prompt that should be echoed back northbound when

33 - /// `echo=true`.

34 + /// Prompt text that should be echoed back northbound when `echo=true`.

35 pub echo: Option<String>,

36 /// Whether the caller requested output logprobs on completion choices.

37 pub requested_logprobs: Option<u32>,

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

52 request: CompletionRequest,

53 lora_resolution: &LoraModelResolution,

54 ctx: ResolvedRequestContext,

55 + tokenizer: &dyn Tokenizer,

56 ) -> Result<PreparedRequest, ApiError> {

57 validate::validate_request_compat(&request, &lora_resolution.model_names)?;

@@ -92,7 +94,7 @@ pub(super) fn prepare_completion_request(

94 } else {

95 request.max_tokens

96 };

95 - let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();

97 + let echo = completion_echo_text(&request, tokenizer)?;

99 let structured_outputs =

100 convert_from_response_format_value(&request.response_format, &request.structured_outputs)?;

@@ -159,11 +161,39 @@ pub(super) fn prepare_completion_request(

161 })

162 }

163

164 +fn completion_echo_text(

165 + request: &CompletionRequest,

166 + tokenizer: &dyn Tokenizer,

167 +) -> Result<Option<String>, ApiError> {

168 + if !request.echo {

169 + return Ok(None);

170 + }

171 +

172 + match &request.prompt {

173 + Prompt::Text(prompt) => Ok(Some(prompt.clone())),

174 + Prompt::TokenIds(token_ids) if request.return_token_ids.unwrap_or(false) => {

175 + Ok(Some(String::new()))

176 + }

177 + Prompt::TokenIds(token_ids) => {

178 + tokenizer.decode(token_ids, false).map(Some).map_err(|error| {

179 + ApiError::invalid_request(

180 + format!(

181 + "Failed to decode token-ID prompt for echo: {}",

182 + error.to_report_string()

183 + ),

184 + Some("prompt"),

185 + )

186 + })

187 + }

188 + }

189 +}

190 +

191 #[cfg(test)]

192 mod tests {

193 use axum::http::HeaderMap;

194 use serde_json::json;

195 use vllm_text::Prompt;

196 + use vllm_text::tokenizer::Tokenizer;

197

198 use super::prepare_completion_request;

199 use crate::lora::LoraModelResolution;

@@ -181,6 +211,34 @@ mod tests {

211 }

212 }

213

214 + #[derive(Debug)]

215 + struct TestTokenizer;

216 +

217 + impl Tokenizer for TestTokenizer {

218 + fn encode(

219 + &self,

220 + text: &str,

221 + _add_special_tokens: bool,

222 + ) -> vllm_text::tokenizer::Result<Vec<u32>> {

223 + Ok(text.bytes().map(u32::from).collect())

224 + }

225 +

226 + fn decode(

227 + &self,

228 + token_ids: &[u32],

229 + _skip_special_tokens: bool,

230 + ) -> vllm_text::tokenizer::Result<String> {

231 + Ok(

232 + String::from_utf8_lossy(&token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>())

233 + .into_owned(),

234 + )

235 + }

236 +

237 + fn token_to_id(&self, _token: &str) -> Option<u32> {

238 + None

239 + }

240 + }

241 +

242 fn base_request_json() -> serde_json::Value {

243 json!({

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

@@ -238,6 +296,7 @@ mod tests {

296 request,

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

298 ResolvedRequestContext::default(),

299 + &TestTokenizer,

300 )

301 .expect("prepare");

302

@@ -280,6 +339,7 @@ mod tests {

339 request,

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

341 ResolvedRequestContext::default(),

342 + &TestTokenizer,

343 )

344 .expect("prepare")

345 .text_request

@@ -313,6 +373,7 @@ mod tests {

373 request,

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

375 ResolvedRequestContext::default(),

376 + &TestTokenizer,

377 )

378 .expect("prepare");

379

@@ -337,6 +398,7 @@ mod tests {

398 request,

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

400 ResolvedRequestContext::default(),

401 + &TestTokenizer,

402 )

403 .expect("prepare");

404

@@ -359,6 +421,7 @@ mod tests {

421 request,

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

423 ResolvedRequestContext::default(),

424 + &TestTokenizer,

425 )

426 .expect("prepare");

427

@@ -382,6 +445,7 @@ mod tests {

445 request,

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

447 ResolvedRequestContext::default(),

448 + &TestTokenizer,

449 )

450 .expect("prepare");

451

@@ -406,6 +470,7 @@ mod tests {

470 request,

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

472 ResolvedRequestContext::default(),

473 + &TestTokenizer,

474 )

475 .expect("prepare");

476

@@ -432,6 +497,7 @@ mod tests {

497 request,

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

499 ResolvedRequestContext::default(),

500 + &TestTokenizer,

501 )

502 .expect("prepare");

503

@@ -443,22 +509,27 @@ mod tests {

509 }

510

511 #[test]

446 - fn prepare_completion_request_rejects_token_id_prompt_echo() {

512 + fn prepare_completion_request_decodes_token_id_prompt_echo() {

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

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

449 - "prompt": [11, 22, 33],

515 + "prompt": [104, 101, 108, 108, 111],

516 "stream": true,

517 "echo": true

518 }))

519 .expect("parse request");

520

455 - assert!(

456 - prepare_completion_request(

457 - request,

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

459 - ResolvedRequestContext::default(),

460 - )

461 - .is_err()

521 + let prepared = prepare_completion_request(

522 + request,

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

524 + ResolvedRequestContext::default(),

525 + &TestTokenizer,

526 + )

527 + .expect("prepare");

528 +

529 + assert_eq!(prepared.options.echo, Some("hello".to_string()));

530 + assert_eq!(

531 + prepared.text_request.prompt,

532 + Prompt::TokenIds(vec![104, 101, 108, 108, 111])

533 );

534 }

535

@@ -477,6 +548,7 @@ mod tests {

548 request,

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

550 ResolvedRequestContext::default(),

551 + &TestTokenizer,

552 )

553 .expect("prepare");

554 assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));

@@ -501,6 +573,7 @@ mod tests {

573 request,

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

575 request_context(&headers, None),

576 + &TestTokenizer,

577 )

578 .expect("prepare");

579 assert_eq!(prepared.text_request.data_parallel_rank, Some(3));

@@ -519,6 +592,7 @@ mod tests {

592 request,

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

594 ResolvedRequestContext::default(),

595 + &TestTokenizer,

596 )

597 .expect("prepare");

598 assert_eq!(prepared.text_request.data_parallel_rank, None);

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

+15 -9

@@ -1,5 +1,3 @@

1 -use vllm_text::Prompt;

2 -

1 use super::types::CompletionRequest;

2 use crate::error::{ApiError, bail_invalid_request};

@@ -33,13 +31,6 @@ pub(super) fn validate_request_compat(

31 );

32 }

36 - if request.echo && matches!(request.prompt, Prompt::TokenIds(_)) {

37 - bail_invalid_request!(

38 - param = "echo",

39 - "echo is not supported with token-ID prompts."

40 - );

41 - }

42 -

34 if request.suffix.is_some() {

35 bail_invalid_request!(param = "suffix", "suffix is not supported.");

36 }

@@ -183,6 +174,21 @@ mod tests {

174 );

175 }

176

177 + #[test]

178 + fn validate_request_compat_accepts_token_id_prompt_echo() {

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

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

181 + "prompt": [104, 101, 108, 108, 111],

182 + "stream": true,

183 + "echo": true,

184 + }))

185 + .expect("parse request");

186 +

187 + assert!(

188 + validate_request_compat(&request, &served_names(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok()

189 + );

190 + }

191 +

192 #[test]

193 fn validate_request_compat_rejects_prompt_only_without_echo() {

194 let request = CompletionRequest {

rust/src/server/src/routes/tests.rs

+166

@@ -3052,6 +3052,107 @@ async fn non_stream_completions_echo_prepends_prompt_text() {

3052 assert_eq!(json["usage"]["completion_tokens"], 3);

3053 }

3054

3055 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

3056 +#[serial]

3057 +async fn non_stream_completions_echo_decodes_token_id_prompt_text() {

3058 + let prompt_token_ids = bytes_to_token_ids(b"hello");

3059 + let expected_prompt_token_ids = prompt_token_ids.clone();

3060 + let (app, engine_task) = test_app_with_backend_and_engine_request_check(

3061 + Arc::new(FakeChatBackend::new()),

3062 + move |request| {

3063 + assert_eq!(

3064 + request.prompt_token_ids.as_deref(),

3065 + Some(expected_prompt_token_ids.as_slice())

3066 + );

3067 + },

3068 + )

3069 + .await;

3070 + let response = app

3071 + .clone()

3072 + .call(

3073 + Request::builder()

3074 + .method("POST")

3075 + .uri("/v1/completions")

3076 + .header("content-type", "application/json")

3077 + .body(Body::from(

3078 + json!({

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

3080 + "prompt": prompt_token_ids,

3081 + "echo": true,

3082 + "stream": false

3083 + })

3084 + .to_string(),

3085 + ))

3086 + .expect("build request"),

3087 + )

3088 + .await

3089 + .expect("call app");

3090 +

3091 + assert_eq!(response.status(), StatusCode::OK);

3092 +

3093 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

3094 + engine_task.await.expect("mock engine task");

3095 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

3096 +

3097 + assert_eq!(json["choices"][0]["text"], "hellohi");

3098 + assert_eq!(json["usage"]["prompt_tokens"], 5);

3099 + assert_eq!(json["usage"]["completion_tokens"], 3);

3100 +}

3101 +

3102 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

3103 +#[serial]

3104 +async fn non_stream_completions_token_id_echo_return_token_ids_keeps_prompt_ids_separate() {

3105 + let prompt_token_ids = bytes_to_token_ids(b"hello");

3106 + let expected_prompt_token_ids = prompt_token_ids.clone();

3107 + let (app, engine_task) = test_app_with_backend_and_engine_request_check(

3108 + Arc::new(FakeChatBackend::new()),

3109 + move |request| {

3110 + assert_eq!(

3111 + request.prompt_token_ids.as_deref(),

3112 + Some(expected_prompt_token_ids.as_slice())

3113 + );

3114 + },

3115 + )

3116 + .await;

3117 + let response = app

3118 + .clone()

3119 + .call(

3120 + Request::builder()

3121 + .method("POST")

3122 + .uri("/v1/completions")

3123 + .header("content-type", "application/json")

3124 + .body(Body::from(

3125 + json!({

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

3127 + "prompt": prompt_token_ids,

3128 + "echo": true,

3129 + "return_token_ids": true,

3130 + "stream": false

3131 + })

3132 + .to_string(),

3133 + ))

3134 + .expect("build request"),

3135 + )

3136 + .await

3137 + .expect("call app");

3138 +

3139 + assert_eq!(response.status(), StatusCode::OK);

3140 +

3141 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

3142 + engine_task.await.expect("mock engine task");

3143 + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");

3144 +

3145 + assert_eq!(json["choices"][0]["text"], "hi");

3146 + assert_eq!(

3147 + json["choices"][0]["prompt_token_ids"],

3148 + json!(bytes_to_token_ids(b"hello"))

3149 + );

3150 + assert_eq!(

3151 + json["choices"][0]["token_ids"],

3152 + json!(bytes_to_token_ids(b"hi!"))

3153 + );

3154 +}

3155 +

3156 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

3157 #[serial]

3158 async fn non_stream_completions_include_logprobs() {

@@ -4112,6 +4213,71 @@ async fn completions_echo_stream_emits_separate_prompt_chunk() {

4213 assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);

4214 }

4215

4216 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

4217 +#[serial]

4218 +async fn completions_echo_stream_decodes_token_id_prompt_chunk() {

4219 + let prompt_token_ids = bytes_to_token_ids(b"hello");

4220 + let expected_prompt_token_ids = prompt_token_ids.clone();

4221 + let (app, engine_task) = test_app_with_backend_and_engine_request_check(

4222 + Arc::new(FakeChatBackend::new()),

4223 + move |request| {

4224 + assert_eq!(

4225 + request.prompt_token_ids.as_deref(),

4226 + Some(expected_prompt_token_ids.as_slice())

4227 + );

4228 + },

4229 + )

4230 + .await;

4231 + let response = app

4232 + .clone()

4233 + .call(

4234 + Request::builder()

4235 + .method("POST")

4236 + .uri("/v1/completions")

4237 + .header("content-type", "application/json")

4238 + .body(Body::from(

4239 + json!({

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

4241 + "prompt": prompt_token_ids,

4242 + "echo": true,

4243 + "stream": true,

4244 + "stream_options": {"include_usage": true}

4245 + })

4246 + .to_string(),

4247 + ))

4248 + .expect("build request"),

4249 + )

4250 + .await

4251 + .expect("call app");

4252 +

4253 + assert_eq!(response.status(), StatusCode::OK);

4254 +

4255 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

4256 + engine_task.await.expect("mock engine task");

4257 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");

4258 + let payloads = sse_data_payloads(&text);

4259 + let hello_index = payloads

4260 + .iter()

4261 + .position(|payload| payload.contains("\"text\":\"hello\""))

4262 + .expect("prompt echo chunk");

4263 + let h_index = payloads

4264 + .iter()

4265 + .position(|payload| payload.contains("\"text\":\"h\""))

4266 + .expect("first generation chunk");

4267 +

4268 + assert!(hello_index < h_index, "{text}");

4269 +

4270 + let usage_chunk: serde_json::Value = serde_json::from_str(

4271 + payloads

4272 + .iter()

4273 + .find(|payload| payload.contains("\"usage\":"))

4274 + .expect("usage chunk"),

4275 + )

4276 + .expect("usage chunk json");

4277 + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 5);

4278 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);

4279 +}

4280 +

4281 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

4282 #[serial]

4283 async fn reasoning_blocks_are_mapped_to_reasoning_sse_chunks() {