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

[Rust Frontend] Handle `continue_final_message` with renderer sentinel (#47844)

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

Bugen Zhao committed Jul 8, 2026 at 19:57 UTC 2cae98dfa59bc58d939f526465ff025670d002f0
3 files changed +200 -17
.buildkite/test-amd.yaml
+1 -1
@@ -2026,7 +2026,7 @@ steps:
2026 - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
2027 - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
2028 - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
2029 - - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
2029 + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"
2030
2031 - label: Rust Frontend Core Correctness # TBD
2032 timeout_in_minutes: 180
.buildkite/test_areas/rust_frontend.yaml
+1 -1
@@ -64,7 +64,7 @@ steps:
64 - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
65 # - pytest -v -s entrypoints/serve/dev/test_sleep.py
66 # /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op).
67 - - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
67 + - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info"
68
69 - label: Rust Frontend Core Correctness
70 timeout_in_minutes: 30
rust/src/chat/src/renderer/hf/mod.rs
+198 -15
@@ -155,11 +155,23 @@ impl HfChatRenderer {
155 effective_template: &CompiledChatTemplate,
156 request: &ChatRequest,
157 ) -> Result<RenderedPrompt> {
158 - let messages = to_template_messages(
158 + let mut messages = to_template_messages(
159 &request.messages,
160 effective_template.content_format(),
161 self.multimodal.as_ref(),
162 )?;
163 +
164 + // Handling of `continue_final_message`:
165 + // Append a sentinel tag to the final message content, render as usual, then
166 + // truncate the rendered prompt at the tag so any template suffix after the
167 + // final message content (e.g. the end-of-turn marker) is dropped.
168 + let final_message_text = if request.chat_options.continue_final_message() {
169 + let final_message = messages.last_mut().ok_or(Error::EmptyMessages)?;
170 + Some(append_continue_final_message_tag(final_message)?)
171 + } else {
172 + None
173 + };
174 +
175 let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools));
176 trace!(
177 message_count = messages.len(),
@@ -183,6 +195,13 @@ impl HfChatRenderer {
195 })
196 .map_err(|error| Error::ChatTemplate(error.to_report_string()))?;
197
198 + let prompt = match &final_message_text {
199 + Some(final_message_text) => {
200 + truncate_prompt_at_continue_final_message_tag(prompt, final_message_text)?
201 + }
202 + None => prompt,
203 + };
204 +
205 trace!(
206 prompt_len = prompt.len(),
207 prompt, "rendered chat template prompt"
@@ -429,6 +448,74 @@ fn to_template_string_content(
448 }
449 }
450
451 +/// Sentinel appended to the final message content when `continue_final_message`
452 +/// is requested, used to locate the truncation point in the rendered prompt.
453 +///
454 +/// Same literal as `transformers`. Occurrences of this string earlier in the
455 +/// prompt are harmless because truncation uses the rightmost match, and the
456 +/// appended sentinel ends up last as long as the template renders messages in
457 +/// order.
458 +const CONTINUE_FINAL_MESSAGE_TAG: &str = "CONTINUE_FINAL_MESSAGE_TAG ";
459 +
460 +/// Append [`CONTINUE_FINAL_MESSAGE_TAG`] to the trailing text of the final
461 +/// message, returning the original text for post-render validation.
462 +// TODO: transformers v5 also allows continuing a non-`content` field (e.g.
463 +// `reasoning_content`) by passing a field name; only the boolean form is
464 +// supported here.
465 +fn append_continue_final_message_tag(message: &mut TemplateMessage) -> Result<String> {
466 + let text = match &mut message.content {
467 + TemplateContent::String(text) => Some(text),
468 + // Pick the last text part in the message.
469 + TemplateContent::OpenAi(parts) => parts.iter_mut().rev().find_map(|part| match part {
470 + TemplateContentPart::Text { text } => Some(text),
471 + TemplateContentPart::Image => None,
472 + }),
473 + };
474 + let text = text.ok_or_else(|| {
475 + Error::ChatTemplate(
476 + "continue_final_message is set but there is no text to continue \
477 + in the final message"
478 + .to_string(),
479 + )
480 + })?;
481 +
482 + let original = text.clone();
483 + text.push_str(CONTINUE_FINAL_MESSAGE_TAG);
484 + Ok(original)
485 +}
486 +
487 +/// Truncate the rendered prompt at [`CONTINUE_FINAL_MESSAGE_TAG`] so that it
488 +/// ends exactly with the final message content, dropping any template suffix
489 +/// such as end-of-turn markers.
490 +fn truncate_prompt_at_continue_final_message_tag(
491 + mut rendered: String,
492 + final_message_text: &str,
493 +) -> Result<String> {
494 + let tag_loc = rendered
495 + .rfind(CONTINUE_FINAL_MESSAGE_TAG.trim_end())
496 + .filter(|_| rendered.contains(final_message_text.trim()));
497 + let Some(tag_loc) = tag_loc else {
498 + return Err(Error::ChatTemplate(format!(
499 + "continue_final_message is set but the final message does not appear \
500 + in the prompt after applying the chat template! This can happen if \
501 + the chat template deletes portions of the final message. Final \
502 + message to continue: {}",
503 + final_message_text.trim(),
504 + )));
505 + };
506 +
507 + if rendered[tag_loc..].starts_with(CONTINUE_FINAL_MESSAGE_TAG) {
508 + // The template preserved spacing, so a plain cut at the tag suffices.
509 + rendered.truncate(tag_loc);
510 + } else {
511 + // The template trimmed the trailing spacing of the message content, so
512 + // apply the same trimming to the retained prefix.
513 + rendered.truncate(tag_loc);
514 + rendered.truncate(rendered.trim_end().len());
515 + }
516 + Ok(rendered)
517 +}
518 +
519 fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
520 tools
521 .iter()
@@ -548,28 +635,124 @@ mod tests {
635 ChatRole::Assistant,
636 "The capital of",
637 )]);
638 + let template =
639 + "{% if continue_final_message %}continue:{% endif %}{{ messages[0].content }}";
640
552 - assert_eq!(
553 - render(
554 - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
555 - &request,
556 - )
557 - .unwrap(),
558 - "new"
559 - );
641 + assert_eq!(render(Some(template), &request).unwrap(), "The capital of");
642
643 request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
644
645 assert_eq!(
564 - render(
565 - Some("{% if continue_final_message %}continue{% else %}new{% endif %}"),
566 - &request,
567 - )
568 - .unwrap(),
569 - "continue"
646 + render(Some(template), &request).unwrap(),
647 + "continue:The capital of"
648 );
649 }
650
651 + #[test]
652 + fn continue_final_message_truncates_template_suffix() {
653 + let mut request = sample_request(vec![
654 + ChatMessage::text(ChatRole::User, "What is the capital of France?"),
655 + ChatMessage::text(ChatRole::Assistant, "The capital of"),
656 + ]);
657 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
658 +
659 + // The Qwen3 template is unaware of `continue_final_message`; the
660 + // end-of-turn marker it appends must still be stripped.
661 + let rendered = render(Some(QWEN3_0_6B_TEMPLATE), &request).unwrap();
662 +
663 + expect![[r#"
664 + <|im_start|>user
665 + What is the capital of France?<|im_end|>
666 + <|im_start|>assistant
667 + <think>
668 +
669 + </think>
670 +
671 + The capital of"#]]
672 + .assert_eq(&rendered);
673 + }
674 +
675 + #[test]
676 + fn continue_final_message_trims_like_the_template_does() {
677 + let mut request = sample_request(vec![ChatMessage::text(ChatRole::Assistant, "Sure, ")]);
678 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
679 +
680 + // The template trims the trailing spacing of the message content, so
681 + // the truncated prompt must be trimmed the same way.
682 + let rendered = render(
683 + Some("{{ messages[0].content.strip() }}<|im_end|>"),
684 + &request,
685 + )
686 + .unwrap();
687 +
688 + assert_eq!(rendered, "Sure,");
689 + }
690 +
691 + #[test]
692 + fn continue_final_message_appends_to_last_text_part() {
693 + // The renderer itself is role-agnostic like transformers (the
694 + // assistant-final restriction is enforced by request validation
695 + // upstream), so a multimodal user message exercises the part
696 + // selection: the sentinel must attach to the last *text* part,
697 + // skipping the trailing image.
698 + let mut request = sample_request(vec![ChatMessage::user(vec![
699 + ChatContentPart::text("Sure,"),
700 + ChatContentPart::image_url("data:image/png;base64,test"),
701 + ])]);
702 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
703 +
704 + let rendered = render_mm(
705 + "{% for item in messages[0].content %}{% if item.type == 'image' %}<image>{% else %}{{ item.text }}{% endif %}{% endfor %}<|im_end|>",
706 + &request,
707 + ChatTemplateContentFormatOption::OpenAi,
708 + )
709 + .unwrap()
710 + .prompt;
711 +
712 + // Anything rendered after the continued text (here the image
713 + // placeholder and the end marker) is truncated away, matching
714 + // transformers.
715 + assert_eq!(rendered, Prompt::Text("Sure,".to_string()));
716 + }
717 +
718 + #[test]
719 + fn continue_final_message_composes_with_aware_templates() {
720 + // A template that reads `continue_final_message` and skips its own
721 + // end-of-turn marker must produce the same prompt as an unaware one:
722 + // the sentinel truncation degenerates to a cut at the very end.
723 + let mut request = sample_request(vec![
724 + ChatMessage::text(ChatRole::User, "hi"),
725 + ChatMessage::text(ChatRole::Assistant, "Sure,"),
726 + ]);
727 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
728 +
729 + let aware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}{% if not (loop.last and continue_final_message) %}<|im_end|>\n{% endif %}{% endfor %}";
730 + let unaware = "{% for m in messages %}<|im_start|>{{ m.role }}\n{{ m.content }}<|im_end|>\n{% endfor %}";
731 +
732 + let expected = "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\nSure,";
733 + assert_eq!(render(Some(aware), &request).unwrap(), expected);
734 + assert_eq!(render(Some(unaware), &request).unwrap(), expected);
735 + }
736 +
737 + #[test]
738 + fn continue_final_message_errors_when_template_drops_final_message() {
739 + let mut request = sample_request(vec![
740 + ChatMessage::text(ChatRole::User, "hi"),
741 + ChatMessage::text(ChatRole::Assistant, "Sure,"),
742 + ]);
743 + request.chat_options.generation_prompt_mode = GenerationPromptMode::ContinueFinalAssistant;
744 +
745 + let error = render(
746 + Some(
747 + "{% for m in messages %}{% if m.role == 'user' %}{{ m.content }}{% endif %}{% endfor %}",
748 + ),
749 + &request,
750 + )
751 + .unwrap_err();
752 +
753 + assert!(matches!(error, Error::ChatTemplate(_)));
754 + }
755 +
756 #[test]
757 fn chat_template_flattens_text_parts_for_string_templates() {
758 let request = sample_request(vec![ChatMessage::user(vec![