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

[Rust Frontend] Extend renderer/parser roundtrip tests to support token ids (#47110)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 30, 2026 at 21:27 UTC bdbd7278b6b169c2b188bc1e3f9541468a9e6f4b
2 files changed +160 -79
rust/src/chat/src/output/harmony/tests.rs
-14
@@ -1,14 +1,8 @@
1 -//! Harmony output tests share the upstream `openai-harmony` tiktoken cache.
2 -//!
3 -//! Use a file lock for tests that load the encoding so `cargo nextest` cannot
4 -//! start multiple processes that concurrently populate the same cache file.
5 -
1 use std::sync::Arc;
2
3 use futures::executor::block_on;
4 use futures::{TryStreamExt as _, stream};
5 use openai_harmony::chat::{Message, Role};
11 -use serial_test::file_serial;
6 use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, Finished};
7
8 use super::*;
@@ -91,7 +85,6 @@ fn request_with_tools() -> ChatRequest {
85 }
86
87 #[test]
94 -#[file_serial(harmony_tiktoken_cache)]
88 fn interrupted_final_message_is_preserved() {
89 let tokens = completion_tokens(&[text_message("final", "hello")]);
90 let events = block_on(collect_events(
@@ -127,7 +120,6 @@ fn interrupted_final_message_is_preserved() {
120 }
121
122 #[test]
130 -#[file_serial(harmony_tiktoken_cache)]
123 fn eos_flush_preserves_trailing_replacement_text() {
124 let mut tokens = completion_tokens(&[text_message("final", "Hi")]);
125 tokens.pop();
@@ -153,7 +145,6 @@ fn eos_flush_preserves_trailing_replacement_text() {
145 }
146
147 #[test]
156 -#[file_serial(harmony_tiktoken_cache)]
148 fn interrupted_analysis_message_is_preserved() {
149 let tokens = completion_tokens(&[text_message("analysis", "think")]);
150 let events = block_on(collect_events(
@@ -189,7 +180,6 @@ fn interrupted_analysis_message_is_preserved() {
180 }
181
182 #[test]
192 -#[file_serial(harmony_tiktoken_cache)]
183 fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
184 let tokens = completion_tokens(&[
185 text_message("commentary", "Let me check."),
@@ -217,7 +207,6 @@ fn commentary_preamble_is_visible_but_commentary_tool_payload_is_not() {
207 }
208
209 #[test]
220 -#[file_serial(harmony_tiktoken_cache)]
210 fn multiple_messages_get_newline_separators() {
211 let tokens = completion_tokens(&[
212 text_message("analysis", "first think"),
@@ -249,7 +238,6 @@ fn multiple_messages_get_newline_separators() {
238 }
239
240 #[test]
252 -#[file_serial(harmony_tiktoken_cache)]
241 fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
242 let tokens = completion_tokens(&[tool_message(
243 "get_weather",
@@ -302,7 +290,6 @@ fn tool_calls_stream_arguments_and_finish_with_local_id_shape() {
290 }
291
292 #[test]
305 -#[file_serial(harmony_tiktoken_cache)]
293 fn semantic_events_precede_same_update_logprobs() {
294 let tokens = completion_tokens(&[text_message("final", "hello")]);
295 let events = block_on(collect_events(
@@ -353,7 +340,6 @@ fn rejects_generic_parser_overrides() {
340 }
341
342 #[test]
356 -#[file_serial(harmony_tiktoken_cache)]
343 fn allows_auto_auto_only() {
344 validate_harmony_parser_overrides(&ParserSelection::Auto, &ParserSelection::Auto).unwrap();
345 let _ = HarmonyChatOutputProcessor::new(&ChatRequest::for_test()).unwrap();
rust/src/chat/tests/roundtrip.rs
+160 -65
@@ -1,8 +1,8 @@
1 -//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
1 +//! Roundtrip tests for the real chat-template and output-processor pairing.
2 //!
3 //! The invariant under test is that a structured assistant message rendered as history can be
4 //! parsed from the generated assistant completion and then rendered back to the exact same
5 -//! assistant-completion text.
5 +//! assistant completion.
6
7 use std::pin::Pin;
8 use std::sync::Arc;
@@ -18,6 +18,10 @@ use vllm_chat::{
18 RendererSelection, load_model_backends,
19 };
20 use vllm_text::{DecodedTextEvent, Finished, Prompt};
21 +use vllm_tokenizer::Tokenizer;
22 +
23 +const TEXT_COMPLETION_CHUNK_CHARS: usize = 7;
24 +const TOKEN_COMPLETION_CHUNK_TOKENS: usize = 1;
25
26 /// One model/parser configuration used to run the fixed roundtrip fixtures.
27 #[derive(Clone)]
@@ -191,14 +195,28 @@ impl RoundtripCase {
195 sort_json_keys: false,
196 }
197 }
198 +
199 + /// GPT-OSS Harmony token-id renderer and native Harmony output processor.
200 + fn gpt_oss() -> Self {
201 + Self {
202 + model_id: "openai/gpt-oss-20b",
203 + assistant_stop_suffix: "", // not applicable for token-id cases
204 + tool_call_parser: ParserSelection::Auto,
205 + reasoning_parser: ParserSelection::Auto,
206 + thinking_behavior: ThinkingBehavior::Always { value: true },
207 + json_fmt: compact_json_fmt(),
208 + sort_json_keys: false,
209 + }
210 + }
211 }
212
213 macro_rules! roundtrip_tests {
197 - ($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
214 + ($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => {
215 paste::paste! {
216 $(
217 $(
218 #[tokio::test]
219 + $(#[$fixture_attr])*
220 #[file_serial([<hf_ $case>])]
221 async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
222 [<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
@@ -217,9 +235,9 @@ roundtrip_tests! {
235 glm47 => [reasoning_and_content, tool_call_mix],
236 seed_oss => [reasoning_and_content],
237 step3p5 => [reasoning_and_content],
220 -
238 gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call
239 kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history
240 + gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call
241 }
242
243 /// Run the fixed reasoning+content fixture for one model/parser case.
@@ -421,10 +439,10 @@ struct RoundtripResult {
439 parsed_message: AssistantMessage,
440 /// Assistant-completion suffix cut from rendering the expected assistant as
441 /// history.
424 - closed_completion: String,
442 + closed_completion: Prompt,
443 /// Assistant-completion suffix cut after rendering the parsed assistant
444 /// back as history.
427 - rerendered_closed_completion: String,
445 + rerendered_closed_completion: Prompt,
446 }
447
448 /// Render, parse, and rerender one assistant turn through the production
@@ -436,60 +454,59 @@ async fn run_roundtrip(
454 assistant: AssistantMessage,
455 ) -> Result<RoundtripResult> {
456 let renderer = backends.chat_backend.chat_renderer();
439 - let (prompt, closed_completion_text) =
440 - render_closed_completion(renderer.as_ref(), request, &assistant)?;
441 - let completion_body = closed_completion_text
442 - .strip_suffix(case.assistant_stop_suffix)
443 - .with_context(|| {
444 - format!(
445 - "closed assistant completion did not end with {:?}: {:?}",
446 - case.assistant_stop_suffix, closed_completion_text
447 - )
448 - })?;
449 -
450 - let parsed_message =
451 - parse_completion(case, backends, request, &prompt, completion_body).await?;
452 - let (_, rerendered_closed_completion) =
453 - render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
457 + let rendered = render_closed_completion(renderer.as_ref(), request, &assistant)?;
458 +
459 + let parsed_message = parse_completion(case, backends, request, &rendered).await?;
460 + let rerendered = render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
461
462 Ok(RoundtripResult {
463 parsed_message,
457 - closed_completion: closed_completion_text,
458 - rerendered_closed_completion,
464 + closed_completion: rendered.completion,
465 + rerendered_closed_completion: rerendered.completion,
466 })
467 }
468
469 +/// Rendered prompt/completion artifacts at the renderer boundary.
470 +struct RenderedTurn {
471 + prompt: Prompt,
472 + completion: Prompt,
473 +}
474 +
475 /// Render `history` as a production prompt and `history + assistant` as closed
476 /// history, then return the production prompt and assistant-completion suffix.
477 fn render_closed_completion(
478 renderer: &dyn vllm_chat::ChatRenderer,
479 base_request: &ChatRequest,
480 assistant: &AssistantMessage,
468 -) -> Result<(String, String)> {
481 +) -> Result<RenderedTurn> {
482 let mut prompt_request = base_request.clone();
483 prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
471 - let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
484 + let prompt = renderer.render(&prompt_request).context("failed to render prompt")?.prompt;
485
486 let mut full_request = base_request.clone();
487 full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
488 full_request.messages.push(ChatMessage::from(assistant.clone()));
476 - let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
477 -
478 - ensure!(
479 - full.starts_with(&prompt),
480 - "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
481 - );
482 - let completion = full[prompt.len()..].to_string();
483 -
484 - Ok((prompt, completion))
485 -}
489 + let full = renderer.render(&full_request).context("failed to render full prompt")?.prompt;
490 +
491 + let completion = match (&prompt, full) {
492 + (Prompt::Text(prompt), Prompt::Text(full)) => {
493 + ensure!(
494 + full.starts_with(prompt),
495 + "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
496 + );
497 + Prompt::Text(full[prompt.len()..].to_string())
498 + }
499 + (Prompt::TokenIds(prompt), Prompt::TokenIds(full)) => {
500 + ensure!(
501 + full.starts_with(prompt),
502 + "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
503 + );
504 + Prompt::TokenIds(full[prompt.len()..].to_vec())
505 + }
506 + (prompt, full) => bail!("prompt kind changed between renders: {prompt:?} vs {full:?}"),
507 + };
508
487 -/// Render one chat request and require a text prompt.
488 -fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
489 - match renderer.render(request)?.prompt {
490 - Prompt::Text(text) => Ok(text),
491 - other => bail!("roundtrip tests expect text prompts, got {other:?}"),
492 - }
509 + Ok(RenderedTurn { prompt, completion })
510 }
511
512 /// Feed one rendered assistant completion body into the real output processor
@@ -498,13 +515,15 @@ async fn parse_completion(
515 case: &RoundtripCase,
516 backends: &vllm_chat::LoadedModelBackends,
517 base_request: &ChatRequest,
501 - prompt: &str,
502 - completion_body: &str,
518 + rendered: &RenderedTurn,
519 ) -> Result<AssistantMessage> {
520 let tokenizer = backends.text_backend.tokenizer();
505 - let prompt_token_ids = tokenizer
506 - .encode(prompt, base_request.add_special_tokens)
507 - .context("failed to encode rendered prompt")?;
521 + let prompt_token_ids = match &rendered.prompt {
522 + Prompt::Text(prompt) => tokenizer
523 + .encode(prompt, base_request.add_special_tokens)
524 + .context("failed to encode rendered prompt")?,
525 + Prompt::TokenIds(token_ids) => token_ids.clone(),
526 + };
527
528 let mut request = base_request.clone();
529 let processor = backends.chat_backend.new_chat_output_processor(
@@ -515,7 +534,12 @@ async fn parse_completion(
534 },
535 )?;
536
518 - let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
537 + let decoded = decoded_completion_stream(
538 + tokenizer.as_ref(),
539 + prompt_token_ids,
540 + &rendered.completion,
541 + case.assistant_stop_suffix,
542 + )?;
543 let mut events = processor.process(decoded)?;
544
545 while let Some(event) = events.next().await {
@@ -538,16 +562,46 @@ async fn parse_completion(
562 /// split into small chunks to exercise streaming parser state across marker
563 /// and JSON boundaries.
564 fn decoded_completion_stream(
565 + tokenizer: &dyn Tokenizer,
566 prompt_token_ids: Vec<u32>,
542 - completion_body: &str,
543 -) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
544 - let prompt_token_count = prompt_token_ids.len();
567 + completion: &Prompt,
568 + assistant_stop_suffix: &str,
569 +) -> Result<Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>>> {
570 let mut events = vec![DecodedTextEvent::Start {
546 - prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
571 + prompt_token_ids: Arc::from(prompt_token_ids.clone().into_boxed_slice()),
572 prompt_logprobs: None,
573 }];
574
550 - let chunks = split_by_chars(completion_body, 7);
575 + let chunks = match completion {
576 + Prompt::Text(text) => {
577 + let body = text.strip_suffix(assistant_stop_suffix).with_context(|| {
578 + format!(
579 + "closed assistant completion did not end with {:?}: {:?}",
580 + assistant_stop_suffix, text
581 + )
582 + })?;
583 + split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
584 + .into_iter()
585 + .map(|delta| DecodedCompletionChunk {
586 + delta,
587 + token_ids: Vec::new(), // unused for text-level roundtrip cases
588 + })
589 + .collect()
590 + }
591 + Prompt::TokenIds(token_ids) => {
592 + ensure!(
593 + assistant_stop_suffix.is_empty(),
594 + "token-id roundtrip cases do not support text stop suffixes"
595 + );
596 + incremental_decode_chunks(
597 + tokenizer,
598 + &prompt_token_ids,
599 + token_ids,
600 + TOKEN_COMPLETION_CHUNK_TOKENS,
601 + )?
602 + }
603 + };
604 +
605 if chunks.is_empty() {
606 events.push({
607 DecodedTextEvent::TextDelta {
@@ -555,11 +609,7 @@ fn decoded_completion_stream(
609 token_ids: Vec::new(),
610 logprobs: None,
611 finished: Some(Finished {
558 - usage: vllm_llm::TokenUsage {
559 - prompt_token_count: 0,
560 - output_token_count: 0,
561 - cached_token_count: 0,
562 - },
612 + usage: Default::default(),
613 finish_reason: FinishReason::stop_eos(),
614 kv_transfer_params: None,
615 }),
@@ -569,24 +619,26 @@ fn decoded_completion_stream(
619 let last_index = chunks.len() - 1;
620 for (index, chunk) in chunks.into_iter().enumerate() {
621 let finished = (index == last_index).then(|| Finished {
572 - usage: vllm_llm::TokenUsage {
573 - prompt_token_count,
574 - output_token_count: completion_body.chars().count(),
575 - cached_token_count: 0,
576 - },
622 + usage: Default::default(),
623 finish_reason: FinishReason::stop_eos(),
624 kv_transfer_params: None,
625 });
626 events.push(DecodedTextEvent::TextDelta {
581 - delta: chunk,
582 - token_ids: Vec::new(),
627 + delta: chunk.delta,
628 + token_ids: chunk.token_ids,
629 logprobs: None,
630 finished,
631 });
632 }
633 }
634
589 - stream::iter(events).map(Ok).boxed()
635 + Ok(stream::iter(events).map(Ok).boxed())
636 +}
637 +
638 +/// One decoded completion chunk fed into the output processor.
639 +struct DecodedCompletionChunk {
640 + delta: String,
641 + token_ids: Vec<u32>,
642 }
643
644 /// Split text into chunks containing at most `chunk_chars` Unicode scalar
@@ -612,6 +664,49 @@ fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
664 chunks
665 }
666
667 +/// Split token ids into chunks containing at most `chunk_size` ids.
668 +fn split_by_count(token_ids: &[u32], chunk_size: usize) -> Vec<Vec<u32>> {
669 + token_ids.chunks(chunk_size).map(<[u32]>::to_vec).collect()
670 +}
671 +
672 +/// Decode token ids incrementally using the production tokenizer stream.
673 +fn incremental_decode_chunks(
674 + tokenizer: &dyn Tokenizer,
675 + prompt_token_ids: &[u32],
676 + token_ids: &[u32],
677 + chunk_size: usize,
678 +) -> Result<Vec<DecodedCompletionChunk>> {
679 + let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
680 + let mut chunks = Vec::new();
681 + for chunk_token_ids in split_by_count(token_ids, chunk_size) {
682 + let mut delta = String::new();
683 + for token_id in chunk_token_ids.iter().copied() {
684 + decoder.push_token(token_id)?;
685 + while let Some(chunk) = decoder.next_chunk() {
686 + delta.push_str(&chunk);
687 + }
688 + }
689 + chunks.push(DecodedCompletionChunk {
690 + delta,
691 + token_ids: chunk_token_ids,
692 + });
693 + }
694 +
695 + let (last_chunk, _) = decoder.flush(None)?;
696 + if let Some(last_chunk) = last_chunk {
697 + if let Some(delta) = chunks.last_mut() {
698 + delta.delta.push_str(&last_chunk);
699 + } else {
700 + chunks.push(DecodedCompletionChunk {
701 + delta: last_chunk,
702 + token_ids: Vec::new(),
703 + });
704 + }
705 + }
706 +
707 + Ok(chunks)
708 +}
709 +
710 /// Build a chat request fixture with parser-enabling tool-choice semantics.
711 fn roundtrip_request(
712 request_id: impl Into<String>,