[Rust Frontend] Support `generation` blocks in HF chat template (#56378)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
Sep 12, 2026 at 12:56 UTC
eed1f3d0c6043bd494424a22443ee198dd56f657
3 files changed
+207
rust/src/chat/src/renderer/hf/generation.rs
new
+187
@@ -0,0 +1,187 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+//! Support Hugging Face's `{% generation %}` blocks in MiniJinja chat templates.
5
+//!
6
+//! Transformers uses these blocks to track assistant-token masks via scoped call
7
+//! blocks. Inference only needs their rendered text. MiniJinja lacks this tag, so
8
+//! we rewrite `{% generation %}...{% endgeneration %}` to
9
+//! `{% call __hf_generation() %}...{% endcall %}` before compiling the template.
10
+//! The registered `__hf_generation` callback renders and returns `caller()`,
11
+//! preserving the block's local scope just as Transformers' AssistantTracker does.
12
+//!
13
+//! We use MiniJinja's lexer to identify statement tags and replace only their
14
+//! keyword spans. This preserves delimiters and whitespace controls (`-` / `+`),
15
+//! while keeping literals, comments, and raw blocks intact. Original `generation`
16
+//! and `call` pairings are checked before rewriting both to call blocks.
17
+
18
+use minijinja::machinery::{Token, tokenize};
19
+use minijinja::value::{Kwargs, Value};
20
+use minijinja::{Error, ErrorKind, State};
21
+
22
+#[derive(Clone, Copy, PartialEq, Eq)]
23
+enum CallBlockKind {
24
+ Call,
25
+ Generation,
26
+}
27
+
28
+/// Rewrite Hugging Face's assistant-mask blocks as scoped MiniJinja call blocks.
29
+pub(super) fn rewrite_generation_blocks(source: String) -> Result<String, Error> {
30
+ let tokens = tokenize(&source, false, Default::default(), Default::default())
31
+ .collect::<Result<Vec<_>, _>>()?;
32
+ // Track native calls too: rewriting makes both block kinds use `endcall`,
33
+ // so validate the original pairing before that distinction is lost.
34
+ let mut blocks = Vec::new();
35
+ let mut replacements = Vec::new();
36
+
37
+ for (index, pair) in tokens.windows(2).enumerate() {
38
+ let [(Token::BlockStart, _), (Token::Ident(name), span)] = pair else {
39
+ continue;
40
+ };
41
+ match *name {
42
+ "call" => blocks.push(CallBlockKind::Call),
43
+ "endcall" if blocks.pop() == Some(CallBlockKind::Generation) => {
44
+ return Err(Error::new(
45
+ ErrorKind::SyntaxError,
46
+ "expected endgeneration, got endcall",
47
+ ));
48
+ }
49
+ "generation" | "endgeneration"
50
+ if matches!(tokens.get(index + 2), Some((Token::BlockEnd, _))) =>
51
+ {
52
+ let replacement = if *name == "generation" {
53
+ blocks.push(CallBlockKind::Generation);
54
+ "call __hf_generation()"
55
+ } else {
56
+ if blocks.pop() != Some(CallBlockKind::Generation) {
57
+ return Err(Error::new(
58
+ ErrorKind::SyntaxError,
59
+ "unexpected endgeneration",
60
+ ));
61
+ }
62
+ "endcall"
63
+ };
64
+ replacements.push((
65
+ span.start_offset as usize..span.end_offset as usize,
66
+ replacement,
67
+ ));
68
+ }
69
+ _ => {}
70
+ }
71
+ }
72
+ if blocks.contains(&CallBlockKind::Generation) {
73
+ return Err(Error::new(
74
+ ErrorKind::SyntaxError,
75
+ "unclosed generation block",
76
+ ));
77
+ }
78
+
79
+ // Replace only keyword spans, preserving delimiters and their whitespace controls.
80
+ let mut source = source;
81
+ for (span, replacement) in replacements.into_iter().rev() {
82
+ source.replace_range(span, replacement);
83
+ }
84
+ Ok(source)
85
+}
86
+
87
+/// Render the caller with the same local scope as Transformers' AssistantTracker.
88
+pub(super) fn render_generation(state: &State, kwargs: Kwargs) -> Result<Value, Error> {
89
+ let caller: Value = kwargs.get("caller")?;
90
+ kwargs.assert_all_used()?;
91
+ caller.call(state, &[])
92
+}
93
+
94
+#[cfg(test)]
95
+mod tests {
96
+ use expect_test::expect;
97
+ use minijinja::Environment;
98
+
99
+ use super::*;
100
+
101
+ #[test]
102
+ fn generation_blocks_reject_invalid_syntax_and_mismatched_end_tags() {
103
+ let env = Environment::new();
104
+ for template in [
105
+ "{% generation %}unclosed",
106
+ "{% endgeneration %}",
107
+ "{% generation extra %}{% endgeneration %}",
108
+ "{% generation %}{% endgeneration extra %}",
109
+ "{% generation %}{% endcall %}",
110
+ "{% call f() %}{% endgeneration %}",
111
+ "{% generation %}{% endcall %}{% call f() %}{% endgeneration %}",
112
+ "{% generation %}{% if true %}{% endgeneration %}{% endif %}",
113
+ ] {
114
+ let result = rewrite_generation_blocks(template.to_string())
115
+ .and_then(|source| env.template_from_str(&source).map(|_| ()));
116
+ assert!(result.is_err(), "accepted invalid template: {template}");
117
+ }
118
+ }
119
+
120
+ #[test]
121
+ fn rewrite_generation_blocks_preserves_hf_rendering_semantics() {
122
+ let cases = [
123
+ (
124
+ "whitespace",
125
+ "前 {%- generation -%} 中 {%- endgeneration -%} 后",
126
+ ),
127
+ (
128
+ "keep_newlines",
129
+ "A{%+ generation +%}\nB{%+ endgeneration +%}\nC",
130
+ ),
131
+ ("literal", "{{ '{% generation %}' }}"),
132
+ ("raw", "{% raw %}{%- generation -%}{% endraw %}"),
133
+ ("comment", "A{# {% generation %} #}B"),
134
+ (
135
+ "scope",
136
+ "{% set x = 'outer' %}{% generation %}{% set x = 'inner' %}{{ x }}{% endgeneration %}{{ x }}",
137
+ ),
138
+ (
139
+ "nested_calls",
140
+ "{% macro wrap() %}[{{ caller() }}]{% endmacro %}{% generation %}A{% call wrap() %}{% generation %}B{% endgeneration %}{% endcall %}C{% endgeneration %}",
141
+ ),
142
+ ];
143
+ let results = cases.map(|(name, source)| {
144
+ let template = rewrite_generation_blocks(source.to_string()).unwrap();
145
+ let mut env = Environment::new();
146
+ env.set_trim_blocks(true);
147
+ env.set_lstrip_blocks(true);
148
+ env.add_function("__hf_generation", render_generation);
149
+ let rendered =
150
+ env.template_from_str(&template).unwrap().render(Value::UNDEFINED).unwrap();
151
+ (name, rendered)
152
+ });
153
+ expect![[r#"
154
+ [
155
+ (
156
+ "whitespace",
157
+ "前中后",
158
+ ),
159
+ (
160
+ "keep_newlines",
161
+ "A\nB\nC",
162
+ ),
163
+ (
164
+ "literal",
165
+ "{% generation %}",
166
+ ),
167
+ (
168
+ "raw",
169
+ "{%- generation -%}",
170
+ ),
171
+ (
172
+ "comment",
173
+ "AB",
174
+ ),
175
+ (
176
+ "scope",
177
+ "innerouter",
178
+ ),
179
+ (
180
+ "nested_calls",
181
+ "A[B]C",
182
+ ),
183
+ ]
184
+ "#]]
185
+ .assert_debug_eq(&results);
186
+ }
187
+}
rust/src/chat/src/renderer/hf/mod.rs
+17
@@ -25,6 +25,7 @@ use crate::{
25
26
mod error;
27
mod format;
28
+mod generation;
29
mod template;
30
mod tojson;
31
@@ -642,6 +643,22 @@ mod tests {
643
.map_err(|_| unreachable!("HF renderer should return text prompt"))
644
}
645
646
+ #[test]
647
+ fn generation_blocks_allow_content_format_detection_and_request_overrides() {
648
+ let template = "{% for message in messages %}{% generation %}{% for part in message.content %}{{ part.text }}{% endfor %}{% endgeneration %}{% endfor %}";
649
+ let mut request = sample_request(vec![ChatMessage::user("hello")]);
650
+ let default = render(Some(template), &request).unwrap();
651
+ request.chat_options.chat_template = Some(template.to_string());
652
+ let overridden = render(Some("unused"), &request).unwrap();
653
+ expect![[r#"
654
+ (
655
+ "hello",
656
+ "hello",
657
+ )
658
+ "#]]
659
+ .assert_debug_eq(&(default, overridden));
660
+ }
661
+
662
fn render_mm(
663
template: &str,
664
request: &ChatRequest,
rust/src/chat/src/renderer/hf/template.rs
+3
@@ -20,6 +20,7 @@ use super::error::TemplateError;
20
use super::format::{
21
ChatTemplateContentFormat, ChatTemplateContentFormatOption, detect_chat_template_content_format,
22
};
23
+use super::generation::{render_generation, rewrite_generation_blocks};
24
use super::tojson::hf_tojson_filter;
25
use crate::renderer::hf::{TemplateMessage, TemplateTool};
26
@@ -36,6 +37,7 @@ fn build_environment(template: String) -> Result<Environment<'static>> {
37
38
env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
39
env.add_filter("tojson", hf_tojson_filter);
40
+ env.add_function("__hf_generation", render_generation);
41
42
Ok(env)
43
}
@@ -107,6 +109,7 @@ pub(super) struct CompiledChatTemplate {
109
impl CompiledChatTemplate {
110
/// Compile the given chat template string into a [`CompiledChatTemplate`].
111
pub fn new(template: String, content_format: ChatTemplateContentFormatOption) -> Result<Self> {
112
+ let template = rewrite_generation_blocks(template)?;
113
let content_format = match content_format {
114
ChatTemplateContentFormatOption::Auto => detect_chat_template_content_format(&template),
115
ChatTemplateContentFormatOption::String => ChatTemplateContentFormat::String,