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

[Hy4] support Hy4-preview model (#54160)

Signed-off-by: chengvjiang <chengvjiang@tencent.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: chengvjiang <chengvjiang@tencent.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: russellfeng <russellfeng@tencent.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Cheng Jiang committed Aug 29, 2026 at 12:07 UTC b2f685834a6456197e7033966fdef52a23f1abcd
47 files changed +6260 -217
docs/models/supported_models.md
+2 -1
@@ -405,7 +405,8 @@ th {
405 | `GraniteMoeSWAForCausalLM` | Granite MoE SWA | `ibm-granite/granite-swash-3b-a600m` | ✅︎ | ✅︎ |
406 | `GraniteSWAForCausalLM` | Granite SWA | `ibm-granite/granite-swash-2b` | ✅︎ | ✅︎ |
407 | `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | |
408 -| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
408 +| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview`, `tencent/Hy3`, `tencent/Hy3-FP8` | ✅︎ | ✅︎ |
409 +| `HYV4ForCausalLM` | HY4 | `tencent/Hy4-preview`, `tencent/Hy4-preview-FP8` | ✅︎ | ✅︎ |
410 | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
411 | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
412 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
rust/src/chat/src/lib.rs
+2 -2
@@ -405,7 +405,7 @@ mod tests {
405 )
406 .unwrap_err();
407
408 - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
408 + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string());
409 }
410
411 #[test]
@@ -416,6 +416,6 @@ mod tests {
416 )
417 .unwrap_err();
418
419 - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, hy_v3, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
419 + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
420 }
421 }
rust/src/chat/src/parser/reasoning/mod.rs
+3
@@ -26,6 +26,7 @@ pub mod names {
26 pub const INKLING: &str = "inkling";
27 pub const GLM45: &str = "glm45";
28 pub const HY_V3: &str = "hy_v3";
29 + pub const HY_V4: &str = "hy_v4";
30 pub const KIMI: &str = "kimi";
31 pub const KIMI_K2: &str = "kimi_k2";
32 pub const KIMI_K3: &str = "kimi_k3";
@@ -69,6 +70,7 @@ impl ReasoningParserFactory {
70 .register_unified_dummy(names::INKLING)
71 .register_parser::<Glm45ReasoningParser>(names::GLM45)
72 .register_unified_dummy(names::HY_V3)
73 + .register_unified_dummy(names::HY_V4)
74 .register_parser::<KimiReasoningParser>(names::KIMI)
75 .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
76 .register_unified_dummy(names::KIMI_K3)
@@ -94,6 +96,7 @@ impl ReasoningParserFactory {
96 .register_pattern("glm-4.6", names::GLM45)
97 .register_pattern("glm-4.5", names::GLM45)
98 .register_pattern("hy3", names::HY_V3)
99 + .register_pattern("hy4", names::HY_V4)
100 .register_pattern("kimi-k2", names::KIMI_K2)
101 .register_pattern("kimi", names::KIMI)
102 // step3p5 patterns must precede `step3`: substring matching would
rust/src/chat/src/parser/tool/mod.rs
+3
@@ -29,6 +29,7 @@ pub mod names {
29 pub const GRANITE4: &str = "granite4";
30 pub const HERMES: &str = "hermes";
31 pub const HY_V3: &str = "hy_v3";
32 + pub const HY_V4: &str = "hy_v4";
33 // Matches the Python CLI name `--tool-call-parser internlm`, which Python
34 // also routes to `Internlm2ToolParser` despite the version-agnostic name.
35 pub const INTERNLM: &str = "internlm";
@@ -77,6 +78,7 @@ impl ToolParserFactory {
78 .register_parser::<Granite4ToolParser>(names::GRANITE4)
79 .register_parser::<HermesToolParser>(names::HERMES)
80 .register_unified_dummy(names::HY_V3)
81 + .register_unified_dummy(names::HY_V4)
82 .register_parser::<Internlm2ToolParser>(names::INTERNLM)
83 .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
84 .register_unified_dummy(names::KIMI_K3)
@@ -100,6 +102,7 @@ impl ToolParserFactory {
102 .register_pattern("qwen3", names::QWEN3_XML)
103 .register_pattern("hermes", names::HERMES)
104 .register_pattern("hy3", names::HY_V3)
105 + .register_pattern("hy4", names::HY_V4)
106 // Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
107 // and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
108 // routes to Llama), `internlm3-*` (also Llama-architecture per
rust/src/chat/src/parser/tool/tests.rs
+4
@@ -183,6 +183,10 @@ fn factory_new_resolves_default_patterns() {
183 factory.resolve_name_for_model("tencent/Hy3-preview"),
184 Some(names::HY_V3)
185 );
186 + assert_eq!(
187 + factory.resolve_name_for_model("tencent/Hy4-preview"),
188 + Some(names::HY_V4)
189 + );
190 assert_eq!(
191 factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
192 Some(names::MINIMAX_M3)
rust/src/chat/src/parser/unified.rs
+34 -2
@@ -6,8 +6,8 @@
6 use std::sync::LazyLock;
7
8 pub use vllm_parser::unified::{
9 - Gemma4UnifiedParser, HyV3UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser,
10 - UnifiedParser,
9 + Gemma4UnifiedParser, HyV3UnifiedParser, HyV4UnifiedParser, InklingUnifiedParser,
10 + KimiK3UnifiedParser, UnifiedParser,
11 };
12 use vllm_tokenizer::DynTokenizer;
13
@@ -18,6 +18,7 @@ use crate::request::ChatTool;
18 pub mod names {
19 pub const GEMMA4: &str = "gemma4";
20 pub const HY_V3: &str = "hy_v3";
21 + pub const HY_V4: &str = "hy_v4";
22 pub const INKLING: &str = "inkling";
23 pub const KIMI_K3: &str = "kimi_k3";
24 }
@@ -44,6 +45,7 @@ impl UnifiedParserFactory {
45
46 factory.register_parser::<Gemma4UnifiedParser>(names::GEMMA4);
47 factory.register_parser::<HyV3UnifiedParser>(names::HY_V3);
48 + factory.register_parser::<HyV4UnifiedParser>(names::HY_V4);
49 factory.register_parser::<InklingUnifiedParser>(names::INKLING);
50 factory.register_parser::<KimiK3UnifiedParser>(names::KIMI_K3);
51
@@ -51,6 +53,7 @@ impl UnifiedParserFactory {
53 .register_pattern("gemma-4", names::GEMMA4)
54 .register_pattern("gemma4", names::GEMMA4)
55 .register_pattern("hy3", names::HY_V3)
56 + .register_pattern("hy4", names::HY_V4)
57 .register_pattern("inkling", names::INKLING)
58 .register_pattern("kimi-k3", names::KIMI_K3)
59 .register_pattern("kimi_k3", names::KIMI_K3);
@@ -147,4 +150,33 @@ mod tests {
150 );
151 factory.create(names::KIMI_K3, &[], Arc::new(tokenizer)).unwrap();
152 }
153 +
154 + #[test]
155 + fn factory_registers_hy_v4() {
156 + let factory = UnifiedParserFactory::new();
157 + let tokenizer = [
158 + "<think:opensource>",
159 + "</think:opensource>",
160 + "<tool_calls:opensource>",
161 + "</tool_calls:opensource>",
162 + "<tool_call:opensource>",
163 + "</tool_call:opensource>",
164 + "<arg_key:opensource>",
165 + "</arg_key:opensource>",
166 + "<arg_value:opensource>",
167 + "</arg_value:opensource>",
168 + ]
169 + .into_iter()
170 + .enumerate()
171 + .fold(TestTokenizer::new(), |tokenizer, (index, token)| {
172 + tokenizer.with_regular_token(token, 1000 + index as u32)
173 + });
174 +
175 + assert!(factory.contains(names::HY_V4));
176 + assert_eq!(
177 + factory.resolve_name_for_model("tencent/Hy4-preview"),
178 + Some(names::HY_V4)
179 + );
180 + factory.create(names::HY_V4, &[], Arc::new(tokenizer)).unwrap();
181 + }
182 }
rust/src/parser/src/reasoning/hy.rs renamed
+7 -7
@@ -5,13 +5,13 @@ use vllm_tokenizer::DynTokenizer;
5
6 use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningError, ReasoningParser, Result};
7
8 -/// Internal HY3 reasoning stage used by the unified HY3 parser.
9 -pub(crate) struct HyV3ReasoningParser {
8 +/// Internal HY reasoning stage used by the unified HY parsers.
9 +pub(crate) struct HyReasoningParser {
10 inner: DelimitedReasoningParser,
11 }
12
13 -impl HyV3ReasoningParser {
14 - /// Create a HY3 reasoning parser for the tokenizer-specific marker suffix.
13 +impl HyReasoningParser {
14 + /// Create a HY reasoning parser for the tokenizer-specific marker suffix.
15 pub(crate) fn new(tokenizer: DynTokenizer, suffix: &str) -> Result<Self> {
16 Ok(Self {
17 inner: DelimitedReasoningParser::new(
@@ -24,15 +24,15 @@ impl HyV3ReasoningParser {
24 }
25 }
26
27 -impl ReasoningParser for HyV3ReasoningParser {
28 - // Suffix discovery belongs to `HyV3UnifiedParser` so its reasoning and
27 +impl ReasoningParser for HyReasoningParser {
28 + // Suffix discovery belongs to the unified HY parser so its reasoning and
29 // tool delimiters always use the same tokenizer-derived value.
30 fn create(_tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
31 where
32 Self: Sized + 'static,
33 {
34 Err(ReasoningError::DummyUnifiedParser {
35 - name: "hy_v3".to_string(),
35 + name: "hy".to_string(),
36 })
37 }
38
rust/src/parser/src/reasoning/mod.rs
+2 -2
@@ -20,7 +20,7 @@
20 mod cohere_cmd;
21 mod deepseek_r1;
22 mod delimited;
23 -mod hy_v3;
23 +mod hy;
24 mod kimi;
25 mod minimax_m3;
26 mod qwen3;
@@ -33,7 +33,7 @@ use vllm_tokenizer::DynTokenizer;
33 pub use self::cohere_cmd::CohereCmdReasoningParser;
34 pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
35 pub(crate) use self::delimited::{DelimitedReasoningParser, last_reasoning_boundary};
36 -pub(crate) use self::hy_v3::HyV3ReasoningParser;
36 +pub(crate) use self::hy::HyReasoningParser;
37 pub use self::kimi::KimiReasoningParser;
38 pub use self::minimax_m3::MiniMaxM3ReasoningParser;
39 pub use self::qwen3::Qwen3ReasoningParser;
rust/src/parser/src/tool/hy.rs renamed
+250 -118
@@ -11,33 +11,56 @@ use winnow::prelude::*;
11 use winnow::stream::Partial;
12 use winnow::token::{literal, rest, take_until};
13
14 -use self::structural_tag::HyV3StructuralTagBuilder;
14 +use self::structural_tag::HyStructuralTagBuilder;
15 use super::parameters::ToolSchemas;
16 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
17 use super::{Result, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput};
18 use crate::tool::{StructuralTagBuilder, Tool};
19
20 +/// Wire-level HY tool-call dialect.
21 +#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22 +pub(crate) enum HyDialect {
23 + V3,
24 + V4,
25 +}
26 +
27 +impl HyDialect {
28 + fn separator(self) -> &'static str {
29 + match self {
30 + Self::V3 => "\n",
31 + Self::V4 => "",
32 + }
33 + }
34 +
35 + fn label(self) -> &'static str {
36 + match self {
37 + Self::V3 => "HY3",
38 + Self::V4 => "HY4",
39 + }
40 + }
41 +}
42 +
43 #[derive(Debug, Clone, PartialEq, Eq)]
21 -pub(crate) struct HyV3ToolMarkers {
44 +pub(crate) struct HyToolMarkers {
45 tool_calls_start: String,
46 tool_calls_end: String,
47 tool_call_start: String,
48 tool_call_end: String,
26 - tool_sep: String,
49 + tool_sep: Option<String>,
50 arg_key_start: String,
51 arg_key_end: String,
52 arg_value_start: String,
53 arg_value_end: String,
54 }
55
33 -impl HyV3ToolMarkers {
34 - pub(crate) fn new(suffix: &str) -> Self {
56 +impl HyToolMarkers {
57 + pub(crate) fn new(suffix: &str, dialect: HyDialect) -> Self {
58 Self {
59 tool_calls_start: format!("<tool_calls{suffix}>"),
60 tool_calls_end: format!("</tool_calls{suffix}>"),
61 tool_call_start: format!("<tool_call{suffix}>"),
62 tool_call_end: format!("</tool_call{suffix}>"),
40 - tool_sep: format!("<tool_sep{suffix}>"),
63 + tool_sep: (dialect == HyDialect::V3).then(|| format!("<tool_sep{suffix}>")),
64 arg_key_start: format!("<arg_key{suffix}>"),
65 arg_key_end: format!("</arg_key{suffix}>"),
66 arg_value_start: format!("<arg_value{suffix}>"),
@@ -47,31 +70,32 @@ impl HyV3ToolMarkers {
70
71 pub(crate) fn iter(&self) -> impl Iterator<Item = &str> {
72 [
50 - self.tool_calls_start.as_str(),
51 - self.tool_calls_end.as_str(),
52 - self.tool_call_start.as_str(),
53 - self.tool_call_end.as_str(),
54 - self.tool_sep.as_str(),
55 - self.arg_key_start.as_str(),
56 - self.arg_key_end.as_str(),
57 - self.arg_value_start.as_str(),
58 - self.arg_value_end.as_str(),
73 + Some(self.tool_calls_start.as_str()),
74 + Some(self.tool_calls_end.as_str()),
75 + Some(self.tool_call_start.as_str()),
76 + Some(self.tool_call_end.as_str()),
77 + self.tool_sep.as_deref(),
78 + Some(self.arg_key_start.as_str()),
79 + Some(self.arg_key_end.as_str()),
80 + Some(self.arg_value_start.as_str()),
81 + Some(self.arg_value_end.as_str()),
82 ]
83 .into_iter()
84 + .flatten()
85 }
86 }
87
64 -type HyV3Input<'i> = Partial<&'i str>;
88 +type HyInput<'i> = Partial<&'i str>;
89
90 #[derive(Debug, Clone, PartialEq, Eq)]
67 -enum HyV3Mode {
91 +enum HyMode {
92 Text,
93 ToolBlock { tool_call_end_scan: MarkerScanState },
94 Done,
95 }
96
97 #[derive(Debug, Clone, PartialEq, Eq)]
74 -enum HyV3Event {
98 +enum HyEvent {
99 Text {
100 len: usize,
101 },
@@ -84,7 +108,7 @@ enum HyV3Event {
108 IgnoredRest,
109 }
110
87 -/// Tool parser for HY3 XML-style tool calls.
111 +/// Tool parser for HY XML-style tool calls.
112 ///
113 /// Example tool call content:
114 ///
@@ -97,43 +121,45 @@ enum HyV3Event {
121 /// ```
122 ///
123 /// Arguments are emitted only after a full `<tool_call>` block is parsed.
100 -/// HY3 marker tokens are added-vocabulary tokens rather than tokenizer special
124 +/// HY marker tokens are added-vocabulary tokens rather than tokenizer special
125 /// tokens, so the default `preserve_special_tokens() == false` is sufficient.
102 -pub(crate) struct HyV3ToolParser {
126 +pub(crate) struct HyToolParser {
127 buffer: String,
104 - mode: HyV3Mode,
128 + mode: HyMode,
129 emitted_tool_count: usize,
130 tool_parameters: ToolSchemas,
107 - markers: Arc<HyV3ToolMarkers>,
108 - structural_tag_builder: HyV3StructuralTagBuilder,
131 + dialect: HyDialect,
132 + markers: Arc<HyToolMarkers>,
133 + structural_tag_builder: HyStructuralTagBuilder,
134 }
135
111 -impl HyV3ToolParser {
112 - /// Create a HY3 tool parser.
113 - pub(crate) fn new(tools: &[Tool], suffix: &str) -> Self {
114 - let markers = Arc::new(HyV3ToolMarkers::new(suffix));
136 +impl HyToolParser {
137 + /// Create a HY tool parser for one wire dialect.
138 + pub(crate) fn new(tools: &[Tool], suffix: &str, dialect: HyDialect) -> Self {
139 + let markers = Arc::new(HyToolMarkers::new(suffix, dialect));
140 Self {
141 buffer: String::new(),
117 - mode: HyV3Mode::Text,
142 + mode: HyMode::Text,
143 emitted_tool_count: 0,
144 tool_parameters: ToolSchemas::from_tools(tools),
145 + dialect,
146 markers: Arc::clone(&markers),
121 - structural_tag_builder: HyV3StructuralTagBuilder::new(markers),
147 + structural_tag_builder: HyStructuralTagBuilder::new(markers, dialect),
148 }
149 }
150
125 - /// Apply one parsed HY3 event to parser state and output.
126 - fn apply_event(&mut self, event: HyV3Event, output: &mut ToolParserOutput) -> Result<()> {
151 + /// Apply one parsed HY event to parser state and output.
152 + fn apply_event(&mut self, event: HyEvent, output: &mut ToolParserOutput) -> Result<()> {
153 match event {
128 - HyV3Event::Text { len: consumed_len } => {
154 + HyEvent::Text { len: consumed_len } => {
155 output.push_text(&self.buffer[..consumed_len]);
156 }
131 - HyV3Event::ToolBlockStart => {
132 - self.mode = HyV3Mode::ToolBlock {
157 + HyEvent::ToolBlockStart => {
158 + self.mode = HyMode::ToolBlock {
159 tool_call_end_scan: MarkerScanState::default(),
160 };
161 }
136 - HyV3Event::ToolCall { name, raw_params } => {
162 + HyEvent::ToolCall { name, raw_params } => {
163 let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
164 let arguments = serde_json::to_string(&arguments)
165 .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
@@ -145,22 +171,22 @@ impl HyV3ToolParser {
171 });
172 self.emitted_tool_count += 1;
173 }
148 - HyV3Event::ToolBlockEnd => self.mode = HyV3Mode::Done,
149 - HyV3Event::IgnoredRest => {}
174 + HyEvent::ToolBlockEnd => self.mode = HyMode::Done,
175 + HyEvent::IgnoredRest => {}
176 }
177 Ok(())
178 }
179 }
180
155 -impl ToolParser for HyV3ToolParser {
156 - // Suffix discovery belongs to `HyV3UnifiedParser` so its reasoning and
181 +impl ToolParser for HyToolParser {
182 + // Suffix discovery belongs to the unified HY parser so its reasoning and
183 // tool delimiters always use the same tokenizer-derived value.
184 fn create(_tools: &[Tool]) -> Result<Box<dyn ToolParser>>
185 where
186 Self: Sized + 'static,
187 {
188 Err(ToolParserError::DummyUnifiedParser {
163 - name: "hy_v3".to_string(),
189 + name: "hy".to_string(),
190 })
191 }
192
@@ -172,7 +198,7 @@ impl ToolParser for HyV3ToolParser {
198 self.buffer.push_str(chunk);
199
200 while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
175 - parse_next_hy_v3_event(input, &mut self.mode, &self.markers)
201 + parse_next_hy_event(input, &mut self.mode, &self.markers, self.dialect)
202 })? {
203 self.apply_event(event, output)?;
204 self.buffer.drain(..consumed_len);
@@ -184,113 +210,145 @@ impl ToolParser for HyV3ToolParser {
210 fn finish(&mut self) -> Result<ToolParserOutput> {
211 let mut output = ToolParserOutput::default();
212 match self.mode {
187 - HyV3Mode::Text => output.push_text(&self.buffer),
188 - HyV3Mode::ToolBlock { .. } => return Err(parsing_failed!("incomplete HY3 tool call")),
189 - HyV3Mode::Done => {}
213 + HyMode::Text => output.push_text(&self.buffer),
214 + HyMode::ToolBlock { .. } => {
215 + return Err(parsing_failed!(
216 + "incomplete {} tool call",
217 + self.dialect.label()
218 + ));
219 + }
220 + HyMode::Done => {}
221 }
222 let _ = self.reset();
223 Ok(output)
224 }
225
226 fn reset(&mut self) -> String {
196 - self.mode = HyV3Mode::Text;
227 + self.mode = HyMode::Text;
228 self.emitted_tool_count = 0;
229 std::mem::take(&mut self.buffer)
230 }
231 }
232
202 -/// Parse a HY3 event for the current parser mode.
203 -fn parse_next_hy_v3_event(
204 - input: &mut HyV3Input<'_>,
205 - mode: &mut HyV3Mode,
206 - markers: &HyV3ToolMarkers,
207 -) -> ModalResult<HyV3Event> {
233 +/// Parse a HY event for the current parser mode.
234 +fn parse_next_hy_event(
235 + input: &mut HyInput<'_>,
236 + mode: &mut HyMode,
237 + markers: &HyToolMarkers,
238 + dialect: HyDialect,
239 +) -> ModalResult<HyEvent> {
240 match mode {
209 - HyV3Mode::Text => parse_text_event(input, markers),
210 - HyV3Mode::ToolBlock { tool_call_end_scan } => {
211 - parse_tool_block_event(input, tool_call_end_scan, markers)
241 + HyMode::Text => parse_text_event(input, markers),
242 + HyMode::ToolBlock { tool_call_end_scan } => {
243 + parse_tool_block_event(input, tool_call_end_scan, markers, dialect)
244 }
213 - HyV3Mode::Done => ignored_rest_event(input),
245 + HyMode::Done => ignored_rest_event(input),
246 }
247 }
248
217 -/// Parse a text-mode HY3 event.
218 -fn parse_text_event(
219 - input: &mut HyV3Input<'_>,
220 - markers: &HyV3ToolMarkers,
221 -) -> ModalResult<HyV3Event> {
249 +/// Parse a text-mode HY event.
250 +fn parse_text_event(input: &mut HyInput<'_>, markers: &HyToolMarkers) -> ModalResult<HyEvent> {
251 alt((
223 - |input: &mut HyV3Input<'_>| tool_block_start_event(input, markers),
224 - |input: &mut HyV3Input<'_>| safe_text_event(input, markers),
252 + |input: &mut HyInput<'_>| tool_block_start_event(input, markers),
253 + |input: &mut HyInput<'_>| safe_text_event(input, markers),
254 ))
255 .parse_next(input)
256 }
257
229 -/// Parse a HY3 tool-block start marker.
258 +/// Parse a HY tool-block start marker.
259 fn tool_block_start_event(
231 - input: &mut HyV3Input<'_>,
232 - markers: &HyV3ToolMarkers,
233 -) -> ModalResult<HyV3Event> {
260 + input: &mut HyInput<'_>,
261 + markers: &HyToolMarkers,
262 +) -> ModalResult<HyEvent> {
263 literal(markers.tool_calls_start.as_str())
235 - .value(HyV3Event::ToolBlockStart)
264 + .value(HyEvent::ToolBlockStart)
265 .parse_next(input)
266 }
267
239 -/// Parse a safe text run before the next HY3 marker.
240 -fn safe_text_event(input: &mut HyV3Input<'_>, markers: &HyV3ToolMarkers) -> ModalResult<HyV3Event> {
241 - safe_text_len(input, &markers.tool_calls_start).map(|len| HyV3Event::Text { len })
268 +/// Parse a safe text run before the next HY marker.
269 +fn safe_text_event(input: &mut HyInput<'_>, markers: &HyToolMarkers) -> ModalResult<HyEvent> {
270 + safe_text_len(input, &markers.tool_calls_start).map(|len| HyEvent::Text { len })
271 }
272
244 -/// Parse one event inside a HY3 tool block.
273 +/// Parse one event inside a HY tool block.
274 fn parse_tool_block_event(
246 - input: &mut HyV3Input<'_>,
275 + input: &mut HyInput<'_>,
276 tool_call_end_scan: &mut MarkerScanState,
248 - markers: &HyV3ToolMarkers,
249 -) -> ModalResult<HyV3Event> {
277 + markers: &HyToolMarkers,
278 + dialect: HyDialect,
279 +) -> ModalResult<HyEvent> {
280 alt((
251 - |input: &mut HyV3Input<'_>| tool_block_end_event(input, markers),
252 - |input: &mut HyV3Input<'_>| tool_call_event(input, tool_call_end_scan, markers),
281 + |input: &mut HyInput<'_>| tool_block_end_event(input, markers),
282 + |input: &mut HyInput<'_>| tool_call_event(input, tool_call_end_scan, markers, dialect),
283 ))
284 .parse_next(input)
285 }
286
257 -/// Parse a HY3 tool-block end marker.
258 -fn tool_block_end_event(
259 - input: &mut HyV3Input<'_>,
260 - markers: &HyV3ToolMarkers,
261 -) -> ModalResult<HyV3Event> {
287 +/// Parse a HY tool-block end marker.
288 +fn tool_block_end_event(input: &mut HyInput<'_>, markers: &HyToolMarkers) -> ModalResult<HyEvent> {
289 (ws0, literal(markers.tool_calls_end.as_str()))
263 - .value(HyV3Event::ToolBlockEnd)
290 + .value(HyEvent::ToolBlockEnd)
291 .parse_next(input)
292 }
293
267 -/// Parse a complete HY3 tool-call block.
294 +/// Parse a complete HY tool-call block.
295 fn tool_call_event(
269 - input: &mut HyV3Input<'_>,
296 + input: &mut HyInput<'_>,
297 tool_call_end_scan: &mut MarkerScanState,
271 - markers: &HyV3ToolMarkers,
272 -) -> ModalResult<HyV3Event> {
273 - let (name, body) = seq!(
298 + markers: &HyToolMarkers,
299 + dialect: HyDialect,
300 +) -> ModalResult<HyEvent> {
301 + let (body,) = seq!(
302 _: ws0,
303 _: literal(markers.tool_call_start.as_str()),
276 - take_until(0.., markers.tool_sep.as_str()),
277 - _: literal(markers.tool_sep.as_str()),
304 take_until_marker(markers.tool_call_end.as_str(), tool_call_end_scan),
305 _: literal(markers.tool_call_end.as_str()),
306 )
307 .parse_next(input)?;
282 - let raw_params = parse_tool_call_params(body, markers)?;
308 + let mut body_input = body;
309 + let (name, params) = parse_tool_call_body(&mut body_input, markers, dialect)?;
310 + let raw_params = parse_tool_call_params(params, markers)?;
311
284 - Ok(HyV3Event::ToolCall {
312 + Ok(HyEvent::ToolCall {
313 name: name.trim().to_string(),
314 raw_params,
315 })
316 }
317
290 -/// Parse all parameter blocks inside a complete HY3 tool call.
318 +/// Parse a complete HY tool-call body according to its dialect.
319 +fn parse_tool_call_body<'i>(
320 + input: &mut &'i str,
321 + markers: &HyToolMarkers,
322 + dialect: HyDialect,
323 +) -> ModalResult<(&'i str, &'i str)> {
324 + match dialect {
325 + HyDialect::V3 => {
326 + let tool_sep = markers.tool_sep.as_deref().expect("HY3 has a tool separator");
327 + terminated(
328 + seq!(
329 + take_until(0.., tool_sep),
330 + _: literal(tool_sep),
331 + rest,
332 + ),
333 + eof,
334 + )
335 + .parse_next(input)
336 + }
337 + HyDialect::V4 => terminated(
338 + alt((
339 + seq!(take_until(0.., markers.arg_key_start.as_str()), rest,),
340 + rest.map(|name| (name, "")),
341 + )),
342 + eof,
343 + )
344 + .parse_next(input),
345 + }
346 +}
347 +
348 +/// Parse all parameter blocks inside a complete HY tool call.
349 fn parse_tool_call_params(
350 tool_call_body: &str,
293 - markers: &HyV3ToolMarkers,
351 + markers: &HyToolMarkers,
352 ) -> ModalResult<Vec<(String, String)>> {
353 let mut input = tool_call_body;
354 delimited(
@@ -304,8 +362,8 @@ fn parse_tool_call_params(
362 .parse_next(&mut input)
363 }
364
307 -/// Parse a HY3 argument key/value block.
308 -fn parameter(input: &mut &str, markers: &HyV3ToolMarkers) -> ModalResult<(String, String)> {
365 +/// Parse a HY argument key/value block.
366 +fn parameter(input: &mut &str, markers: &HyToolMarkers) -> ModalResult<(String, String)> {
367 let (name, value) = seq!(
368 _: literal(markers.arg_key_start.as_str()),
369 take_until(0.., markers.arg_key_end.as_str()),
@@ -320,9 +378,9 @@ fn parameter(input: &mut &str, markers: &HyV3ToolMarkers) -> ModalResult<(String
378 Ok((name.trim().to_string(), value.to_string()))
379 }
380
323 -/// Parse ignored rest after the HY3 tool block ends.
324 -fn ignored_rest_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
325 - rest.value(HyV3Event::IgnoredRest).parse_next(input)
381 +/// Parse ignored rest after the HY tool block ends.
382 +fn ignored_rest_event(input: &mut HyInput<'_>) -> ModalResult<HyEvent> {
383 + rest.value(HyEvent::IgnoredRest).parse_next(input)
384 }
385
386 #[cfg(test)]
@@ -331,7 +389,7 @@ mod tests {
389 use serde_json::{Value, json};
390 use thiserror_ext::AsReport;
391
334 - use super::{HyV3ToolParser, ToolParser};
392 + use super::{HyDialect, HyToolParser, ToolParser};
393 use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
394 use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
395
@@ -354,14 +412,14 @@ mod tests {
412
413 #[test]
414 fn hy_v3_does_not_preserve_special_tokens() {
357 - let parser = HyV3ToolParser::new(&test_tools(), "");
415 + let parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
416
417 assert!(!parser.preserve_special_tokens());
418 }
419
420 #[test]
421 fn hy_v3_parse_complete_without_tool_call_keeps_text() {
364 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
422 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
423 let output = parser.parse_complete("This is a plain response.").unwrap();
424
425 assert_eq!(output.normal_text(), "This is a plain response.");
@@ -370,7 +428,7 @@ mod tests {
428
429 #[test]
430 fn hy_v3_parse_complete_extracts_zero_arg_inline_tool_call() {
373 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
431 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
432 let output = parser
433 .parse_complete(
434 "<tool_calls><tool_call>get_current_date<tool_sep></tool_call></tool_calls>",
@@ -385,7 +443,7 @@ mod tests {
443
444 #[test]
445 fn hy_v3_parse_complete_extracts_zero_arg_newline_tool_call() {
388 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
446 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
447 let output = parser
448 .parse_complete(
449 "<tool_calls>\n<tool_call>get_current_date<tool_sep>\n</tool_call>\n</tool_calls>",
@@ -398,7 +456,7 @@ mod tests {
456
457 #[test]
458 fn hy_v3_parse_complete_extracts_arguments_on_same_line() {
401 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
459 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
460 let output = parser
461 .parse_complete(
462 "<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</arg_value><arg_key>date</arg_key><arg_value>2026-03-30</arg_value></tool_call></tool_calls>",
@@ -413,7 +471,7 @@ mod tests {
471
472 #[test]
473 fn hy_v3_parse_complete_extracts_arguments_with_newlines() {
416 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
474 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
475 let output = parser
476 .parse_complete(&build_tool_calls(&[build_tool_call(
477 "get_weather",
@@ -429,7 +487,7 @@ mod tests {
487
488 #[test]
489 fn hy_v3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
432 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
490 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
491 let output = parser
492 .parse_complete(&format!(
493 "Checking.{} trailing text",
@@ -443,7 +501,7 @@ mod tests {
501
502 #[test]
503 fn hy_v3_parse_complete_extracts_multiple_tool_calls_in_one_block() {
446 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
504 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
505 let output = parser
506 .parse_complete(&build_tool_calls(&[
507 build_tool_call(
@@ -486,7 +544,7 @@ mod tests {
544
545 #[test]
546 fn hy_v3_parse_complete_converts_schema_types() {
489 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
547 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
548 let output = parser
549 .parse_complete(&build_tool_calls(&[build_tool_call(
550 "convert",
@@ -512,7 +570,7 @@ mod tests {
570
571 #[test]
572 fn hy_v3_streaming_without_tool_call_emits_text_incrementally() {
515 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
573 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
574 let mut output = ToolParserOutput::default();
575
576 output.append(parser.parse_chunk("This is ").unwrap());
@@ -526,7 +584,7 @@ mod tests {
584
585 #[test]
586 fn hy_v3_streaming_extracts_zero_arg_tool_call() {
529 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
587 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
588 let chunks = [
589 "<tool_calls>",
590 "\n<tool_call>",
@@ -545,7 +603,7 @@ mod tests {
603
604 #[test]
605 fn hy_v3_streaming_extracts_arguments() {
548 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
606 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
607 let chunks = [
608 "<tool_calls>",
609 "\n<tool_call>",
@@ -571,7 +629,7 @@ mod tests {
629
630 #[test]
631 fn hy_v3_streaming_preserves_prefix_text() {
574 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
632 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
633 let chunks = [
634 "Checking.",
635 "<tool_calls>",
@@ -601,7 +659,7 @@ mod tests {
659 ),
660 ]);
661 let chunks = split_by_chars(&input, 9);
604 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
662 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
663
664 let output = collect_stream(&mut parser, &chunks);
665
@@ -617,7 +675,7 @@ mod tests {
675 build_tool_calls(&[build_tool_call("get_weather", &[("city", "Beijing")])])
676 );
677 let chunks = split_by_chars(&input, 5);
620 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
678 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
679
680 let output = collect_stream(&mut parser, &chunks);
681
@@ -628,7 +686,7 @@ mod tests {
686
687 #[test]
688 fn hy_v3_streaming_does_not_emit_incomplete_tool_call() {
631 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
689 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
690 let mut output = ToolParserOutput::default();
691
692 parser
@@ -644,7 +702,7 @@ mod tests {
702
703 #[test]
704 fn hy_v3_finish_fails_incomplete_tool_call() {
647 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
705 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
706 parser.parse_chunk("<tool_calls><tool_call>get_weather<tool_sep>").unwrap();
707
708 let error = parser.finish().unwrap_err();
@@ -655,7 +713,7 @@ mod tests {
713
714 #[test]
715 fn hy_v3_malformed_tool_call_fails_fast() {
658 - let mut parser = HyV3ToolParser::new(&test_tools(), "");
716 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V3);
717 let error = parser
718 .parse_complete(
719 "<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</tool_call></tool_calls>",
@@ -664,4 +722,78 @@ mod tests {
722
723 assert!(error.to_report_string().starts_with("tool parser parsing failed:"));
724 }
725 +
726 + #[test]
727 + fn hy_v4_parse_complete_extracts_compact_calls_with_and_without_arguments() {
728 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V4);
729 + let output = parser
730 + .parse_complete(
731 + "<tool_calls><tool_call>get_weather<arg_key>city</arg_key><arg_value>Beijing</arg_value><arg_key>date</arg_key><arg_value>2026-03-30</arg_value></tool_call><tool_call>get_current_date</tool_call></tool_calls>",
732 + )
733 + .unwrap();
734 +
735 + expect![[r#"
736 + ToolParserOutput {
737 + events: [
738 + ToolCall(
739 + ToolCallDelta {
740 + tool_index: 0,
741 + name: Some(
742 + "get_weather",
743 + ),
744 + arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}",
745 + },
746 + ),
747 + ToolCall(
748 + ToolCallDelta {
749 + tool_index: 1,
750 + name: Some(
751 + "get_current_date",
752 + ),
753 + arguments: "{}",
754 + },
755 + ),
756 + ],
757 + }
758 + "#]]
759 + .assert_debug_eq(&output);
760 + }
761 +
762 + #[test]
763 + fn hy_v4_streaming_handles_markers_split_across_chunks() {
764 + let input = "prefix<tool_calls><tool_call>get_weather<arg_key>city</arg_key><arg_value>上海</arg_value></tool_call><tool_call>get_current_date</tool_call></tool_calls>";
765 + let chunks = split_by_chars(input, 7);
766 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V4);
767 +
768 + let output = collect_stream(&mut parser, &chunks);
769 +
770 + assert_eq!(output.normal_text(), "prefix");
771 + assert_eq!(output.calls().len(), 2);
772 + assert_eq!(parsed_arguments(&output, 0), json!({ "city": "上海" }));
773 + assert_eq!(output.calls()[1].name.as_deref(), Some("get_current_date"));
774 + assert_eq!(parsed_arguments(&output, 1), json!({}));
775 + }
776 +
777 + #[test]
778 + fn hy_v4_finish_reports_its_dialect_for_incomplete_calls() {
779 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V4);
780 + parser.parse_chunk("<tool_calls><tool_call>get_weather").unwrap();
781 +
782 + let error = parser.finish().unwrap_err();
783 +
784 + expect!["tool parser parsing failed: incomplete HY4 tool call"]
785 + .assert_eq(&error.to_report_string());
786 + }
787 +
788 + #[test]
789 + fn hy_v4_malformed_argument_pair_fails_fast() {
790 + let mut parser = HyToolParser::new(&test_tools(), "", HyDialect::V4);
791 + let error = parser
792 + .parse_complete(
793 + "<tool_calls><tool_call>get_weather<arg_key>city</arg_key>Beijing</tool_call></tool_calls>",
794 + )
795 + .unwrap_err();
796 +
797 + assert!(error.to_report_string().starts_with("tool parser parsing failed:"));
798 + }
799 }
rust/src/parser/src/tool/hy/structural_tag.rs renamed
+86 -23
@@ -1,7 +1,7 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -//! Structural-tag grammar for HY3 XML-style tool calls.
4 +//! Structural-tag grammar for HY XML-style tool calls.
5
6 use std::collections::HashSet;
7 use std::sync::Arc;
@@ -12,31 +12,40 @@ use xgrammar_structural_tag::builders::{StructuralTagBuilder, StructuralTagConte
12 use xgrammar_structural_tag::format::{Format, StructuralTag, TagFormat};
13 use xgrammar_structural_tag::tool::{BuilderToolChoice, FunctionToolParam};
14
15 -use super::HyV3ToolMarkers;
15 +use super::{HyDialect, HyToolMarkers};
16
17 -/// HY3 structural-tag builder using tokenizer-specific structural markers.
17 +/// HY structural-tag builder using tokenizer-specific structural markers.
18 #[derive(Debug, Clone)]
19 -pub(super) struct HyV3StructuralTagBuilder {
20 - markers: Arc<HyV3ToolMarkers>,
19 +pub(super) struct HyStructuralTagBuilder {
20 + markers: Arc<HyToolMarkers>,
21 + dialect: HyDialect,
22 }
23
23 -impl HyV3StructuralTagBuilder {
24 - pub(super) fn new(markers: Arc<HyV3ToolMarkers>) -> Self {
25 - Self { markers }
24 +impl HyStructuralTagBuilder {
25 + pub(super) fn new(markers: Arc<HyToolMarkers>, dialect: HyDialect) -> Self {
26 + Self { markers, dialect }
27 }
28
29 fn argument_pair(&self, key: &str) -> Format {
30 let excludes = self.markers.iter().collect::<Vec<_>>();
30 - Format::sequence(vec![
31 + let separator = self.dialect.separator();
32 + let mut elements = vec![
33 Format::const_string(&self.markers.arg_key_start),
34 Format::const_string(key),
35 Format::const_string(&self.markers.arg_key_end),
34 - Format::const_string("\n"),
36 + ];
37 + if !separator.is_empty() {
38 + elements.push(Format::const_string(separator));
39 + }
40 + elements.extend([
41 Format::const_string(&self.markers.arg_value_start),
42 Format::any_text_excluding(&excludes),
43 Format::const_string(&self.markers.arg_value_end),
38 - Format::const_string("\n"),
39 - ])
44 + ]);
45 + if !separator.is_empty() {
46 + elements.push(Format::const_string(separator));
47 + }
48 + Format::sequence(elements)
49 }
50
51 fn tool_call(&self, tool: &FunctionToolParam) -> TagFormat {
@@ -56,14 +65,21 @@ impl HyV3StructuralTagBuilder {
65 }
66
67 let content = if elements.is_empty() {
59 - Format::any_text()
68 + match self.dialect {
69 + HyDialect::V3 => Format::any_text(),
70 + HyDialect::V4 => Format::const_string(""),
71 + }
72 } else {
73 Format::sequence(elements)
74 };
75 + let tool_sep = self.markers.tool_sep.as_deref().unwrap_or_default();
76 TagFormat::new(
77 format!(
65 - "{}{}{}\n",
66 - self.markers.tool_call_start, tool.function.name, self.markers.tool_sep
78 + "{}{}{}{}",
79 + self.markers.tool_call_start,
80 + tool.function.name,
81 + tool_sep,
82 + self.dialect.separator(),
83 ),
84 content,
85 self.markers.tool_call_end.clone(),
@@ -72,15 +88,16 @@ impl HyV3StructuralTagBuilder {
88
89 fn tool_calls(&self, tools: &[FunctionToolParam], choice: BuilderToolChoice) -> Format {
90 let mut calls = tools.iter().map(|tool| self.tool_call(tool)).collect::<Vec<_>>();
75 - let begin = format!("{}\n", self.markers.tool_calls_start);
76 - let end = format!("\n{}", self.markers.tool_calls_end);
91 + let separator = self.dialect.separator();
92 + let begin = format!("{}{separator}", self.markers.tool_calls_start);
93 + let end = format!("{separator}{}", self.markers.tool_calls_end);
94
95 match choice {
96 BuilderToolChoice::Auto if calls.is_empty() => Format::any_text(),
97 BuilderToolChoice::Auto => {
98 let outer = TagFormat::new(
99 begin,
83 - Format::tags_with_separator(calls, "\n", true, false),
100 + Format::tags_with_separator(calls, separator, true, false),
101 end,
102 );
103 Format::triggered_tags(&[&self.markers.tool_calls_start], vec![outer])
@@ -92,14 +109,14 @@ impl HyV3StructuralTagBuilder {
109 ]),
110 BuilderToolChoice::Required => Format::sequence(vec![
111 Format::const_string(begin),
95 - Format::tags_with_separator(calls, "\n", true, false),
112 + Format::tags_with_separator(calls, separator, true, false),
113 Format::const_string(end),
114 ]),
115 }
116 }
117 }
118
102 -impl StructuralTagBuilder for HyV3StructuralTagBuilder {
119 +impl StructuralTagBuilder for HyStructuralTagBuilder {
120 fn build(&self, ctx: StructuralTagContext<'_>) -> Result<StructuralTag> {
121 Ok(StructuralTag::new(
122 self.tool_calls(ctx.function_tools, ctx.tool_choice),
@@ -152,8 +169,8 @@ mod tests {
169 FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
170 };
171
155 - use super::HyV3StructuralTagBuilder;
156 - use crate::tool::HyV3ToolMarkers;
172 + use super::HyStructuralTagBuilder;
173 + use crate::tool::{HyDialect, HyToolMarkers};
174
175 fn tool(name: &str, parameters: Value) -> ToolParam {
176 ToolParam::Function(FunctionToolParam::new(
@@ -165,9 +182,18 @@ mod tests {
182 suffix: &str,
183 tools: &[ToolParam],
184 choice: ToolChoice,
185 + ) -> xgrammar_structural_tag::format::StructuralTag {
186 + build_for_dialect(HyDialect::V3, suffix, tools, choice)
187 + }
188 +
189 + fn build_for_dialect(
190 + dialect: HyDialect,
191 + suffix: &str,
192 + tools: &[ToolParam],
193 + choice: ToolChoice,
194 ) -> xgrammar_structural_tag::format::StructuralTag {
195 build_structural_tag(
170 - HyV3StructuralTagBuilder::new(Arc::new(HyV3ToolMarkers::new(suffix))),
196 + HyStructuralTagBuilder::new(Arc::new(HyToolMarkers::new(suffix, dialect)), dialect),
197 tools,
198 choice,
199 StructuralTagOptions::default().with_reasoning(false),
@@ -196,6 +222,25 @@ mod tests {
222 expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>\n"},{"type":"tags_with_separator","tags":[{"begin":"<tool_call:opensource>get_weather<tool_sep:opensource>\n","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]},{"type":"star","content":{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"days"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"\n"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<tool_sep:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"},{"type":"const_string","value":"\n"}]}}]},"end":"</tool_call:opensource>"}],"separator":"\n","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"\n</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
223 }
224
225 + #[test]
226 + fn required_uses_suffixed_compact_hy4_skeleton_and_bounded_values() {
227 + let tag = build_for_dialect(
228 + HyDialect::V4,
229 + ":opensource",
230 + &[tool(
231 + "get_weather",
232 + json!({
233 + "type": "object",
234 + "properties": { "city": { "type": "string" } },
235 + "required": ["city"]
236 + }),
237 + )],
238 + ToolChoice::required(),
239 + );
240 +
241 + expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls:opensource>"},{"type":"tags_with_separator","tags":[{"begin":"<tool_call:opensource>get_weather","content":{"type":"sequence","elements":[{"type":"sequence","elements":[{"type":"const_string","value":"<arg_key:opensource>"},{"type":"const_string","value":"city"},{"type":"const_string","value":"</arg_key:opensource>"},{"type":"const_string","value":"<arg_value:opensource>"},{"type":"any_text","excludes":["<tool_calls:opensource>","</tool_calls:opensource>","<tool_call:opensource>","</tool_call:opensource>","<arg_key:opensource>","</arg_key:opensource>","<arg_value:opensource>","</arg_value:opensource>"]},{"type":"const_string","value":"</arg_value:opensource>"}]}]},"end":"</tool_call:opensource>"}],"separator":"","at_least_one":true,"stop_after_first":false},{"type":"const_string","value":"</tool_calls:opensource>"}]}}"#]].assert_eq(&tag.to_json_string().unwrap());
242 + }
243 +
244 #[test]
245 fn optional_only_schema_keeps_declared_key_alternatives() {
246 let tag = build(
@@ -235,6 +280,24 @@ mod tests {
280 expect![[r#"{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"const_string","value":"<tool_calls>\n"},{"type":"tag","begin":"<tool_call>lookup<tool_sep>\n","content":{"type":"any_text","excludes":[]},"end":"</tool_call>"},{"type":"const_string","value":"\n</tool_calls>"}]}}"#]].assert_eq(&forced);
281 }
282
283 + #[test]
284 + fn hy_v4_zero_argument_tool_has_empty_compact_body() {
285 + let tag = build_for_dialect(
286 + HyDialect::V4,
287 + "",
288 + &[tool(
289 + "get_current_date",
290 + json!({ "type": "object", "properties": {} }),
291 + )],
292 + ToolChoice::function("get_current_date"),
293 + );
294 + let value = serde_json::to_value(tag).unwrap();
295 +
296 + expect![[r#"{"type":"tag","begin":"<tool_call>get_current_date","content":{"type":"const_string","value":""},"end":"</tool_call>"}"#]].assert_eq(
297 + &serde_json::to_string(&value["format"]["elements"][1]).unwrap(),
298 + );
299 + }
300 +
301 #[test]
302 fn missing_required_property_is_still_emitted() {
303 let tag = build(
rust/src/parser/src/tool/mod.rs
+2 -2
@@ -8,7 +8,7 @@ pub(crate) mod error;
8 mod deepseek_dsml;
9 pub(crate) mod deepseek_json;
10 mod glm_xml;
11 -mod hy_v3;
11 +mod hy;
12 pub(crate) mod json;
13 mod kimi_k2;
14 mod minimax_m2;
@@ -24,7 +24,7 @@ pub use deepseek_dsml::{DeepSeekV4ToolParser, DeepSeekV32ToolParser};
24 pub use deepseek_json::{DeepSeekV3ToolParser, DeepSeekV31ToolParser};
25 pub use error::{Result, ToolParserError};
26 pub use glm_xml::{Glm45MoeToolParser, Glm47MoeToolParser};
27 -pub(crate) use hy_v3::{HyV3ToolMarkers, HyV3ToolParser};
27 +pub(crate) use hy::{HyDialect, HyToolMarkers, HyToolParser};
28 pub use json::{
29 Granite4ToolParser, HermesToolParser, Internlm2ToolParser, Llama3JsonToolParser,
30 MistralToolParser, Phi4MiniJsonToolParser, Qwen3XmlToolParser,
rust/src/parser/src/unified/hy/hy_v3.rs renamed
+9 -39
@@ -1,20 +1,12 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -use vllm_tokenizer::{DynTokenizer, Tokenizer};
4 +use vllm_tokenizer::DynTokenizer;
5
6 -use super::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id};
7 -use crate::reasoning::HyV3ReasoningParser;
8 -use crate::tool::{HyV3ToolMarkers, HyV3ToolParser, StructuralTagBuilder, Tool};
9 -
10 -const HY_V3_MARKER_STEMS: &[&str] = &[
11 - "think",
12 - "tool_calls",
13 - "tool_call",
14 - "tool_sep",
15 - "arg_key",
16 - "arg_value",
17 -];
6 +use super::detect_hy_token_suffix;
7 +use crate::reasoning::HyReasoningParser;
8 +use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, StructuralTagBuilder, Tool};
9 +use crate::unified::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id};
10
11 /// Unified reasoning and tool parser for HY3 output.
12 pub struct HyV3UnifiedParser {
@@ -24,14 +16,14 @@ pub struct HyV3UnifiedParser {
16 impl HyV3UnifiedParser {
17 /// Create a HY3 parser using the suffix encoded in tokenizer added tokens.
18 pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
27 - let suffix = detect_token_suffix(tokenizer.as_ref());
28 - let markers = HyV3ToolMarkers::new(&suffix);
19 + let suffix = detect_hy_token_suffix(tokenizer.as_ref());
20 + let markers = HyToolMarkers::new(&suffix, HyDialect::V3);
21 for marker in markers.iter() {
22 token_id(tokenizer.as_ref(), marker)?;
23 }
24
33 - let reasoning = HyV3ReasoningParser::new(tokenizer, &suffix)?;
34 - let tool = HyV3ToolParser::new(tools, &suffix);
25 + let reasoning = HyReasoningParser::new(tokenizer, &suffix)?;
26 + let tool = HyToolParser::new(tools, &suffix, HyDialect::V3);
27 Ok(Self {
28 inner: CombinedParser::new(Some(Box::new(reasoning)), Some(Box::new(tool))),
29 })
@@ -75,28 +67,6 @@ impl UnifiedParser for HyV3UnifiedParser {
67 }
68 }
69
78 -/// Detect the HY3 structural-token suffix from tokenizer added vocabulary.
79 -fn detect_token_suffix(tokenizer: &dyn Tokenizer) -> String {
80 - tokenizer
81 - .added_vocab()
82 - .iter()
83 - .filter_map(|(token, id)| marker_suffix(token).map(|suffix| (*id, suffix)))
84 - .min_by_key(|(id, _)| *id)
85 - .map(|(_, suffix)| suffix.to_string())
86 - .unwrap_or_default()
87 -}
88 -
89 -/// Extract the suffix from one opening or closing HY3 structural token.
90 -fn marker_suffix(token: &str) -> Option<&str> {
91 - let body = token.strip_prefix('<')?.strip_suffix('>')?;
92 - let body = body.strip_prefix('/').unwrap_or(body);
93 -
94 - HY_V3_MARKER_STEMS.iter().find_map(|stem| {
95 - let suffix = body.strip_prefix(stem)?;
96 - (suffix.is_empty() || suffix.starts_with(':')).then_some(suffix)
97 - })
98 -}
99 -
70 #[cfg(test)]
71 mod tests {
72 use std::sync::Arc;
rust/src/parser/src/unified/hy/hy_v4.rs new
+194
@@ -0,0 +1,194 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use vllm_tokenizer::DynTokenizer;
5 +
6 +use super::detect_hy_token_suffix;
7 +use crate::reasoning::HyReasoningParser;
8 +use crate::tool::{HyDialect, HyToolMarkers, HyToolParser, StructuralTagBuilder, Tool};
9 +use crate::unified::{CombinedParser, Result, UnifiedParser, UnifiedParserOutput, token_id};
10 +
11 +/// Unified reasoning and tool parser for HY4 output.
12 +pub struct HyV4UnifiedParser {
13 + inner: CombinedParser,
14 +}
15 +
16 +impl HyV4UnifiedParser {
17 + /// Create a HY4 parser using the suffix encoded in tokenizer added tokens.
18 + pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Self> {
19 + let suffix = detect_hy_token_suffix(tokenizer.as_ref());
20 + let markers = HyToolMarkers::new(&suffix, HyDialect::V4);
21 + for marker in markers.iter() {
22 + token_id(tokenizer.as_ref(), marker)?;
23 + }
24 +
25 + let reasoning = HyReasoningParser::new(tokenizer, &suffix)?;
26 + let tool = HyToolParser::new(tools, &suffix, HyDialect::V4);
27 + Ok(Self {
28 + inner: CombinedParser::new(Some(Box::new(reasoning)), Some(Box::new(tool))),
29 + })
30 + }
31 +}
32 +
33 +impl UnifiedParser for HyV4UnifiedParser {
34 + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result<Box<dyn UnifiedParser>>
35 + where
36 + Self: Sized + 'static,
37 + {
38 + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box<dyn UnifiedParser>)
39 + }
40 +
41 + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
42 + self.inner.initialize(prompt_token_ids)
43 + }
44 +
45 + fn preserve_special_tokens(&self) -> bool {
46 + self.inner.preserve_special_tokens()
47 + }
48 +
49 + fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
50 + self.inner.structural_tag_builder()
51 + }
52 +
53 + fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
54 + self.inner.tool_call_id(tool_index)
55 + }
56 +
57 + fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> {
58 + self.inner.parse_into(delta, output)
59 + }
60 +
61 + fn finish(&mut self) -> Result<UnifiedParserOutput> {
62 + self.inner.finish()
63 + }
64 +
65 + fn reset(&mut self) -> String {
66 + self.inner.reset()
67 + }
68 +}
69 +
70 +#[cfg(test)]
71 +mod tests {
72 + use std::sync::Arc;
73 +
74 + use serde_json::json;
75 + use vllm_tokenizer::{Tokenizer, test_utils::TestTokenizer};
76 + use xgrammar_structural_tag::builders::StructuralTagOptions;
77 + use xgrammar_structural_tag::{
78 + FunctionDefinition, FunctionToolParam, ToolChoice, ToolParam, build_structural_tag,
79 + };
80 +
81 + use super::{HyV4UnifiedParser, UnifiedParser};
82 + use crate::tool::Tool;
83 + use crate::unified::{UnifiedParserEvent, UnifiedParserOutput};
84 +
85 + fn tokenizer() -> TestTokenizer {
86 + [
87 + "<think:opensource>",
88 + "</think:opensource>",
89 + "<tool_calls:opensource>",
90 + "</tool_calls:opensource>",
91 + "<tool_call:opensource>",
92 + "</tool_call:opensource>",
93 + "<arg_key:opensource>",
94 + "</arg_key:opensource>",
95 + "<arg_value:opensource>",
96 + "</arg_value:opensource>",
97 + ]
98 + .into_iter()
99 + .enumerate()
100 + .fold(TestTokenizer::new(), |tokenizer, (index, token)| {
101 + tokenizer.with_regular_token(token, 1000 + index as u32)
102 + })
103 + }
104 +
105 + fn tools() -> Vec<Tool> {
106 + vec![
107 + Tool {
108 + name: "get_weather".to_string(),
109 + description: None,
110 + parameters: json!({
111 + "type": "object",
112 + "properties": { "city": { "type": "string" } },
113 + }),
114 + strict: None,
115 + },
116 + Tool {
117 + name: "get_current_date".to_string(),
118 + description: None,
119 + parameters: json!({ "type": "object", "properties": {} }),
120 + strict: None,
121 + },
122 + ]
123 + }
124 +
125 + #[test]
126 + fn parses_suffixed_reasoning_and_compact_tool_calls_as_one_stream() {
127 + let tokenizer = Arc::new(tokenizer());
128 + let think_start_id = tokenizer.token_to_id("<think:opensource>").unwrap();
129 + let mut parser = HyV4UnifiedParser::new(&tools(), tokenizer).unwrap();
130 + parser.initialize(&[think_start_id]).unwrap();
131 +
132 + let chunks = [
133 + "reasoning</think:open",
134 + "source>answer<tool_calls:opensource><tool_call:opensource>get_weather",
135 + "<arg_key:opensource>city</arg_key:opensource>",
136 + "<arg_value:opensource>Beijing</arg_value:opensource>",
137 + "</tool_call:opensource><tool_call:opensource>get_current_date",
138 + "</tool_call:opensource></tool_calls:opensource>",
139 + ];
140 + let mut output = UnifiedParserOutput::default();
141 + for chunk in chunks {
142 + parser.parse_into(chunk, &mut output).unwrap();
143 + }
144 + output.append(parser.finish().unwrap());
145 +
146 + assert_eq!(
147 + output.events,
148 + vec![
149 + UnifiedParserEvent::Reasoning("reasoning".to_string()),
150 + UnifiedParserEvent::Text("answer".to_string()),
151 + UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta {
152 + tool_index: 0,
153 + name: Some("get_weather".to_string()),
154 + arguments: r#"{"city":"Beijing"}"#.to_string(),
155 + }),
156 + UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta {
157 + tool_index: 1,
158 + name: Some("get_current_date".to_string()),
159 + arguments: "{}".to_string(),
160 + }),
161 + ]
162 + );
163 + }
164 +
165 + #[test]
166 + fn structural_tag_uses_compact_hy4_skeleton() {
167 + let parser = HyV4UnifiedParser::new(&tools(), Arc::new(tokenizer())).unwrap();
168 + let structural_tools = [ToolParam::Function(FunctionToolParam::new(
169 + FunctionDefinition::new("get_weather").with_parameters(json!({
170 + "type": "object",
171 + "properties": { "city": { "type": "string" } },
172 + "required": ["city"]
173 + })),
174 + ))];
175 +
176 + let tag = build_structural_tag(
177 + parser.structural_tag_builder().unwrap(),
178 + &structural_tools,
179 + ToolChoice::required(),
180 + StructuralTagOptions::default().with_reasoning(false),
181 + )
182 + .unwrap()
183 + .to_json_string()
184 + .unwrap();
185 +
186 + assert!(tag.contains("<tool_calls:opensource>"));
187 + assert!(tag.contains("<tool_call:opensource>get_weather"));
188 + assert!(tag.contains("<arg_key:opensource>"));
189 + assert!(tag.contains("</tool_call:opensource>"));
190 + assert!(tag.contains("</tool_calls:opensource>"));
191 + assert!(!tag.contains("tool_sep"));
192 + assert!(!tag.contains(r#"\n"#));
193 + }
194 +}
rust/src/parser/src/unified/hy/mod.rs new
+39
@@ -0,0 +1,39 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +mod hy_v3;
5 +mod hy_v4;
6 +
7 +pub use hy_v3::HyV3UnifiedParser;
8 +pub use hy_v4::HyV4UnifiedParser;
9 +
10 +const HY_MARKER_STEMS: &[&str] = &[
11 + "think",
12 + "tool_calls",
13 + "tool_call",
14 + "tool_sep",
15 + "arg_key",
16 + "arg_value",
17 +];
18 +
19 +/// Detect the HY structural-token suffix from tokenizer added vocabulary.
20 +fn detect_hy_token_suffix(tokenizer: &dyn vllm_tokenizer::Tokenizer) -> String {
21 + tokenizer
22 + .added_vocab()
23 + .iter()
24 + .filter_map(|(token, id)| hy_marker_suffix(token).map(|suffix| (*id, suffix)))
25 + .min_by_key(|(id, _)| *id)
26 + .map(|(_, suffix)| suffix.to_string())
27 + .unwrap_or_default()
28 +}
29 +
30 +/// Extract the suffix from one opening or closing HY structural token.
31 +fn hy_marker_suffix(token: &str) -> Option<&str> {
32 + let body = token.strip_prefix('<')?.strip_suffix('>')?;
33 + let body = body.strip_prefix('/').unwrap_or(body);
34 +
35 + HY_MARKER_STEMS.iter().find_map(|stem| {
36 + let suffix = body.strip_prefix(stem)?;
37 + (suffix.is_empty() || suffix.starts_with(':')).then_some(suffix)
38 + })
39 +}
rust/src/parser/src/unified/mod.rs
+2 -2
@@ -5,13 +5,13 @@
5
6 mod combined;
7 mod gemma4;
8 -mod hy_v3;
8 +mod hy;
9 mod inkling;
10 mod kimi_k3;
11
12 pub use combined::CombinedParser;
13 pub use gemma4::Gemma4UnifiedParser;
14 -pub use hy_v3::HyV3UnifiedParser;
14 +pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser};
15 pub use inkling::InklingUnifiedParser;
16 pub use kimi_k3::{KimiK3StructuralTagBuilder, KimiK3UnifiedParser};
17 use thiserror::Error;
tests/models/registry.py
+7
@@ -331,6 +331,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
331 "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),
332 "HunYuanMoEV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-A13B-Instruct"),
333 "HYV3ForCausalLM": _HfExamplesInfo("tencent/Hy3-preview", trust_remote_code=True),
334 + # Internal checkpoint, not publicly available yet.
335 + "HYV4ForCausalLM": _HfExamplesInfo("tencent/Hy4-preview", trust_remote_code=True),
336 "HyperCLOVAXForCausalLM": _HfExamplesInfo(
337 "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
338 min_transformers_version="5.9.0",
@@ -1723,6 +1725,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1725 "tencent/Hy3-preview",
1726 speculative_model="tencent/Hy3-preview",
1727 ),
1728 + # Internal checkpoint, not publicly available yet.
1729 + "HYV4MTPModel": _HfExamplesInfo(
1730 + "tencent/Hy4-preview",
1731 + speculative_model="tencent/Hy4-preview",
1732 + ),
1733 "InklingMTPModel": _HfExamplesInfo(
1734 "thinkingmachines/Inkling-NVFP4",
1735 speculative_model="thinkingmachines/Inkling-NVFP4",
tests/models/test_registry.py
+5
@@ -59,6 +59,11 @@ def test_registry_imports(model_arch):
59 ):
60 pytest.skip("Dots3 NOTE is only supported on CUDA")
61
62 + if model_arch in ("HYV4ForCausalLM", "HYV4MTPModel") and not (
63 + current_platform.is_cuda()
64 + ):
65 + pytest.skip("HY V4 is only supported on CUDA")
66 +
67 # Ensure all model classes can be imported successfully
68 model_cls = ModelRegistry._try_load_model_cls(model_arch)
69 assert model_cls is not None
tests/models/utils.py
+5
@@ -552,6 +552,11 @@ def dummy_hf_overrides(
552
553 text_config.update(update_dict)
554
555 + # Keep per-layer metadata consistent with the reduced layer count.
556 + # HY V4 uses this to decide which sparse-attention indexers are shared.
557 + if indexer_types := getattr(text_config, "indexer_types", None):
558 + text_config.update({"indexer_types": indexer_types[:num_hidden_layers]})
559 +
560 # Update n_layers and moe configs for Moondream3 model
561 if model_arch in ("Moondream3ForCausalLM", "HfMoondream"):
562 text_config.update(
vllm/config/speculative.py
+8
@@ -57,6 +57,7 @@ MTPModelTypes = Literal[
57 "pangu_ultra_moe_mtp",
58 "step3p5_mtp",
59 "hy_v3_mtp",
60 + "hy_v4_mtp",
61 "gemma4_mtp",
62 "inkling_mtp",
63 ]
@@ -917,6 +918,13 @@ class SpeculativeConfig:
918 {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]}
919 )
920
921 + if hf_config.model_type == "hy_v4":
922 + hf_config.model_type = "hy_v4_mtp"
923 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
924 + hf_config.update(
925 + {"n_predict": n_predict, "architectures": ["HYV4MTPModel"]}
926 + )
927 +
928 if hf_config.model_type in ("inkling_mm_model", "inkling_model"):
929 mtp_config = getattr(hf_config, "mtp_config", None) or {}
930 hf_config = getattr(hf_config, "text_config", hf_config)
vllm/config/vllm.py
+2
@@ -79,6 +79,8 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
79 "DeepseekV4ForCausalLM",
80 "DeepSeekV4MTPModel",
81 "GlmMoeDsaForCausalLM",
82 + "HYV4ForCausalLM",
83 + "HYV4MTPModel",
84 "InklingForCausalLM",
85 "InklingForConditionalGeneration",
86 "KimiK3ForConditionalGeneration",
vllm/envs.py
+10
@@ -321,6 +321,7 @@ if TYPE_CHECKING:
321 VLLM_GPU_NIC_PCIE_MAPPING: str = ""
322 VLLM_NIC_SELECTION_VARS: str = ""
323 VLLM_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None
324 + VLLM_ENABLE_HPC_OPS: bool = False
325
326
327 def get_default_cache_root():
@@ -2155,6 +2156,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
2156 # Each entry is VAR_NAME or VAR_NAME:<suffix> (suffix appended to
2157 # RDMA device name). Must be set together with VLLM_GPU_NIC_PCIE_MAPPING.
2158 "VLLM_NIC_SELECTION_VARS": lambda: os.getenv("VLLM_NIC_SELECTION_VARS", ""),
2159 + # If set to 1, enable the HPC fused kernels (requires the hpc package
2160 + # (.so) and an sm100/sm103 device). Covers:
2161 + # * the HY V4 iHC ops -- each of the eager HYV4HCPreLayer /
2162 + # HYV4HCPostLayer / HYV4HCHeadLayer bodies becomes one kernel launch;
2163 + # * the gated-MLA output gating (attn_out * sigmoid(gate projection)),
2164 + # fused into gated_mla_gemm; elementwise gating only.
2165 + # Each op additionally checks its own shape / dtype constraints and falls
2166 + # back to the eager path when they do not hold.
2167 + "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))),
2168 }
2169
2170
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+20 -5
@@ -79,12 +79,27 @@ class TrtLlmFp8ExpertsBase:
79 activation_key,
80 activation_format,
81 )
82 - if not supported or moe_config.num_experts <= 2048:
82 + if not supported:
83 return supported, reason
84 - return False, (
85 - "FlashInfer TRTLLM routing supports at most 2048 experts, "
86 - f"but got {moe_config.num_experts}"
87 - )
84 + # FlashInfer accepts gemm1_alpha/beta/clamp_limit only for MXFP8; the
85 + # DeepSeek-FP8 block-scale kernel rejects them and the per-tensor kernel
86 + # has no clamp at all, so a model-requested SwiGLU clamp cannot be
87 + # honored on those paths.
88 + if (
89 + moe_config.swiglu_limit is not None
90 + or moe_config.swiglu_alpha is not None
91 + or moe_config.swiglu_beta is not None
92 + ) and (weight_key, activation_key) != (kMxfp8Static, kMxfp8Dynamic):
93 + return False, (
94 + "the TRTLLM FP8 kernels apply the SwiGLU alpha/beta/clamp "
95 + "parameters only for MXFP8 weights"
96 + )
97 + if moe_config.num_experts > 2048:
98 + return False, (
99 + "FlashInfer TRTLLM routing supports at most 2048 experts, "
100 + f"but got {moe_config.num_experts}"
101 + )
102 + return True, None
103
104 def __init__(
105 self,
vllm/model_executor/layers/hpc/__init__.py
+14
@@ -1,11 +1,25 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from vllm.model_executor.layers.hpc.gated_mla import (
5 + hpc_gated_mla_gemm,
6 + hpc_gated_mla_supported,
7 +)
8 +from vllm.model_executor.layers.hpc.hpc_ihc import (
9 + HpcIHCHead,
10 + HpcIHCPost,
11 + HpcIHCPre,
12 +)
13 from vllm.model_executor.layers.hpc.hpc_module import HpcModule
14 from vllm.model_executor.layers.hpc.rope_norm import HpcRopeNorm, QkNormPolicy
15
16 __all__ = [
17 + "HpcIHCHead",
18 + "HpcIHCPost",
19 + "HpcIHCPre",
20 "HpcModule",
21 "HpcRopeNorm",
22 "QkNormPolicy",
23 + "hpc_gated_mla_gemm",
24 + "hpc_gated_mla_supported",
25 ]
vllm/model_executor/layers/hpc/gated_mla.py new
+113
@@ -0,0 +1,113 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Enablement gate for the HPC fused gated-MLA GEMM.
4 +
5 +``hpc.gated_mla_gemm`` fuses the MLA output gating -- the gate projection, the
6 +sigmoid and the elementwise product -- into a single kernel launch, replacing
7 +the eager ``attn_out * sigmoid(hidden_states @ weight.T)`` sequence.
8 +
9 +Kernel constraints:
10 + - Requires VLLM_ENABLE_HPC_OPS=1
11 + - Requires the hpc package (.so) built for the current arch
12 + - Only sm100 / sm103 (compute capability 100, 103)
13 + - All three operands must be bfloat16 and contiguous
14 + - The gate width must be a multiple of 256
15 + - Only elementwise gating (headwise has no fused form)
16 +"""
17 +
18 +import torch
19 +
20 +from vllm import envs
21 +from vllm.logger import init_logger
22 +from vllm.utils.hpc import has_hpc
23 +
24 +logger = init_logger(__name__)
25 +
26 +# hpc.gated_mla_gemm fuses `attn_out * sigmoid(x @ weight.T)` into one launch.
27 +# It is built for sm100/sm103 and needs all three operands bf16 + contiguous,
28 +# with the gate width a multiple of 256.
29 +_HPC_GATED_GEMM_N_ALIGN = 256
30 +_HPC_GATED_GEMM_CAPABILITIES: frozenset[int] = frozenset({100, 103})
31 +
32 +
33 +def hpc_gated_mla_supported(
34 + gating_type: str | None, attn_output_gate: torch.nn.Module | None
35 +) -> bool:
36 + """Gate for running the MLA output gating through hpc.gated_mla_gemm."""
37 + if not envs.VLLM_ENABLE_HPC_OPS:
38 + return False
39 +
40 + if attn_output_gate is None:
41 + return False
42 +
43 + if not has_hpc():
44 + logger.warning_once(
45 + "HPC gated MLA disabled: 'hpc' package is not installed. "
46 + "Please install the HPC library (.so) to enable the fused kernel."
47 + )
48 + return False
49 +
50 + from vllm.platforms import current_platform
51 +
52 + if not current_platform.is_cuda():
53 + logger.warning_once("HPC gated MLA disabled: only CUDA is supported.")
54 + return False
55 +
56 + capability = current_platform.get_device_capability()
57 + if capability is None or capability.to_int() not in _HPC_GATED_GEMM_CAPABILITIES:
58 + logger.warning_once(
59 + "HPC gated MLA disabled: compute capability %s not in %s.",
60 + capability,
61 + _HPC_GATED_GEMM_CAPABILITIES,
62 + )
63 + return False
64 +
65 + # Headwise gating broadcasts one scalar per head over v_head_dim; the
66 + # kernel only implements the elementwise product.
67 + if gating_type != "elementwise":
68 + logger.warning_once(
69 + "HPC gated MLA disabled: gating_type '%s' has no fused form.",
70 + gating_type,
71 + )
72 + return False
73 +
74 + weight = getattr(attn_output_gate, "weight", None)
75 + if weight is None or weight.dtype != torch.bfloat16:
76 + logger.warning_once(
77 + "HPC gated MLA disabled: gate weight dtype is %s, the kernel needs "
78 + "bfloat16 (keep the layer out of quantization).",
79 + None if weight is None else weight.dtype,
80 + )
81 + return False
82 +
83 + if weight.shape[0] % _HPC_GATED_GEMM_N_ALIGN != 0:
84 + logger.warning_once(
85 + "HPC gated MLA disabled: gate width %d is not a multiple of %d.",
86 + weight.shape[0],
87 + _HPC_GATED_GEMM_N_ALIGN,
88 + )
89 + return False
90 +
91 + logger.info_once("HPC gated MLA enabled by set VLLM_ENABLE_HPC_OPS.")
92 + return True
93 +
94 +
95 +def hpc_gated_mla_gemm(
96 + x: torch.Tensor,
97 + weight: torch.Tensor,
98 + atten_output: torch.Tensor,
99 +) -> torch.Tensor:
100 + """Fused ``atten_output * sigmoid(x @ weight.T)``.
101 +
102 + Args:
103 + x: ``[m, k]`` bfloat16, contiguous.
104 + weight: ``[n, k]`` bfloat16, contiguous and row-major (the kernel
105 + transposes it). ``n`` must be a multiple of 256.
106 + atten_output: ``[m, n]`` bfloat16, contiguous.
107 +
108 + Returns:
109 + ``[m, n]`` bfloat16.
110 + """
111 + from hpc.gemm import gated_mla_gemm
112 +
113 + return gated_mla_gemm(x, weight, atten_output)
vllm/model_executor/layers/hpc/hpc_ihc.py new
+261
@@ -0,0 +1,261 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HPC fused iHC (independent Hyper-Connections) kernels for HY V4.
4 +
5 +Replaces the eager HYV4HCPreLayer / HYV4HCPostLayer / HYV4HCHeadLayer bodies
6 +with single-kernel HPC implementations:
7 +
8 + pre : rms square sum + w projection + sigmoid gates + hc-dim weighted sum
9 + post : y[n, i, :] = H_post[n, i] * x[n, :] + residual[n, i, :]
10 + head : same as pre but hc_mult projection rows and no H_post output
11 +
12 +The eager path issues 20 / 5 / 15 kernels for pre / post / head; each HPC op is
13 +a single kernel.
14 +
15 +Constraints:
16 + - Requires VLLM_ENABLE_HPC_OPS=1
17 + - Requires the hpc package (.so) built for the current arch
18 + - Only sm100 / sm103 (compute capability 100, 103)
19 + - Only hc_mult == 4 and hidden_size in {4096, 6144}
20 + - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint
21 + keeps them out of fp8 quantization via modules_to_not_convert)
22 +
23 +The fused modules own no parameters: they hold references to the weights of the
24 +eager layer they replace, bound in ``process_weights_after_loading`` (invoked
25 +generically for every ``HpcModule`` by the model loader).
26 +
27 +NOTE: The reference implementation also offers a cross-layer post+pre fusion
28 +(``HpcIHCPostPre``) that folds one segment's post into the next segment's pre.
29 +It requires reworking the decoder-layer forward scheduling and is not ported
30 +here. TODO: add it once the decoder dataflow is restructured for it.
31 +"""
32 +
33 +import torch
34 +
35 +from vllm import envs
36 +from vllm.logger import init_logger
37 +from vllm.model_executor.layers.hpc.hpc_module import HpcModule
38 +from vllm.utils.hpc import has_hpc
39 +
40 +logger = init_logger(__name__)
41 +
42 +# Shapes the HPC kernels are instantiated for. Other shapes launch nothing,
43 +# so gate on this rather than relying on a silent no-op.
44 +_SUPPORTED_HIDDEN_SIZES: frozenset[int] = frozenset({4096, 6144})
45 +_SUPPORTED_HC_MULTS: frozenset[int] = frozenset({4})
46 +_SUPPORTED_CAPABILITIES: frozenset[int] = frozenset({100, 103})
47 +
48 +
49 +def _ihc_supported(hc_mult: int, hidden_size: int) -> bool:
50 + """Shared gate for all three iHC ops."""
51 + if not envs.VLLM_ENABLE_HPC_OPS:
52 + return False
53 +
54 + if not has_hpc():
55 + logger.warning_once(
56 + "HPC iHC disabled: 'hpc' package is not installed. "
57 + "Please install the HPC library (.so) to enable fused kernels."
58 + )
59 + return False
60 +
61 + from vllm.platforms import current_platform
62 +
63 + if not current_platform.is_cuda():
64 + logger.warning_once("HPC iHC disabled: only CUDA is supported.")
65 + return False
66 +
67 + capability = current_platform.get_device_capability()
68 + if capability is None or capability.to_int() not in _SUPPORTED_CAPABILITIES:
69 + logger.warning_once(
70 + "HPC iHC disabled: compute capability %s not in %s.",
71 + capability,
72 + _SUPPORTED_CAPABILITIES,
73 + )
74 + return False
75 +
76 + if hc_mult not in _SUPPORTED_HC_MULTS:
77 + logger.warning_once(
78 + "HPC iHC disabled: hc_mult=%d not in %s.", hc_mult, _SUPPORTED_HC_MULTS
79 + )
80 + return False
81 +
82 + if hidden_size not in _SUPPORTED_HIDDEN_SIZES:
83 + logger.warning_once(
84 + "HPC iHC disabled: hidden_size=%d not in %s.",
85 + hidden_size,
86 + _SUPPORTED_HIDDEN_SIZES,
87 + )
88 + return False
89 +
90 + if envs.VLLM_BATCH_INVARIANT:
91 + # Batch-invariant mode relies on torch's deterministic reductions; the
92 + # HPC kernel reduces in a different order, so keep the two regimes
93 + # apart.
94 + logger.warning_once(
95 + "HPC iHC disabled: not supported under VLLM_BATCH_INVARIANT."
96 + )
97 + return False
98 +
99 + logger.info_once("HPC iHC enabled by set VLLM_ENABLE_HPC_OPS.")
100 + return True
101 +
102 +
103 +class HpcIHCPre(HpcModule):
104 + """Fused iHC pre block.
105 +
106 + Computes, in one kernel:
107 + x_flat = x.flatten(1)
108 + r = rsqrt(x_flat.square().mean(-1) + rms_norm_eps)
109 + mixes = (x_flat @ w.T) * r
110 + H_pre = sigmoid(mixes[:, :hc] * hc_scale[0] + hc_base[:hc]) + hc_eps
111 + H_post = magnitude * sigmoid(
112 + mixes[:, hc:] * hc_scale[1] + hc_base[hc:]) + hc_eps
113 + y = sum_i H_pre[:, i] * x[:, i, :]
114 +
115 + Args:
116 + hc_mult: HC expand ratio.
117 + hidden_size: Model hidden dimension.
118 + magnitude: H_post multiplier (config.hc_magnitude).
119 + hc_eps: Epsilon added to both gates (config.hc_eps).
120 + norm_eps: Epsilon inside the rsqrt (config.rms_norm_eps).
121 + fallback_op: The eager HYV4HCPreLayer to source weights from.
122 + norm_owner: Optional RMSNorm that immediately follows this pre block.
123 + When given, its weight/eps are folded into the kernel so a single
124 + launch covers pre + RMSNorm (the caller then skips the separate
125 + layernorm). None keeps the original pre-only behaviour.
126 + """
127 +
128 + def __init__(
129 + self,
130 + hc_mult: int,
131 + hidden_size: int,
132 + magnitude: float,
133 + hc_eps: float,
134 + norm_eps: float,
135 + fallback_op: torch.nn.Module,
136 + norm_owner: torch.nn.Module | None = None,
137 + ) -> None:
138 + super().__init__()
139 + self.hc_mult = hc_mult
140 + self.hidden_size = hidden_size
141 + self.magnitude = magnitude
142 + self.hc_eps = hc_eps
143 + self.norm_eps = norm_eps
144 + # Stash the owner outside nn.Module's attribute machinery: callers pass
145 + # their own `self` here, and registering that as a submodule would make
146 + # the module tree cyclic (state_dict() then hits RecursionError).
147 + object.__setattr__(self, "_fallback_op", fallback_op)
148 + # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module).
149 + object.__setattr__(self, "_norm_owner", norm_owner)
150 + self.weight = None
151 + self.hc_scale = None
152 + self.hc_base = None
153 + self.rms_weight = None
154 + self.rms_eps = 0.0
155 +
156 + @classmethod
157 + def support(cls, hc_mult: int, hidden_size: int) -> bool:
158 + return _ihc_supported(hc_mult, hidden_size)
159 +
160 + def process_weights_after_loading(self, model=None):
161 + # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d].
162 + # It stays float32 -- the HPC kernel consumes fp32 weights directly, so
163 + # there is no conversion (and no precision loss) here.
164 + owner = self._fallback_op
165 + self.weight = owner.hc_fn.weight
166 + self.hc_scale = owner.hc_scale
167 + self.hc_base = owner.hc_base
168 + if self._norm_owner is not None:
169 + self.rms_weight = self._norm_owner.weight
170 + self.rms_eps = self._norm_owner.variance_epsilon
171 +
172 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
173 + import hpc
174 +
175 + return hpc.fuse_ihc_pre(
176 + x,
177 + self.weight,
178 + self.hc_scale,
179 + self.hc_base,
180 + self.norm_eps,
181 + self.hc_eps,
182 + self.magnitude,
183 + self.rms_weight,
184 + self.rms_eps,
185 + )
186 +
187 +
188 +class HpcIHCPost(HpcModule):
189 + """Fused iHC post block: H_post gating plus multi-channel residual add.
190 +
191 + Unlike mHC there is no comb matrix, so each output channel only needs its
192 + own residual channel and the whole thing is one fused multiply-add per
193 + element.
194 + """
195 +
196 + def __init__(self, hc_mult: int, hidden_size: int) -> None:
197 + super().__init__()
198 + self.hc_mult = hc_mult
199 + self.hidden_size = hidden_size
200 +
201 + @classmethod
202 + def support(cls, hc_mult: int, hidden_size: int) -> bool:
203 + return _ihc_supported(hc_mult, hidden_size)
204 +
205 + def forward(
206 + self, x: torch.Tensor, residual: torch.Tensor, H_post: torch.Tensor
207 + ) -> torch.Tensor:
208 + import hpc
209 +
210 + return hpc.fuse_ihc_post(x, residual, H_post)
211 +
212 +
213 +class HpcIHCHead(HpcModule):
214 + """Fused iHC head block: merge the hc channels back into one hidden state.
215 +
216 + Same structure as HpcIHCPre but the projection emits only hc_mult gate
217 + logits and there is no H_post output. Called once per forward, after the
218 + last decoder layer.
219 + """
220 +
221 + def __init__(
222 + self,
223 + hc_mult: int,
224 + hidden_size: int,
225 + hc_eps: float,
226 + norm_eps: float,
227 + fallback_op: torch.nn.Module,
228 + ) -> None:
229 + super().__init__()
230 + self.hc_mult = hc_mult
231 + self.hidden_size = hidden_size
232 + self.hc_eps = hc_eps
233 + self.norm_eps = norm_eps
234 + # See HpcIHCPre: keep the owner out of the module tree.
235 + object.__setattr__(self, "_fallback_op", fallback_op)
236 + self.weight = None
237 + self.hc_scale = None
238 + self.hc_base = None
239 +
240 + @classmethod
241 + def support(cls, hc_mult: int, hidden_size: int) -> bool:
242 + return _ihc_supported(hc_mult, hidden_size)
243 +
244 + def process_weights_after_loading(self, model=None):
245 + # hc_head_fn weight is [hc_mult, hc_mult * d], float32.
246 + owner = self._fallback_op
247 + self.weight = owner.hc_head_fn.weight
248 + self.hc_scale = owner.hc_head_scale
249 + self.hc_base = owner.hc_head_base
250 +
251 + def forward(self, x: torch.Tensor) -> torch.Tensor:
252 + import hpc
253 +
254 + return hpc.fuse_ihc_head(
255 + x,
256 + self.weight,
257 + self.hc_scale,
258 + self.hc_base,
259 + self.norm_eps,
260 + self.hc_eps,
261 + )
vllm/model_executor/layers/logits_processor.py
+7 -1
@@ -15,6 +15,7 @@ from vllm.distributed import (
15 )
16 from vllm.logger import init_logger
17 from vllm.model_executor.custom_op import PluggableLayer
18 +from vllm.model_executor.layers.linear import UnquantizedLinearMethod
19 from vllm.model_executor.layers.vocab_parallel_embedding import (
20 UnquantizedEmbeddingMethod,
21 VocabParallelEmbedding,
@@ -144,7 +145,12 @@ class LogitsProcessor(PluggableLayer):
145 lm_head, hidden_states, bias=embedding_bias
146 )
147
147 - if not isinstance(lm_head.quant_method, UnquantizedEmbeddingMethod):
148 + # A quant config that excludes lm_head hands out UnquantizedLinearMethod
149 + # rather than UnquantizedEmbeddingMethod, so accept both: either way the
150 + # weight is plain and `lm_head.weight` can be cast directly.
151 + if not isinstance(
152 + lm_head.quant_method, (UnquantizedEmbeddingMethod, UnquantizedLinearMethod)
153 + ):
154 raise ValueError(
155 "A head_dtype different from the model dtype is only "
156 "supported for an unquantized lm_head."
vllm/model_executor/models/deepseek_v2.py
+10 -2
@@ -874,12 +874,20 @@ def _try_load_fp8_indexer_wk(
874 # We have both weight and scale: dequantize FP8 to BF16.
875 weight_fp8, scale_inv = entry["weight"], entry["scale"]
876 del buf[layer_prefix]
877 + if scale_inv.dtype in (torch.uint8, torch.float8_e8m0fnu):
878 + # MXFP8 checkpoints store power-of-two E8M0 scales, decode them to float.
879 + scale_inv = scale_inv.view(torch.float8_e8m0fnu).to(torch.float32)
880 if scale_inv.ndim == 1:
881 # Per-channel scale: one scale per row of [out, in]
882 group_shape = GroupShape(1, weight_fp8.shape[1])
883 else:
881 - block_size = weight_fp8.shape[1] // scale_inv.shape[1]
882 - group_shape = GroupShape(block_size, block_size)
884 + # Derive the block size independently per dim so that per-channel
885 + # ([out, 1]), MX ([out, in / 32]) and 2D block ([out / 128, in / 128])
886 + # scale layouts are all handled.
887 + group_shape = GroupShape(
888 + weight_fp8.shape[0] // scale_inv.shape[0],
889 + weight_fp8.shape[1] // scale_inv.shape[1],
890 + )
891 weight_bf16 = scaled_dequantize(
892 weight_fp8,
893 scale_inv,
vllm/model_executor/models/registry.py
+2
@@ -128,6 +128,7 @@ _TEXT_GENERATION_MODELS = {
128 "GraniteSWAForCausalLM": ("granite", "GraniteForCausalLM"),
129 "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),
130 "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"),
131 + "HYV4ForCausalLM": ("vllm.models.hy_v4", "HYV4ForCausalLM"),
132 "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"),
133 "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"),
134 "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),
@@ -663,6 +664,7 @@ _SPECULATIVE_DECODING_MODELS = {
664 "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
665 "InternS2MobiusMTP": ("interns2_mobius", "InternS2MobiusMTP"),
666 "HYV3MTPModel": ("hy_v3_mtp", "HYV3MTP"),
667 + "HYV4MTPModel": ("vllm.models.hy_v4", "HYV4MTP"),
668 "KimiK3MTPModel": ("vllm.models.kimi_k3", "KimiK3MTP"),
669 # Temporarily disabled.
670 # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.
vllm/models/hy_v4/__init__.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HY V4 (``hy_v4``) model — hardware-isolated entry point.
4 +
5 +HY V4 combines three architectural pieces:
6 +
7 +- **iHC** (independent Hyper-Connections): the single residual stream is
8 + replaced by ``hc_mult`` parallel residual channels, gated per sub-block.
9 +- **MLA + lightning indexer**: multi-head latent attention with an optional
10 + DSA-style sparse top-k selection, plus an output gate and a learnable sink.
11 +- **MoE**: sigmoid-routed experts with a clamped SwiGLU and shared experts.
12 +
13 +The package is organized like `vllm.models.deepseek_v32`: this module is the
14 +only public entry point and dispatches on the current platform, so registry
15 +entries never reach into a platform subpackage.
16 +
17 +Only NVIDIA is supported for now. The port also drops the reference
18 +implementation's HPC/TPCP fusion paths, which depend on infrastructure that
19 +does not exist in this tree.
20 +"""
21 +
22 +from vllm.platforms import current_platform
23 +
24 +if current_platform.is_rocm():
25 + raise NotImplementedError("hy_v4 does not yet support ROCm.")
26 +elif current_platform.is_xpu():
27 + raise NotImplementedError("hy_v4 does not yet support XPU.")
28 +else:
29 + # Covers Blackwell (sm100) and all other CUDA devices.
30 + from .nvidia.model import HYV4ForCausalLM
31 + from .nvidia.mtp import HYV4MTP
32 +
33 +__all__ = [
34 + "HYV4ForCausalLM",
35 + "HYV4MTP",
36 +]
vllm/models/hy_v4/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/hy_v4/nvidia/attention.py new
+765
@@ -0,0 +1,765 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MLA attention and lightning indexer for HY V4 (NVIDIA).
4 +
5 +NOTE: The reference implementation also offers a fused MLA preprocessing
6 +wrapper, which is not available here, so this port keeps the eager MLA path
7 +with an FP8 indexer cache. TODO: restore the fused preprocessing once the
8 +kernel lands. The HPC gated-MLA GEMM *is* wired up: it replaces the eager
9 +output gating when the hpc package is installed, ``VLLM_ENABLE_HPC_OPS=1``
10 +and the elementwise/dtype/alignment constraints hold.
11 +
12 +The per-head learnable sink is supported through `.flashmla_sparse`, which
13 +subclasses the platform's sparse MLA backend to forward ``attn_sink``.
14 +"""
15 +
16 +from typing import cast
17 +
18 +import regex as re
19 +import torch
20 +from torch import nn
21 +from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig
22 +
23 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
24 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
25 +from vllm.config.cache import CacheDType
26 +from vllm.distributed import get_tensor_model_parallel_world_size
27 +from vllm.logger import init_logger
28 +from vllm.model_executor.layers.attention import MLAAttention
29 +from vllm.model_executor.layers.hpc import hpc_gated_mla_gemm, hpc_gated_mla_supported
30 +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
31 +from vllm.model_executor.layers.linear import (
32 + ColumnParallelLinear,
33 + MergedColumnParallelLinear,
34 + ReplicatedLinear,
35 + RowParallelLinear,
36 +)
37 +from vllm.model_executor.layers.quantization import QuantizationConfig
38 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
39 + per_token_group_quant_fp8,
40 +)
41 +from vllm.model_executor.layers.rotary_embedding import get_rope
42 +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
43 +from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
44 +from vllm.platforms import current_platform
45 +from vllm.v1.attention.backend import AttentionBackend, AttentionType
46 +from vllm.v1.attention.selector import get_attn_backend
47 +
48 +logger = init_logger(__name__)
49 +
50 +_SPARSE_LAYER_TYPES = ("sparse_attention", "sparse", "deepseek_sparse_attention")
51 +_WEIGHT_LAYER_INDEX_RE = re.compile(r"(?:^|\.)layers\.(\d+)(?:\.|$)")
52 +
53 +
54 +def compute_skip_topk_layers(config: PretrainedConfig) -> set[int]:
55 + """Return the backbone layers that reuse a previous layer's top-k indices.
56 +
57 + A "shared" indexer layer performs sparse attention with the indices computed
58 + by the closest preceding "full" indexer layer, so it does not build its own
59 + indexer and its checkpoint indexer weights must be skipped.
60 +
61 + Args:
62 + config: The model config.
63 +
64 + Returns:
65 + The set of layer indices that share another layer's top-k indices.
66 +
67 + Raises:
68 + ValueError: If ``indexer_types`` has the wrong length or an unknown
69 + entry, or if ``index_topk_freq`` is not a positive integer.
70 + """
71 + if not hasattr(config, "index_topk"):
72 + return set()
73 +
74 + num_hidden_layers = config.num_hidden_layers
75 + indexer_types = getattr(config, "indexer_types", None)
76 + if indexer_types is not None:
77 + if len(indexer_types) != num_hidden_layers:
78 + raise ValueError(
79 + "indexer_types must contain one entry per hidden layer: "
80 + f"expected {num_hidden_layers}, got {len(indexer_types)}."
81 + )
82 + invalid_types = sorted(set(indexer_types) - {"full", "shared"})
83 + if invalid_types:
84 + raise ValueError(
85 + f"indexer_types only supports 'full' and 'shared', got {invalid_types}."
86 + )
87 + return {
88 + layer_idx
89 + for layer_idx, indexer_type in enumerate(indexer_types)
90 + if indexer_type == "shared"
91 + }
92 +
93 + freq = getattr(config, "index_topk_freq", 1)
94 + if not isinstance(freq, int) or freq <= 0:
95 + raise ValueError(f"index_topk_freq must be a positive integer, got {freq!r}.")
96 + pattern = getattr(config, "index_topk_pattern", None)
97 + offset = getattr(config, "index_skip_topk_offset", 2)
98 + skip_layers: set[int] = set()
99 + for layer_idx in range(num_hidden_layers):
100 + if pattern is None:
101 + if max(layer_idx - offset + 1, 0) % freq != 0:
102 + skip_layers.add(layer_idx)
103 + elif 0 <= layer_idx < len(pattern) and pattern[layer_idx] == "S":
104 + skip_layers.add(layer_idx)
105 + return skip_layers
106 +
107 +
108 +def is_skip_topk_indexer_weight(weight_name: str, skip_topk_layers: set[int]) -> bool:
109 + """Return whether an indexer weight belongs to a top-k sharing layer.
110 +
111 + Args:
112 + weight_name: Checkpoint weight name.
113 + skip_topk_layers: Result of `compute_skip_topk_layers`.
114 +
115 + Returns:
116 + True when the weight is an indexer weight of a layer that has no
117 + indexer module and therefore must be dropped.
118 + """
119 + if ".indexer." not in weight_name or not skip_topk_layers:
120 + return False
121 + match = _WEIGHT_LAYER_INDEX_RE.search(weight_name)
122 + return match is not None and int(match.group(1)) in skip_topk_layers
123 +
124 +
125 +class Indexer(nn.Module):
126 + """Lightning indexer selecting the top-k tokens for sparse MLA."""
127 +
128 + def __init__(
129 + self,
130 + vllm_config: VllmConfig,
131 + config: DeepseekV2Config | DeepseekV3Config,
132 + hidden_size: int,
133 + q_lora_rank: int,
134 + quant_config: QuantizationConfig | None,
135 + cache_config: CacheConfig | None,
136 + topk_indices_buffer: torch.Tensor | None,
137 + prefix: str = "",
138 + ):
139 + super().__init__()
140 + self.vllm_config = vllm_config
141 + self.config = config
142 + self.quant_config = quant_config
143 + self.topk_tokens = config.index_topk
144 + self.n_head = config.index_n_heads
145 + self.head_dim = config.index_head_dim
146 + self.rope_dim = config.qk_rope_head_dim
147 + self.q_lora_rank = q_lora_rank
148 +
149 + # No tensor parallelism, just replicated.
150 + self.wq_b = ReplicatedLinear(
151 + self.q_lora_rank,
152 + self.head_dim * self.n_head,
153 + bias=False,
154 + quant_config=quant_config,
155 + prefix=f"{prefix}.wq_b",
156 + )
157 + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
158 + # FP8 wk weights are upcast to BF16 while loading to keep the fusion.
159 + self.wk_weights_proj = MergedColumnParallelLinear(
160 + hidden_size,
161 + [self.head_dim, self.n_head],
162 + bias=False,
163 + quant_config=None,
164 + disable_tp=True,
165 + prefix=f"{prefix}.wk_weights_proj",
166 + )
167 + self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
168 + self.softmax_scale = self.head_dim**-0.5
169 +
170 + self.scale_fmt = "ue8m0"
171 + self.quant_block_size = 128
172 + self.topk_indices_buffer = topk_indices_buffer
173 +
174 + # FP8 naive cache: values in fp8 plus one fp32 scale per
175 + # ``quant_block_size`` elements.
176 + assert cache_config is not None, "HYV4 indexer requires cache_config"
177 + self.k_cache = DeepseekV32IndexerCache(
178 + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
179 + dtype=torch.uint8,
180 + prefix=f"{prefix}.k_cache",
181 + cache_config=cache_config,
182 + )
183 + self.max_model_len = vllm_config.model_config.max_model_len
184 + self.prefix = prefix
185 +
186 + from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
187 +
188 + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
189 + self.indexer_op = SparseAttnIndexer(
190 + self.k_cache,
191 + self.quant_block_size,
192 + self.scale_fmt,
193 + self.topk_tokens,
194 + self.head_dim,
195 + self.max_model_len,
196 + self.max_total_seq_len,
197 + self.topk_indices_buffer,
198 + )
199 +
200 + def forward(
201 + self,
202 + hidden_states: torch.Tensor,
203 + qr: torch.Tensor,
204 + positions: torch.Tensor,
205 + rotary_emb: nn.Module,
206 + ) -> torch.Tensor:
207 + hidden_states, q_quant, k, weights = self.prepare_inputs(
208 + hidden_states, qr, positions, rotary_emb
209 + )
210 + return self.indexer_op(hidden_states, q_quant, k, weights)
211 +
212 + def prepare_inputs(
213 + self,
214 + hidden_states: torch.Tensor,
215 + qr: torch.Tensor,
216 + positions: torch.Tensor,
217 + rotary_emb: nn.Module,
218 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
219 + """Build the quantized query, key and per-head weights of the indexer."""
220 + q, _ = self.wq_b(qr)
221 + q = q.view(-1, self.n_head, self.head_dim)
222 + # Checkpoint (PTM) layout: pe occupies the LAST rope_dim dims.
223 + q_nope, q_pe = torch.split(
224 + q, [self.head_dim - self.rope_dim, self.rope_dim], dim=-1
225 + )
226 +
227 + kw, _ = self.wk_weights_proj(hidden_states)
228 + k = kw[:, : self.head_dim]
229 + weights = kw[:, self.head_dim :]
230 +
231 + k = self.k_norm(k)
232 + k_nope, k_pe = torch.split(
233 + k, [self.head_dim - self.rope_dim, self.rope_dim], dim=-1
234 + )
235 +
236 + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
237 + # RoPE (NeoX) can introduce extra leading dims, so flatten back to the
238 + # token-major shapes.
239 + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
240 + k_pe = k_pe.reshape(-1, 1, self.rope_dim)
241 +
242 + # Reassemble with the original physical layout: no_pe first, pe last.
243 + q = torch.cat([q_nope, q_pe], dim=-1)
244 + # ``k_pe`` is [num_tokens, 1, rope_dim] (MQA).
245 + k = torch.cat([k_nope, k_pe.squeeze(-2)], dim=-1)
246 +
247 + # Only q is quantized here; k quantization is fused with cache insertion.
248 + q = q.view(-1, self.head_dim)
249 + q_fp8, q_scale = per_token_group_quant_fp8(
250 + q,
251 + self.quant_block_size,
252 + column_major_scales=False,
253 + use_ue8m0=self.scale_fmt is not None,
254 + )
255 + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
256 + q_scale = q_scale.view(-1, self.n_head, 1)
257 +
258 + weights = (
259 + weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
260 + )
261 + weights = weights.squeeze(-1)
262 +
263 + return hidden_states, q_fp8, k, weights
264 +
265 +
266 +class HYV4MLAAttention(nn.Module):
267 + """Multi-head latent attention with optional sparse lightning indexer.
268 +
269 + Main reference: the DeepSeek-V2 paper and the FlashInfer implementation
270 + (https://arxiv.org/abs/2405.04434). HY V4 additionally supports an output
271 + gate (``gated_mla``) and a per-head learnable attention sink.
272 +
273 + The sink is applied by binding the sink-capable backend from
274 + `.flashmla_sparse`; if no backend on this platform can consume sinks, the
275 + weight is still loaded but the bias is disabled with a warning.
276 + """
277 +
278 + def __init__(
279 + self,
280 + vllm_config: VllmConfig,
281 + config: PretrainedConfig,
282 + hidden_size: int,
283 + num_heads: int,
284 + qk_nope_head_dim: int,
285 + qk_rope_head_dim: int,
286 + v_head_dim: int,
287 + q_lora_rank: int | None,
288 + kv_lora_rank: int,
289 + max_position_embeddings: int = 8192,
290 + cache_config: CacheConfig | None = None,
291 + quant_config: QuantizationConfig | None = None,
292 + prefix: str = "",
293 + topk_indices_buffer: torch.Tensor | None = None,
294 + layer_idx: int = 0,
295 + ) -> None:
296 + super().__init__()
297 + self.config = config
298 + self.hidden_size = hidden_size
299 + self.qk_nope_head_dim = qk_nope_head_dim
300 + self.qk_rope_head_dim = qk_rope_head_dim
301 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
302 + self.v_head_dim = v_head_dim
303 + self.layer_idx = layer_idx
304 + self.q_lora_rank = q_lora_rank
305 + self.kv_lora_rank = kv_lora_rank
306 +
307 + self.num_heads = num_heads
308 + tp_size = get_tensor_model_parallel_world_size()
309 + assert num_heads % tp_size == 0
310 + self.num_local_heads = num_heads // tp_size
311 +
312 + self.layer_id = int(prefix.split(".")[-2])
313 + layer_types = getattr(config, "layer_types", None)
314 + requested_sparse = (
315 + hasattr(config, "index_topk")
316 + and layer_types is not None
317 + and self.layer_id < len(layer_types)
318 + and layer_types[self.layer_id] in _SPARSE_LAYER_TYPES
319 + )
320 + # Only actual sparse layers may share another layer's top-k indices.
321 + self.skip_topk = requested_sparse and self.layer_id in compute_skip_topk_layers(
322 + config
323 + )
324 + # The skip pattern only governs backbone layers. MTP/nextn layers
325 + # (layer_id >= num_hidden_layers) always build a full indexer: they
326 + # compute indices at draft step 0 and toggle at runtime.
327 + num_hidden_layers = getattr(config, "num_hidden_layers", None)
328 + is_mtp_layer = (
329 + num_hidden_layers is not None and self.layer_id >= num_hidden_layers
330 + )
331 + self.create_indexer = requested_sparse and (not self.skip_topk or is_mtp_layer)
332 + self.is_sparse = requested_sparse
333 +
334 + # Do not silently degrade sparse layers into dense attention. Probe the
335 + # sparse MLA backend directly and fail fast with the real error.
336 + kv_cache_dtype = cache_config.cache_dtype if cache_config else "auto"
337 + if self.is_sparse:
338 + try:
339 + get_attn_backend(
340 + head_size=self.kv_lora_rank + self.qk_rope_head_dim,
341 + dtype=torch.get_default_dtype(),
342 + kv_cache_dtype=kv_cache_dtype,
343 + use_mla=True,
344 + has_sink=False,
345 + use_sparse=True,
346 + num_heads=self.num_local_heads,
347 + )
348 + except Exception as exc:
349 + raise RuntimeError(
350 + "HYV4 sparse attention was requested, but no valid sparse MLA "
351 + "backend is available for current runtime/config. "
352 + "Refusing to fall back to dense attention."
353 + ) from exc
354 +
355 + self.scaling = self.qk_head_dim**-0.5
356 + self.max_position_embeddings = max_position_embeddings
357 + self.fused_qkv_a_proj = None
358 + self.kv_a_proj_with_mqa = None
359 + if self.q_lora_rank is not None:
360 + # ``q_a_proj`` and ``kv_a_proj_with_mqa`` read the same
361 + # ``hidden_states`` and are both TP-replicated, so they run as one
362 + # GEMM. The checkpoint keeps them separate; `load_weights` merges
363 + # them through ``stacked_params_mapping``.
364 + self.fused_qkv_a_proj = MergedColumnParallelLinear(
365 + self.hidden_size,
366 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
367 + bias=False,
368 + quant_config=quant_config,
369 + prefix=f"{prefix}.fused_qkv_a_proj",
370 + disable_tp=True,
371 + )
372 + else:
373 + self.kv_a_proj_with_mqa = ReplicatedLinear(
374 + self.hidden_size,
375 + self.kv_lora_rank + self.qk_rope_head_dim,
376 + bias=False,
377 + quant_config=quant_config,
378 + prefix=f"{prefix}.kv_a_proj_with_mqa",
379 + )
380 +
381 + self.q_a_layernorm = None
382 + self.q_b_proj = None
383 + self.q_proj = None
384 + if self.q_lora_rank is not None:
385 + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
386 + self.q_b_proj = ColumnParallelLinear(
387 + self.q_lora_rank,
388 + self.num_heads * self.qk_head_dim,
389 + bias=False,
390 + quant_config=quant_config,
391 + prefix=f"{prefix}.q_b_proj",
392 + )
393 + else:
394 + self.q_proj = ColumnParallelLinear(
395 + self.hidden_size,
396 + self.num_heads * self.qk_head_dim,
397 + bias=False,
398 + quant_config=quant_config,
399 + prefix=f"{prefix}.q_proj",
400 + )
401 + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
402 + self.kv_b_proj = ColumnParallelLinear(
403 + self.kv_lora_rank,
404 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
405 + bias=False,
406 + quant_config=quant_config,
407 + prefix=f"{prefix}.kv_b_proj",
408 + )
409 + self.o_proj = RowParallelLinear(
410 + self.num_heads * self.v_head_dim,
411 + self.hidden_size,
412 + bias=False,
413 + quant_config=quant_config,
414 + prefix=f"{prefix}.o_proj",
415 + )
416 + self.rotary_emb = get_rope(
417 + qk_rope_head_dim,
418 + max_position=max_position_embeddings,
419 + rope_parameters=config.rope_parameters,
420 + is_neox_style=False,
421 + )
422 + self.indexer_rope_emb: nn.Module | None
423 + self.indexer: Indexer | None
424 + if self.create_indexer:
425 + # The checkpoint stores indexer q_pe/k_pe in interleaved
426 + # (Megatron/PTM) layout, so the indexer must use interleaved RoPE
427 + # (is_neox_style=False) like the main attention path. Using NeoX
428 + # here loses the relative-position dependence and corrupts the DSA
429 + # top-k selection.
430 + self.indexer_rope_emb = get_rope(
431 + qk_rope_head_dim,
432 + max_position=max_position_embeddings,
433 + rope_parameters=config.rope_parameters,
434 + is_neox_style=False,
435 + )
436 + # The indexer projects its queries from the MLA q_lora activations,
437 + # so a sparse layer requires a query down-projection.
438 + assert q_lora_rank is not None, (
439 + "HYV4 sparse attention requires q_lora_rank to be set"
440 + )
441 + self.indexer = Indexer(
442 + vllm_config,
443 + config,
444 + hidden_size,
445 + q_lora_rank,
446 + quant_config,
447 + cache_config,
448 + topk_indices_buffer,
449 + f"{prefix}.indexer",
450 + )
451 + else:
452 + self.indexer_rope_emb = None
453 + self.indexer = None
454 +
455 + self.gated_mla = bool(getattr(config, "gated_mla", False))
456 + self.linear_gate: ColumnParallelLinear | None
457 + if self.gated_mla:
458 + if config.gating_type == "headwise":
459 + self.gate_projection_size_per_head = 1
460 + elif config.gating_type == "elementwise":
461 + self.gate_projection_size_per_head = self.v_head_dim
462 + else:
463 + raise ValueError(f"Unknown gating type: {config.gating_type}")
464 + self.linear_gate = ColumnParallelLinear(
465 + self.hidden_size,
466 + self.num_heads * self.gate_projection_size_per_head,
467 + bias=False,
468 + quant_config=quant_config,
469 + prefix=f"{prefix}.linear_gate",
470 + )
471 + self.use_hpc_gated_mla = hpc_gated_mla_supported(
472 + config.gating_type, self.linear_gate
473 + )
474 + else:
475 + self.linear_gate = None
476 + self.use_hpc_gated_mla = False
477 + self.prefix = prefix
478 +
479 + # Per-head learnable attention sink. Created BEFORE ``MLAAttention`` so
480 + # it can be forwarded as the ``sinks`` impl kwarg. The parameter always
481 + # holds the local TP shard.
482 + self.learnable_sink = bool(getattr(config, "learnable_sink", False))
483 + sinks = None
484 + sink_backend: type[AttentionBackend] | None = None
485 + if self.learnable_sink:
486 + sink_backend = self._resolve_sink_backend(kv_cache_dtype)
487 + enable_sink = sink_backend is not None
488 + self.learnable_sink_param = nn.Parameter(
489 + torch.empty(
490 + self.num_local_heads,
491 + # The kernels require fp32 sinks; the disabled path keeps
492 + # the checkpoint dtype since the value is never consumed.
493 + dtype=torch.float32 if enable_sink else torch.bfloat16,
494 + )
495 + )
496 + if enable_sink:
497 + sinks = self.learnable_sink_param
498 + self._force_sparse_mqa()
499 +
500 + extra_impl_args = {} if sinks is None else {"sinks": sinks}
501 + self.mla_attn = MLAAttention(
502 + num_heads=self.num_local_heads,
503 + scale=self.scaling,
504 + qk_nope_head_dim=self.qk_nope_head_dim,
505 + qk_rope_head_dim=self.qk_rope_head_dim,
506 + v_head_dim=self.v_head_dim,
507 + q_lora_rank=self.q_lora_rank,
508 + kv_lora_rank=self.kv_lora_rank,
509 + cache_config=cache_config,
510 + quant_config=quant_config,
511 + prefix=f"{prefix}.attn",
512 + kv_b_proj=self.kv_b_proj,
513 + use_sparse=self.is_sparse,
514 + indexer=self.indexer,
515 + topk_indices_buffer=topk_indices_buffer,
516 + attn_backend=sink_backend,
517 + **extra_impl_args,
518 + )
519 +
520 + @property
521 + def topk_indices_buffer(self) -> torch.Tensor | None:
522 + """The sparse top-k index buffer this layer reads from.
523 +
524 + The real consumer is ``mla_attn.impl``, an ``AttentionImpl`` rather than
525 + an ``nn.Module``, so a plain ``named_modules()`` walk cannot reach it.
526 + The MTP proposer rebinds the draft layers onto the target model's buffer
527 + through such a walk, so expose it here: without this the draft attention
528 + would keep reading its own buffer while the indices it needs are written
529 + to the target's.
530 + """
531 + return getattr(self.mla_attn.impl, "topk_indices_buffer", None)
532 +
533 + @topk_indices_buffer.setter
534 + def topk_indices_buffer(self, buffer: torch.Tensor) -> None:
535 + self.mla_attn.impl.topk_indices_buffer = buffer # type: ignore[attr-defined]
536 +
537 + def _resolve_sink_backend(
538 + self, kv_cache_dtype: str
539 + ) -> type[AttentionBackend] | None:
540 + """Return an MLA backend that can apply this layer's learnable sink.
541 +
542 + The sink is part of the architecture, so a backend that cannot apply it
543 + changes the model's output. Resolution order:
544 +
545 + 1. If the backend the selector would pick already advertises
546 + `supports_sink`, keep it — this also honours an explicit
547 + ``--attention-backend`` choice.
548 + 2. Otherwise fall back to the sink-capable ``FLASHMLA_SPARSE`` subclass
549 + in `.flashmla_sparse`, whose kernels accept ``attn_sink``, provided
550 + it validates against the current runtime configuration.
551 + 3. Otherwise give up on the bias rather than failing the load.
552 +
553 + Args:
554 + kv_cache_dtype: The layer's KV cache dtype string.
555 +
556 + Returns:
557 + The backend class to bind, or None when no sink-capable backend is
558 + available; the caller then loads the sink weight but disables the
559 + bias.
560 + """
561 + head_size = self.kv_lora_rank + self.qk_rope_head_dim
562 + dtype = torch.get_default_dtype()
563 + try:
564 + selected_cls = get_attn_backend(
565 + head_size=head_size,
566 + dtype=dtype,
567 + kv_cache_dtype=kv_cache_dtype,
568 + use_mla=True,
569 + use_sparse=self.is_sparse,
570 + num_heads=self.num_local_heads,
571 + )
572 + except Exception as exc:
573 + # Stringify before logging: warning_once dedupes on the arguments,
574 + # and a fresh exception object per layer would defeat it.
575 + logger.warning_once(
576 + "HYV4 failed to select an MLA backend for the learnable sink "
577 + "(%s); the sink parameter is loaded but the sink bias is "
578 + "disabled.",
579 + str(exc),
580 + )
581 + return None
582 +
583 + if selected_cls.supports_sink():
584 + return selected_cls
585 +
586 + from .flashmla_sparse import HYV4FlashMLASparseBackend
587 +
588 + # Mirror how the selector derives the configuration-dependent inputs so
589 + # this check accepts exactly what the backend would accept at runtime.
590 + capability = current_platform.get_device_capability()
591 + if capability is None:
592 + logger.warning_once(
593 + "HYV4 learnable sink is unavailable: the device compute "
594 + "capability is unknown. The sink parameter is loaded but the "
595 + "sink bias is disabled."
596 + )
597 + return None
598 + cache_config = get_current_vllm_config().cache_config
599 + block_size = (
600 + cache_config.block_size
601 + if cache_config is not None and cache_config.user_specified_block_size
602 + else None
603 + )
604 + invalid_reasons = HYV4FlashMLASparseBackend.validate_configuration(
605 + head_size=head_size,
606 + dtype=dtype,
607 + kv_cache_dtype=cast(CacheDType, kv_cache_dtype),
608 + block_size=block_size,
609 + use_mla=True,
610 + has_sink=True,
611 + use_sparse=self.is_sparse,
612 + use_mm_prefix=False,
613 + use_per_head_quant_scales=False,
614 + device_capability=capability,
615 + attn_type=AttentionType.DECODER,
616 + )
617 + if invalid_reasons:
618 + logger.warning_once(
619 + "HYV4 learnable sink is unavailable: the selected backend %s "
620 + "cannot apply sinks and the sink-capable FLASHMLA_SPARSE path "
621 + "is invalid here (%s). The sink parameter is loaded but the "
622 + "sink bias is disabled.",
623 + selected_cls.get_name(),
624 + ", ".join(invalid_reasons),
625 + )
626 + return None
627 +
628 + logger.info_once(
629 + "HYV4 learnable sink enabled: using the sink-capable "
630 + "FLASHMLA_SPARSE impl instead of %s, which cannot apply sinks.",
631 + selected_cls.get_name(),
632 + )
633 + return HYV4FlashMLASparseBackend
634 +
635 + def _force_sparse_mqa(self) -> None:
636 + """Keep every token on the sink-capable sparse MQA path.
637 +
638 + ``_resolve_sink_backend`` only binds the backend that serves decode.
639 + ``MLAAttention`` additionally routes short prefills to a separate dense
640 + MLA prefill backend, and none of those accept ``attn_sink``, so prefill
641 + would silently drop the sink while decode applies it. Such a partially
642 + applied sink is not the trained architecture and corrupts the output, so
643 + opt out of the dense split instead.
644 +
645 + Prefills up to ``index_topk`` keep every token inside the sparse top-k,
646 + making the sparse path numerically equivalent to the dense one apart from
647 + also applying the sink.
648 + """
649 + attention_config = get_current_vllm_config().attention_config
650 + if attention_config.sparse_mla_force_mqa:
651 + return
652 + attention_config.sparse_mla_force_mqa = True
653 + logger.info_once(
654 + "HYV4 learnable sink enabled: forcing sparse MQA for prefill too, "
655 + "as the dense MLA prefill backends cannot apply sinks."
656 + )
657 +
658 + def forward(
659 + self,
660 + positions: torch.Tensor,
661 + hidden_states: torch.Tensor,
662 + llama_4_scaling: torch.Tensor | None = None,
663 + ) -> torch.Tensor:
664 + q_c = None
665 + if self.q_lora_rank is not None:
666 + assert self.fused_qkv_a_proj is not None
667 + assert self.q_a_layernorm is not None
668 + assert self.q_b_proj is not None
669 + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
670 + q_c, kv_lora = qkv_lora.split(
671 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
672 + dim=-1,
673 + )
674 + q_c = self.q_a_layernorm(q_c)
675 + q = self.q_b_proj(q_c)[0]
676 + else:
677 + assert self.q_proj is not None
678 + assert self.kv_a_proj_with_mqa is not None
679 + q = self.q_proj(hidden_states)[0]
680 + kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
681 +
682 + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
683 + kv_c_normed = self.kv_a_layernorm(kv_c)
684 +
685 + q = q.view(-1, self.num_local_heads, self.qk_head_dim)
686 + # Add a head dim of 1 to k_pe.
687 + k_pe = k_pe.unsqueeze(1)
688 + q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
689 + positions, q[..., self.qk_nope_head_dim :], k_pe
690 + )
691 +
692 + if llama_4_scaling is not None:
693 + q *= llama_4_scaling
694 +
695 + output_shape = (
696 + hidden_states.shape[0],
697 + self.num_local_heads * self.v_head_dim,
698 + )
699 + # Single coarse eager break covering the indexer and MLA attention, as
700 + # the breakable cudagraph contract requires: everything that reads
701 + # per-batch metadata runs in one eager segment, so no tensor has to stay
702 + # alive across a capture-segment boundary.
703 + attn_out = torch.empty(
704 + output_shape, dtype=hidden_states.dtype, device=hidden_states.device
705 + )
706 + self._indexer_and_attn(
707 + hidden_states, q_c, positions, q, kv_c_normed, k_pe, attn_out
708 + )
709 +
710 + if self.gated_mla and self.linear_gate is not None:
711 + if self.use_hpc_gated_mla:
712 + # Projection, sigmoid and the product in one launch. The gate
713 + # is column-parallel and unbiased, so its local weight shard
714 + # maps straight onto the local attn_out columns.
715 + assert hidden_states.is_contiguous() and attn_out.is_contiguous()
716 + attn_out = hpc_gated_mla_gemm(
717 + hidden_states,
718 + self.linear_gate.weight,
719 + attn_out,
720 + )
721 + else:
722 + gate_score = self.linear_gate(hidden_states)[0]
723 + if self.config.gating_type == "headwise":
724 + gate_score = gate_score.unsqueeze(-1)
725 + attn_out = attn_out.reshape(
726 + *attn_out.shape[:-1], -1, self.v_head_dim
727 + )
728 + attn_out = attn_out * torch.sigmoid(gate_score)
729 + attn_out = attn_out.reshape(*attn_out.shape[:-2], -1)
730 + else:
731 + attn_out = attn_out * torch.sigmoid(gate_score)
732 +
733 + out, _ = self.o_proj(attn_out)
734 + return out
735 +
736 + @eager_break_during_capture
737 + def _indexer_and_attn(
738 + self,
739 + hidden_states: torch.Tensor,
740 + q_c: torch.Tensor | None,
741 + positions: torch.Tensor,
742 + q: torch.Tensor,
743 + kv_c_normed: torch.Tensor,
744 + k_pe: torch.Tensor,
745 + out: torch.Tensor, # [num_tokens, heads * v_head_dim], written in place
746 + ) -> None:
747 + """Run the lightning indexer and MLA attention in one eager segment.
748 +
749 + Both read per-batch attention metadata, so under the breakable cudagraph
750 + they must not be captured. Keeping them in a single break (instead of one
751 + break each) also means the attention inputs never have to survive a
752 + capture-segment boundary. The nested ``sparse_attn_indexer`` and
753 + ``unified_mla_attention_with_output`` breaks short-circuit here, since
754 + the capture is no longer active inside an eager segment.
755 + """
756 + if self.indexer is not None and self.is_sparse and not self.skip_topk:
757 + self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
758 + out.copy_(
759 + self.mla_attn(
760 + q,
761 + kv_c_normed,
762 + k_pe,
763 + output_shape=out.shape,
764 + )
765 + )
vllm/models/hy_v4/nvidia/flashmla_sparse.py new
+283
@@ -0,0 +1,283 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sink-capable FlashMLA sparse backend for HY V4 (NVIDIA).
4 +
5 +HY V4 adds a per-head learnable attention sink on top of sparse MLA. The
6 +vendored FlashMLA kernels already accept an ``attn_sink`` argument, but vLLM's
7 +shared ``FLASHMLA_SPARSE`` backend neither advertises sink support nor forwards
8 +the tensor, so the bias would be silently dropped.
9 +
10 +This module supplies the missing wiring inside the model package, mirroring how
11 +`vllm.models.deepseek_v4.nvidia.flashinfer_sparse` hands ``sinks`` to the
12 +FlashInfer sparse MLA kernels: subclass the platform backend, declare
13 +`supports_sink`, and thread the sink into every kernel call.
14 +
15 +The subclass intentionally keeps the inherited ``get_name()``
16 +(``"FLASHMLA_SPARSE"``). Several shared code paths key off that exact string —
17 +`_canonicalize_sparse_mla_kv_cache_dtype` promotes a quantized KV cache to
18 +``fp8_ds_mla`` for it, and `FlashMLASparseImpl` asserts that layout — so a new
19 +name would silently change KV cache behaviour. Only``supports_sink`` and the
20 +two kernel wrappers differ from the parent.
21 +"""
22 +
23 +from typing import TYPE_CHECKING
24 +
25 +import torch
26 +
27 +from vllm.logger import init_logger
28 +from vllm.v1.attention.backends.mla.flashmla_sparse import (
29 + FlashMLASparseBackend,
30 + FlashMLASparseImpl,
31 + FlashMLASparseMetadata,
32 +)
33 +from vllm.v1.attention.ops.flashmla import (
34 + flash_mla_sparse_fwd,
35 + flash_mla_with_kvcache,
36 +)
37 +
38 +if TYPE_CHECKING:
39 + from vllm.model_executor.models.deepseek_v2 import Indexer
40 +
41 +logger = init_logger(__name__)
42 +
43 +
44 +class HYV4FlashMLASparseImpl(FlashMLASparseImpl):
45 + """FlashMLA sparse impl that applies HY V4's per-head learnable sink.
46 +
47 + The sink enters as the ``sinks`` impl kwarg of
48 + `vllm.model_executor.layers.attention.MLAAttention` and is consumed by the
49 + FlashMLA kernels, which fold it into the softmax denominator:
50 + ``out *= exp(lse) / (exp(lse) + exp(sink))``.
51 + """
52 +
53 + def __init__(
54 + self,
55 + num_heads: int,
56 + head_size: int,
57 + scale: float,
58 + num_kv_heads: int,
59 + alibi_slopes: list[float] | None,
60 + sliding_window: int | None,
61 + kv_cache_dtype: str,
62 + logits_soft_cap: float | None,
63 + attn_type: str,
64 + kv_sharing_target_layer_name: str | None,
65 + # MLA Specific Arguments
66 + topk_indices_buffer: torch.Tensor | None = None,
67 + indexer: "Indexer | None" = None,
68 + **mla_args,
69 + ) -> None:
70 + # ``SparseMLACommonImpl`` takes explicit keyword arguments only, so the
71 + # sink has to be removed before the base classes see``mla_args``.
72 + sinks: torch.Tensor | None = mla_args.pop("sinks", None)
73 + super().__init__(
74 + num_heads,
75 + head_size,
76 + scale,
77 + num_kv_heads,
78 + alibi_slopes,
79 + sliding_window,
80 + kv_cache_dtype,
81 + logits_soft_cap,
82 + attn_type,
83 + kv_sharing_target_layer_name,
84 + topk_indices_buffer=topk_indices_buffer,
85 + indexer=indexer,
86 + **mla_args,
87 + )
88 + self._validate_sinks(sinks, num_heads)
89 + self.sinks = sinks
90 +
91 + @staticmethod
92 + def _validate_sinks(sinks: torch.Tensor | None, num_heads: int) -> None:
93 + """Reject sink tensors the FlashMLA kernels cannot consume.
94 +
95 + Args:
96 + sinks: Candidate sink tensor, or None when the layer has no sink.
97 + num_heads: Local (TP-sharded) query head count of this layer.
98 +
99 + Raises:
100 + ValueError: If the dtype is not float32 or the shape is not
101 + ``(num_heads,)``.
102 + """
103 + if sinks is None:
104 + return
105 + if sinks.dtype != torch.float32:
106 + raise ValueError(
107 + "HYV4 FlashMLA sparse attention sinks must have dtype "
108 + f"torch.float32, but got {sinks.dtype}."
109 + )
110 + if sinks.ndim != 1 or sinks.shape[0] != num_heads:
111 + raise ValueError(
112 + "HYV4 FlashMLA sparse attention sinks must have shape "
113 + f"({num_heads},), but got {tuple(sinks.shape)}."
114 + )
115 +
116 + def _sinks_for_query(
117 + self,
118 + q: torch.Tensor,
119 + head_dim: int,
120 + kernel_heads: int,
121 + ) -> torch.Tensor | None:
122 + """Return the sink laid out for the kernel's query head count.
123 +
124 + Args:
125 + q: Query tensor, before any head padding.
126 + head_dim: Axis of ``q`` holding the query heads.
127 + kernel_heads: Head count the kernel is invoked with, which may
128 + exceed the query head count because of padding.
129 +
130 + Returns:
131 + The sink tensor padded to ``kernel_heads`` with ``-inf`` (a no-op
132 + sink) for the padded lanes, or None when the layer has no sink.
133 +
134 + Raises:
135 + ValueError: If the sink and query head layouts disagree, or if they
136 + live on different devices.
137 + """
138 + sinks = self.sinks
139 + if sinks is None:
140 + return None
141 +
142 + query_heads = q.shape[head_dim]
143 + if sinks.shape[0] != query_heads:
144 + raise ValueError(
145 + "HYV4 FlashMLA sparse attention sink head count must match the "
146 + f"runtime query layout: sinks={sinks.shape[0]}, "
147 + f"query_heads={query_heads}. The sink must use the same "
148 + "unpadded head layout as the query."
149 + )
150 + if sinks.device != q.device:
151 + raise ValueError(
152 + "HYV4 FlashMLA sparse attention sinks and query must be on the "
153 + f"same device, but got sinks={sinks.device}, query={q.device}."
154 + )
155 + if kernel_heads < query_heads:
156 + raise ValueError(
157 + "HYV4 FlashMLA sparse kernel head count cannot be smaller than "
158 + f"the runtime query layout: query_heads={query_heads}, "
159 + f"kernel_heads={kernel_heads}."
160 + )
161 + if kernel_heads == query_heads:
162 + return sinks
163 +
164 + # Mirror the query padding the kernels require. Reading ``sinks`` here
165 + # (rather than caching a padded copy at construction time) keeps the
166 + # values correct for weights loaded after the module is built, and the
167 + # allocation plus copy are captured in the CUDA graph.
168 + padded_sinks = sinks.new_full((kernel_heads,), float("-inf"))
169 + padded_sinks[:query_heads] = sinks
170 + return padded_sinks
171 +
172 + def _fp8_flash_mla_kernel(
173 + self,
174 + q: torch.Tensor,
175 + kv_c_and_k_pe_cache: torch.Tensor,
176 + topk_indices: torch.Tensor,
177 + kernel_metadata: FlashMLASparseMetadata.FP8KernelMetadata,
178 + ) -> tuple[torch.Tensor, torch.Tensor]:
179 + # q shape: (batch, seq_len, num_heads, head_dim)
180 + actual_num_heads = q.size(2)
181 + padded_num_heads = self.fp8_decode_padded_heads
182 + attn_sink = self._sinks_for_query(q, head_dim=2, kernel_heads=padded_num_heads)
183 +
184 + # Pad query if needed (kernel only supports h_q = 64 or 128)
185 + if actual_num_heads < padded_num_heads:
186 + logger.warning_once(
187 + f"Padding num_heads from {actual_num_heads} to "
188 + f"{padded_num_heads} for FP8 sparse decode kernel"
189 + )
190 + q_padded = q.new_zeros((q.size(0), q.size(1), padded_num_heads, q.size(3)))
191 + q_padded[:, :, :actual_num_heads, :] = q
192 + q = q_padded
193 +
194 + out, lse = flash_mla_with_kvcache(
195 + q=q,
196 + k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),
197 + block_table=kernel_metadata.dummy_block_table,
198 + head_dim_v=512,
199 + cache_seqlens=kernel_metadata.cache_lens,
200 + tile_scheduler_metadata=kernel_metadata.scheduler_metadata,
201 + is_fp8_kvcache=True,
202 + indices=topk_indices,
203 + softmax_scale=self.softmax_scale,
204 + attn_sink=attn_sink,
205 + )
206 +
207 + # Slice output and lse back to actual head count if we padded
208 + if actual_num_heads < padded_num_heads:
209 + out = out[:, :, :actual_num_heads, :]
210 + lse = lse[:, :actual_num_heads, :]
211 +
212 + return out, lse
213 +
214 + def _bf16_flash_mla_kernel(
215 + self,
216 + q: torch.Tensor,
217 + kv_c_and_k_pe_cache: torch.Tensor,
218 + topk_indices: torch.Tensor,
219 + topk_length: torch.Tensor | None = None,
220 + actual_num_heads: int | None = None,
221 + ) -> tuple[torch.Tensor, torch.Tensor]:
222 + num_tokens = q.shape[0]
223 + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
224 + -1, 1, kv_c_and_k_pe_cache.shape[-1]
225 + )
226 +
227 + # NOTE(Chen): kernel requires num_local_head to be a multiple of
228 + # 64 on hopper and 128 on blackwell. Pad from q's head count, not
229 + # self.num_heads: under DCP the heads are all-gathered before this.
230 + if actual_num_heads is None:
231 + actual_num_heads = q.shape[1]
232 + padded_num_heads = (
233 + (actual_num_heads + self.prefill_padding - 1)
234 + // self.prefill_padding
235 + * self.prefill_padding
236 + )
237 + # q_concat_buffer may already include the kernel-required padding.
238 + # The sink remains defined only for the real query heads.
239 + attn_sink = self._sinks_for_query(
240 + q[:, :actual_num_heads], head_dim=1, kernel_heads=padded_num_heads
241 + )
242 +
243 + if q.shape[1] < padded_num_heads:
244 + logger.warning_once(
245 + f"Padding num_heads from {actual_num_heads} to "
246 + f"{padded_num_heads} for BF16 sparse prefill kernel"
247 + )
248 + # Zero (not new_empty) the padded lanes: topk_indices is shared by
249 + # all heads, so the kernel reduces across the head group and NaNs
250 + # from uninitialized memory would leak into the real heads.
251 + q_padded = q.new_zeros((q.shape[0], padded_num_heads, q.shape[2]))
252 + q_padded[:, :actual_num_heads, :] = q
253 + q = q_padded
254 +
255 + topk_indices = topk_indices.view(num_tokens, 1, -1)
256 + output, _, lse = flash_mla_sparse_fwd(
257 + q,
258 + kv_c_and_k_pe_cache,
259 + topk_indices,
260 + self.softmax_scale,
261 + attn_sink=attn_sink,
262 + topk_length=topk_length,
263 + )
264 +
265 + output = output[:, :actual_num_heads, :]
266 + lse = lse[:, :actual_num_heads]
267 + return output, lse
268 +
269 +
270 +class HYV4FlashMLASparseBackend(FlashMLASparseBackend):
271 + """``FLASHMLA_SPARSE`` with attention-sink support for HY V4.
272 +
273 + Keeps the parent's name, metadata and builder; only the impl class and the
274 + sink capability differ. See the module docstring for why the name is reused.
275 + """
276 +
277 + @staticmethod
278 + def get_impl_cls() -> type[HYV4FlashMLASparseImpl]:
279 + return HYV4FlashMLASparseImpl
280 +
281 + @classmethod
282 + def supports_sink(cls) -> bool:
283 + return True
vllm/models/hy_v4/nvidia/hc.py new
+370
@@ -0,0 +1,370 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""iHC (independent Hyper-Connections) layers for HY V4 (NVIDIA).
4 +
5 +iHC replaces the single residual stream of a standard transformer with
6 +``hc_mult`` parallel residual channels. Each decoder sub-block reduces the
7 +channels to one hidden state (``HYV4HCPreLayer``), runs the sub-block, then
8 +scatters the result back over the channels (``HYV4HCPostLayer``). The final
9 +``HYV4HCHeadLayer`` merges the channels before the model's output norm.
10 +
11 +NOTE: Each of the three steps has an optional single-kernel HPC replacement
12 +(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). They are only constructed
13 +when the hpc package is installed, ``VLLM_ENABLE_HPC_OPS=1`` and the shape /
14 +device constraints hold; otherwise the eager path below runs unchanged.
15 +TODO: port the cross-layer post+pre fusion (``HpcIHCPostPre``) as well; it
16 +requires restructuring the decoder-layer forward scheduling.
17 +"""
18 +
19 +import torch
20 +from torch import nn
21 +from transformers import PretrainedConfig
22 +
23 +from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre
24 +from vllm.model_executor.layers.linear import ReplicatedLinear
25 +
26 +
27 +class HYV4HCPreLayer(nn.Module):
28 + """iHC pre-processing layer (2D-activation adaptation).
29 +
30 + Steps:
31 + 1. RMS-normalize the flattened ``[num_tokens, hc * d]`` input.
32 + 2. Project to the pre/post gating logits.
33 + 3. Turn the logits into sigmoid gates.
34 + 4. Reduce over the channel dim with the pre gates.
35 + """
36 +
37 + def __init__(
38 + self,
39 + config: PretrainedConfig,
40 + hidden_dim: int,
41 + hc_mult: int = 4,
42 + magnitude: float = 2.0,
43 + init_std: float = 6e-3,
44 + base_noise_std: float = 0.0,
45 + hc_eps: float = 1e-6,
46 + layernorm_epsilon: float = 1e-5,
47 + prefix: str = "",
48 + ):
49 + super().__init__()
50 + self.config = config
51 + self.hidden_dim = hidden_dim
52 + self.hc_mult = hc_mult
53 + self.magnitude = magnitude
54 + self.hc_eps = hc_eps
55 + self.layernorm_epsilon = layernorm_epsilon
56 + hc_in_dim = hc_mult * hidden_dim
57 + mix_hc = 2 * hc_mult # pre + post only, no comb
58 +
59 + self.hc_fn = ReplicatedLinear(
60 + input_size=hc_in_dim,
61 + output_size=mix_hc,
62 + params_dtype=torch.float32,
63 + bias=False,
64 + prefix=f"{prefix}.hc_fn",
65 + )
66 + self.hc_scale = nn.Parameter(torch.empty(2, dtype=torch.float32))
67 + self.hc_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
68 +
69 + self.reset_parameters(init_std, base_noise_std)
70 +
71 + # Optional single-kernel HPC replacement for the whole forward below.
72 + # ``norm_owner`` is left unset: the RMSNorm that follows the pre block
73 + # lives in the decoder layer and is not folded in yet.
74 + self.hpc_op: HpcIHCPre | None = None
75 + if HpcIHCPre.support(hc_mult, hidden_dim):
76 + self.hpc_op = HpcIHCPre(
77 + hc_mult=hc_mult,
78 + hidden_size=hidden_dim,
79 + magnitude=magnitude,
80 + hc_eps=hc_eps,
81 + norm_eps=layernorm_epsilon,
82 + fallback_op=self,
83 + )
84 +
85 + def reset_parameters(self, init_std: float, base_noise_std: float = 0.0) -> None:
86 + """Initialize the gate scale and per-channel gate bias."""
87 + del init_std # hc_fn is initialized by ReplicatedLinear
88 + nn.init.constant_(self.hc_scale, 0.01)
89 + with torch.no_grad():
90 + self.hc_base[: self.hc_mult].fill_(
91 + -torch.log(torch.tensor(self.hc_mult - 1.0, dtype=self.hc_base.dtype))
92 + )
93 + self.hc_base[self.hc_mult : 2 * self.hc_mult].fill_(0.0)
94 + if base_noise_std > 0.0:
95 + self.hc_base.add_(torch.randn_like(self.hc_base) * base_noise_std)
96 +
97 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
98 + """Reduce the iHC channels and emit the post gates.
99 +
100 + Args:
101 + x: Input of shape ``[num_tokens, hc, d]``.
102 +
103 + Returns:
104 + A tuple of the pre-gated reduction ``[num_tokens, d]`` and the post
105 + gates ``[num_tokens, hc]`` consumed by `HYV4HCPostLayer`.
106 + """
107 + if self.hpc_op is not None:
108 + return self.hpc_op(x)
109 +
110 + shape = x.size() # [num_tokens, hc, d]
111 + hc = self.hc_mult
112 + hc_eps = self.hc_eps
113 +
114 + x_flat = x.flatten(1).float() # [num_tokens, hc*d]
115 +
116 + rsqrt = torch.rsqrt(
117 + x_flat.square().mean(-1, keepdim=True) + self.layernorm_epsilon
118 + )
119 + mixes = self.hc_fn(x_flat)[0] * rsqrt # [num_tokens, 2*hc]
120 +
121 + pre_raw = mixes[..., :hc]
122 + post_raw = mixes[..., hc : 2 * hc]
123 +
124 + pre = (
125 + torch.sigmoid(
126 + pre_raw * self.hc_scale[0].float() + self.hc_base[:hc].float()
127 + )
128 + + hc_eps
129 + )
130 + post = (
131 + self.magnitude
132 + * torch.sigmoid(
133 + post_raw * self.hc_scale[1].float() + self.hc_base[hc : 2 * hc].float()
134 + )
135 + + hc_eps
136 + )
137 +
138 + y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1) # [num_tokens, d]
139 + return y.to(x.dtype), post
140 +
141 +
142 +class HYV4HCPostLayer(nn.Module):
143 + """iHC post-processing layer (2D-activation adaptation).
144 +
145 + Applies post-gating to the sub-block output and adds the multi-channel
146 + residual (no comb mixing)::
147 +
148 + y[n, i, d] = post[n, i] * x[n, d] + residual[n, i, d]
149 + """
150 +
151 + def __init__(self, config: PretrainedConfig):
152 + super().__init__()
153 + self.config = config
154 +
155 + # Optional single-kernel HPC replacement for the whole forward below.
156 + # Only constructed under enable_ihc (see HYV4HCLayer), so hc_mult is
157 + # present; getattr keeps this robust if that ever changes.
158 + self.hpc_op: HpcIHCPost | None = None
159 + hc_mult = getattr(config, "hc_mult", 0)
160 + if HpcIHCPost.support(hc_mult, config.hidden_size):
161 + self.hpc_op = HpcIHCPost(hc_mult=hc_mult, hidden_size=config.hidden_size)
162 +
163 + def forward(
164 + self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor
165 + ) -> torch.Tensor:
166 + """Scatter the sub-block output back onto the iHC channels.
167 +
168 + Args:
169 + x: Attention/MLP output of shape ``[num_tokens, d]``.
170 + residual: Multi-channel residual ``[num_tokens, hc, d]``.
171 + post: Post gates ``[num_tokens, hc]`` from `HYV4HCPreLayer`.
172 +
173 + Returns:
174 + The updated residual channels ``[num_tokens, hc, d]``.
175 + """
176 + if self.hpc_op is not None:
177 + return self.hpc_op(x, residual, post)
178 +
179 + dtype = x.dtype
180 + x = x.float()
181 + residual = residual.float()
182 + post = post.float()
183 +
184 + post_gated = post.unsqueeze(-1) * x.unsqueeze(-2) # [num_tokens, hc, d]
185 + y = post_gated + residual
186 + return y.to(dtype)
187 +
188 +
189 +class HYV4HCHeadLayer(nn.Module):
190 + """iHC head layer (2D-activation adaptation).
191 +
192 + Merges the iHC channels back into a single hidden state before the final
193 + layer norm, using an RMS-normed projection plus sigmoid-gated reduction.
194 + """
195 +
196 + def __init__(
197 + self,
198 + config: PretrainedConfig,
199 + hidden_size: int,
200 + hc_mult: int = 4,
201 + hc_eps: float = 1e-6,
202 + init_std: float = 6e-3,
203 + base_noise_std: float = 0.0,
204 + prefix: str = "",
205 + ):
206 + super().__init__()
207 + self.config = config
208 + self.hidden_size = hidden_size
209 + self.hc_mult = hc_mult
210 + self.hc_eps = hc_eps
211 + self.hc_head_fn = ReplicatedLinear(
212 + input_size=hc_mult * hidden_size,
213 + output_size=hc_mult,
214 + params_dtype=torch.float32,
215 + bias=False,
216 + prefix=f"{prefix}.hc_head_fn",
217 + )
218 + self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32))
219 + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32))
220 +
221 + self.reset_parameters(init_std, base_noise_std)
222 +
223 + # Optional single-kernel HPC replacement for the whole forward below.
224 + self.hpc_op: HpcIHCHead | None = None
225 + if HpcIHCHead.support(hc_mult, hidden_size):
226 + self.hpc_op = HpcIHCHead(
227 + hc_mult=hc_mult,
228 + hidden_size=hidden_size,
229 + hc_eps=hc_eps,
230 + norm_eps=config.rms_norm_eps,
231 + fallback_op=self,
232 + )
233 +
234 + def reset_parameters(
235 + self, init_std: float = 6e-3, base_noise_std: float = 0.0
236 + ) -> None:
237 + """Initialize the head gate scale and per-channel gate bias."""
238 + del init_std # hc_head_fn is initialized by ReplicatedLinear
239 + nn.init.constant_(self.hc_head_scale, 0.01)
240 + with torch.no_grad():
241 + self.hc_head_base.fill_(
242 + -torch.log(
243 + torch.tensor(self.hc_mult - 1.0, dtype=self.hc_head_base.dtype)
244 + )
245 + )
246 + if base_noise_std > 0.0:
247 + self.hc_head_base.add_(
248 + torch.randn_like(self.hc_head_base) * base_noise_std
249 + )
250 +
251 + def forward(self, x: torch.Tensor) -> torch.Tensor:
252 + """Merge the iHC channels into a single hidden state.
253 +
254 + Args:
255 + x: Input of shape ``[num_tokens, hc, d]``.
256 +
257 + Returns:
258 + The merged hidden state ``[num_tokens, d]``.
259 + """
260 + if self.hpc_op is not None:
261 + return self.hpc_op(x)
262 +
263 + shape, x_dtype = x.size(), x.dtype
264 +
265 + x = x.flatten(1).float() # [num_tokens, hc*d]
266 + rsqrt = torch.rsqrt(
267 + x.square().mean(-1, keepdim=True) + self.config.rms_norm_eps
268 + )
269 + mixes = self.hc_head_fn(x)[0] * rsqrt # [num_tokens, hc]
270 + pre = (
271 + torch.sigmoid(
272 + mixes * self.hc_head_scale.float() + self.hc_head_base.float()
273 + )
274 + + self.hc_eps
275 + )
276 + y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1) # [num_tokens, d]
277 + return y.to(x_dtype)
278 +
279 +
280 +class HYV4HCLayer(nn.Module):
281 + """Wrapper owning one iHC boundary (pre + post) of a decoder sub-block."""
282 +
283 + def __init__(
284 + self,
285 + config: PretrainedConfig,
286 + layer_idx: int,
287 + init_std: float = 6e-3,
288 + base_noise_std: float = 0.0,
289 + prefix: str = "",
290 + ):
291 + super().__init__()
292 + self.config = config
293 + self.layer_idx = layer_idx
294 + self.enable_ihc = getattr(config, "enable_ihc", False)
295 + if self.enable_ihc:
296 + self.hc_pre = HYV4HCPreLayer(
297 + config,
298 + config.hidden_size,
299 + config.hc_mult,
300 + config.hc_magnitude,
301 + init_std,
302 + base_noise_std,
303 + config.hc_eps,
304 + config.rms_norm_eps,
305 + prefix=f"{prefix}.hc_pre",
306 + )
307 + self.hc_post = HYV4HCPostLayer(config)
308 +
309 + def prepare_input(self, hidden_states: torch.Tensor) -> torch.Tensor:
310 + """Normalize the sub-block input to 3D when iHC is enabled."""
311 + if not self.enable_ihc:
312 + return hidden_states
313 + return self._prepare_input_to_3d(hidden_states)
314 +
315 + def _prepare_input_to_3d(self, hidden_states: torch.Tensor) -> torch.Tensor:
316 + """Reshape the iHC input to ``[num_tokens, hc, h]``.
317 +
318 + Accepted inputs are ``[num_tokens, hc, h]`` (no-op), ``[num_tokens, h]``
319 + (broadcast over the channels) and ``[num_tokens, hc * h]`` (reshape).
320 + """
321 + if hidden_states.dim() == 3:
322 + return hidden_states
323 + if hidden_states.dim() != 2:
324 + raise RuntimeError(
325 + f"HC expects a 2D/3D tensor, got shape={tuple(hidden_states.shape)}"
326 + )
327 +
328 + d0, d1 = hidden_states.shape
329 + h = self.config.hidden_size
330 + hc = self.config.hc_mult
331 +
332 + if d1 == h:
333 + return hidden_states.unsqueeze(1).repeat(1, hc, 1)
334 +
335 + expected = hc * h
336 + if d1 == expected:
337 + return hidden_states.reshape(d0, hc, h)
338 +
339 + raise RuntimeError(
340 + f"HC expects last dim to be hidden_size ({h})"
341 + f"or hc_mult*hidden_size ({expected}), got {d1}. "
342 + f"hc_mult={hc}, hidden_size={h}."
343 + )
344 +
345 + def pre(
346 + self, hidden_states: torch.Tensor
347 + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
348 + """Reduce the iHC channels and produce the post gates.
349 +
350 + Returns:
351 + A tuple of the reduced hidden states ``[num_tokens, d]``, the post
352 + gates ``[num_tokens, hc]`` (``None`` when iHC is disabled) and the
353 + residual (the untouched input).
354 + """
355 + if not self.enable_ihc:
356 + return hidden_states, None, hidden_states
357 + reduced, post_gates = self.hc_pre(hidden_states)
358 + return reduced, post_gates, hidden_states
359 +
360 + def post(
361 + self,
362 + output_with_bias: torch.Tensor,
363 + residual: torch.Tensor,
364 + post_gates: torch.Tensor | None,
365 + ) -> torch.Tensor:
366 + """Apply post-gating and add the residual."""
367 + if not self.enable_ihc:
368 + return output_with_bias + residual
369 + assert post_gates is not None
370 + return self.hc_post(output_with_bias, residual, post_gates)
vllm/models/hy_v4/nvidia/model.py new
+716
@@ -0,0 +1,716 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# Copyright 2023 The vLLM team.
6 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 +#
8 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 +# and OPT implementations in this library. It has been modified from its
10 +# original forms to accommodate minor architectural differences compared
11 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 +#
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +"""Inference-only HY V4 model compatible with HuggingFace weights (NVIDIA)."""
25 +
26 +import typing
27 +from collections.abc import Callable, Iterable, MutableSequence, Sequence
28 +from itertools import islice
29 +
30 +import regex as re
31 +import torch
32 +from torch import nn
33 +from transformers import PretrainedConfig
34 +
35 +from vllm.config import CacheConfig, VllmConfig
36 +from vllm.distributed import (
37 + get_pp_group,
38 + get_tensor_model_parallel_rank,
39 + get_tensor_model_parallel_world_size,
40 +)
41 +from vllm.logger import init_logger
42 +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
43 +from vllm.model_executor.layers.layernorm import RMSNorm
44 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
45 +from vllm.model_executor.layers.quantization import QuantizationConfig
46 +from vllm.model_executor.layers.vocab_parallel_embedding import (
47 + ParallelLMHead,
48 + VocabParallelEmbedding,
49 +)
50 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51 +from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk
52 +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
53 +from vllm.model_executor.models.utils import (
54 + AutoWeightsLoader,
55 + PPMissingLayer,
56 + WeightsMapper,
57 + get_pp_missing_layer_names,
58 + is_pp_missing_parameter,
59 + make_empty_intermediate_tensors_factory,
60 + make_layers,
61 + maybe_prefix,
62 +)
63 +from vllm.platforms import current_platform
64 +from vllm.sequence import IntermediateTensors
65 +
66 +from .attention import (
67 + HYV4MLAAttention,
68 + compute_skip_topk_layers,
69 + is_skip_topk_indexer_weight,
70 +)
71 +from .hc import HYV4HCHeadLayer, HYV4HCLayer
72 +from .moe import HYV4FeedForward, HYV4MoEFused
73 +
74 +logger = init_logger(__name__)
75 +
76 +
77 +def _normalize_hyv4_config(config: PretrainedConfig) -> PretrainedConfig:
78 + """Populate the aliases consumed by the shared MoE implementation."""
79 + config.router_scaling_factor = config.routed_scaling_factor
80 + config.num_experts = config.n_routed_experts
81 + config.expert_hidden_dim = config.moe_intermediate_size
82 + config.num_shared_experts = config.n_shared_experts
83 + config.route_norm = config.norm_topk_prob
84 + return config
85 +
86 +
87 +class HYV4DecoderLayer(nn.Module):
88 + """One HY V4 decoder layer: MLA attention plus a dense or MoE MLP.
89 +
90 + When``config.enable_ihc`` is set the layer runs on ``hc_mult`` residual
91 + channels and each sub-block is wrapped by an `HYV4HCLayer` boundary;
92 + otherwise it uses the standard single-stream residual.
93 + """
94 +
95 + def __init__(
96 + self,
97 + config: PretrainedConfig,
98 + vllm_config: VllmConfig,
99 + cache_config: CacheConfig | None = None,
100 + quant_config: QuantizationConfig | None = None,
101 + prefix: str = "",
102 + topk_indices_buffer: torch.Tensor | None = None,
103 + ) -> None:
104 + super().__init__()
105 + self.hidden_size = config.hidden_size
106 + self.config = config
107 + self.enable_ihc = getattr(config, "enable_ihc", False)
108 + layer_idx = int(prefix.split(".")[-1])
109 + self.layer_idx = layer_idx
110 +
111 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
112 + self.self_attn = HYV4MLAAttention(
113 + vllm_config=vllm_config,
114 + config=config,
115 + hidden_size=self.hidden_size,
116 + num_heads=config.num_attention_heads,
117 + qk_nope_head_dim=config.qk_nope_head_dim,
118 + qk_rope_head_dim=config.qk_rope_head_dim,
119 + v_head_dim=config.v_head_dim,
120 + q_lora_rank=config.q_lora_rank,
121 + kv_lora_rank=config.kv_lora_rank,
122 + max_position_embeddings=max_position_embeddings,
123 + cache_config=cache_config,
124 + quant_config=quant_config,
125 + prefix=f"{prefix}.self_attn",
126 + layer_idx=layer_idx,
127 + topk_indices_buffer=topk_indices_buffer,
128 + )
129 + self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
130 + if config.mlp_layer_types[layer_idx] == "dense":
131 + self.mlp = HYV4FeedForward(
132 + hidden_size=config.hidden_size,
133 + intermediate_size=config.intermediate_size,
134 + hidden_act=config.hidden_act,
135 + quant_config=quant_config,
136 + prefix=f"{prefix}.mlp",
137 + )
138 + self.block_type = "feedforward"
139 + else:
140 + self.mlp = HYV4MoEFused(
141 + config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
142 + )
143 + self.block_type = "moe"
144 + self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
145 + self.hc_attn_layer = HYV4HCLayer(
146 + config, layer_idx, prefix=f"{prefix}.hc_attn_layer"
147 + )
148 + self.hc_mlp_layer = HYV4HCLayer(
149 + config, layer_idx, prefix=f"{prefix}.hc_mlp_layer"
150 + )
151 +
152 + def forward(
153 + self,
154 + positions: torch.Tensor,
155 + hidden_states: torch.Tensor,
156 + residual: torch.Tensor | None,
157 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
158 + if self.enable_ihc:
159 + return self._forward_ihc(positions, hidden_states)
160 + return self._forward_normal(positions, hidden_states, residual)
161 +
162 + def _forward_normal(
163 + self,
164 + positions: torch.Tensor,
165 + hidden_states: torch.Tensor,
166 + residual: torch.Tensor | None,
167 + ) -> tuple[torch.Tensor, torch.Tensor]:
168 + """Standard single-residual-stream forward (iHC disabled)."""
169 + if residual is not None:
170 + hidden_states = hidden_states + residual
171 + residual = hidden_states
172 + hidden_states = self.input_layernorm(hidden_states)
173 +
174 + hidden_states = self.self_attn(
175 + positions=positions,
176 + hidden_states=hidden_states,
177 + )
178 +
179 + hidden_states = hidden_states + residual
180 + residual = hidden_states
181 + hidden_states = self.post_attention_layernorm(hidden_states)
182 +
183 + hidden_states = self.mlp(hidden_states)
184 +
185 + return hidden_states, residual
186 +
187 + def _forward_ihc(
188 + self,
189 + positions: torch.Tensor,
190 + hidden_states: torch.Tensor,
191 + ) -> tuple[torch.Tensor, None]:
192 + """iHC forward: each sub-block reduces and re-scatters the channels."""
193 + hidden_states = self.hc_attn_layer.prepare_input(hidden_states)
194 + hidden_states, post_gates, residual = self.hc_attn_layer.pre(hidden_states)
195 + hidden_states = self.input_layernorm(hidden_states)
196 + hidden_states = self.self_attn(
197 + positions=positions,
198 + hidden_states=hidden_states,
199 + )
200 + hidden_states = self.hc_attn_layer.post(hidden_states, residual, post_gates)
201 +
202 + hidden_states = self.hc_mlp_layer.prepare_input(hidden_states)
203 + hidden_states, post_gates, residual = self.hc_mlp_layer.pre(hidden_states)
204 + hidden_states = self.post_attention_layernorm(hidden_states)
205 + hidden_states = self.mlp(hidden_states)
206 + hidden_states = self.hc_mlp_layer.post(hidden_states, residual, post_gates)
207 +
208 + # Under iHC the residual is carried inside hidden_states.
209 + return hidden_states, None
210 +
211 +
212 +class HYV4Model(nn.Module):
213 + """HY V4 backbone."""
214 +
215 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
216 + super().__init__()
217 +
218 + config = _normalize_hyv4_config(vllm_config.model_config.hf_config)
219 + cache_config = vllm_config.cache_config
220 + quant_config = vllm_config.quant_config
221 + parallel_config = vllm_config.parallel_config
222 + eplb_config = parallel_config.eplb_config
223 + self.num_redundant_experts = eplb_config.num_redundant_experts
224 + self.device = current_platform.device_type
225 + self.vocab_size = config.vocab_size
226 + self.is_sparse = hasattr(config, "index_topk")
227 + if self.is_sparse:
228 + self.topk_indices_buffer = torch.empty(
229 + vllm_config.scheduler_config.max_num_batched_tokens,
230 + config.index_topk,
231 + dtype=torch.int32,
232 + device=self.device,
233 + )
234 + else:
235 + self.topk_indices_buffer = None
236 + self.config = config
237 + self.quant_config = quant_config
238 + self.enable_ihc = getattr(config, "enable_ihc", False)
239 +
240 + if get_pp_group().is_first_rank or (
241 + self.config.tie_word_embeddings and get_pp_group().is_last_rank
242 + ):
243 + self.embed_tokens = VocabParallelEmbedding(
244 + config.vocab_size,
245 + config.hidden_size,
246 + )
247 + else:
248 + self.embed_tokens = PPMissingLayer()
249 + self.start_layer, self.end_layer, self.layers = make_layers(
250 + config.num_hidden_layers,
251 + lambda prefix: HYV4DecoderLayer(
252 + config=config,
253 + vllm_config=vllm_config,
254 + cache_config=cache_config,
255 + quant_config=quant_config,
256 + prefix=prefix,
257 + topk_indices_buffer=self.topk_indices_buffer,
258 + ),
259 + prefix=f"{prefix}.layers",
260 + )
261 + if get_pp_group().is_last_rank:
262 + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
263 + else:
264 + self.norm = PPMissingLayer()
265 + if self.enable_ihc:
266 + # iHC head layer: merge the residual channels into one stream.
267 + if get_pp_group().is_last_rank:
268 + self.hc_head = HYV4HCHeadLayer(
269 + config,
270 + hidden_size=config.hidden_size,
271 + hc_mult=config.hc_mult,
272 + hc_eps=config.hc_eps,
273 + prefix=f"{prefix}.hc_head",
274 + )
275 + else:
276 + self.hc_head = PPMissingLayer()
277 + self.make_empty_intermediate_tensors = (
278 + make_empty_intermediate_tensors_factory(
279 + ["hidden_states"], config.hc_mult * config.hidden_size
280 + )
281 + )
282 + else:
283 + self.make_empty_intermediate_tensors = (
284 + make_empty_intermediate_tensors_factory(
285 + ["hidden_states", "residual"], config.hidden_size
286 + )
287 + )
288 +
289 + # MoE hyperparameters (consumed by EPLB).
290 + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = []
291 + self.num_expert_groups = 1
292 + self.moe_layers: list[nn.Module] = []
293 + example_layer: HYV4MoEFused | None = None
294 + for layer in self.layers:
295 + if isinstance(layer, PPMissingLayer):
296 + continue
297 +
298 + assert isinstance(layer, HYV4DecoderLayer)
299 + if layer.block_type == "moe":
300 + assert isinstance(layer.mlp, HYV4MoEFused)
301 + example_layer = layer.mlp
302 + self.moe_layers.append(layer.mlp.experts)
303 +
304 + if example_layer is None:
305 + self.num_moe_layers = 0
306 + self.num_logical_experts = 0
307 + self.num_physical_experts = 0
308 + self.num_local_physical_experts = 0
309 + self.num_routed_experts = 0
310 + self.num_redundant_experts = 0
311 + return
312 +
313 + self.num_moe_layers = len(self.moe_layers)
314 + self.num_logical_experts = example_layer.n_logical_experts
315 + self.num_physical_experts = example_layer.n_physical_experts
316 + self.num_local_physical_experts = example_layer.n_local_physical_experts
317 + self.num_routed_experts = example_layer.n_routed_experts
318 + self.num_redundant_experts = example_layer.n_redundant_experts
319 +
320 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
321 + return self.embed_tokens(input_ids)
322 +
323 + def update_physical_experts_metadata(
324 + self,
325 + num_physical_experts: int,
326 + num_local_physical_experts: int,
327 + ) -> None:
328 + assert self.num_local_physical_experts == num_local_physical_experts
329 + self.num_physical_experts = num_physical_experts
330 + self.num_local_physical_experts = num_local_physical_experts
331 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
332 + for layer in self.layers:
333 + if isinstance(layer, PPMissingLayer):
334 + continue
335 + if isinstance(layer.mlp, HYV4MoEFused):
336 + moe = layer.mlp
337 + moe.n_local_physical_experts = num_local_physical_experts
338 + moe.n_physical_experts = num_physical_experts
339 + moe.n_redundant_experts = self.num_redundant_experts
340 + moe.experts.update_expert_map()
341 +
342 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
343 + # (param_name, weight_name, expert_id, shard_id) for weights, fp8
344 + # weight scales and fp8 activation scales.
345 + if not hasattr(self, "_cached_expert_params_mapping"):
346 + self._cached_expert_params_mapping = fused_moe_make_expert_params_mapping(
347 + self,
348 + ckpt_gate_proj_name="gate_proj",
349 + ckpt_down_proj_name="down_proj",
350 + ckpt_up_proj_name="up_proj",
351 + num_experts=self.config.num_experts,
352 + num_redundant_experts=self.num_redundant_experts,
353 + )
354 + return self._cached_expert_params_mapping
355 +
356 + def forward(
357 + self,
358 + input_ids: torch.Tensor,
359 + positions: torch.Tensor,
360 + intermediate_tensors: IntermediateTensors | None = None,
361 + inputs_embeds: torch.Tensor | None = None,
362 + ) -> torch.Tensor | IntermediateTensors:
363 + if get_pp_group().is_first_rank:
364 + if inputs_embeds is not None:
365 + hidden_states = inputs_embeds
366 + else:
367 + hidden_states = self.embed_input_ids(input_ids)
368 + residual = None
369 + else:
370 + assert intermediate_tensors is not None
371 + hidden_states = intermediate_tensors["hidden_states"]
372 + # In iHC mode the flattened [num_tokens, hc*h] tensor from the
373 + # previous PP stage is reshaped back to 3D by the first layer's
374 + # prepare_input, and residual is unused.
375 + residual = None if self.enable_ihc else intermediate_tensors["residual"]
376 +
377 + for layer in islice(self.layers, self.start_layer, self.end_layer):
378 + hidden_states, residual = layer(positions, hidden_states, residual)
379 +
380 + if not get_pp_group().is_last_rank:
381 + if self.enable_ihc:
382 + # hidden_states is [num_tokens, hc, h]; flatten the channel dim
383 + # for PP transfer (matches the 2D receive buffer).
384 + return IntermediateTensors({"hidden_states": hidden_states.flatten(1)})
385 + return IntermediateTensors(
386 + {"hidden_states": hidden_states, "residual": residual}
387 + )
388 + if self.enable_ihc:
389 + hidden_states = self.hc_head(hidden_states)
390 + else:
391 + hidden_states = hidden_states + residual
392 +
393 + return self.norm(hidden_states)
394 +
395 + def load_fused_expert_weights(
396 + self,
397 + name: str,
398 + params_dict: dict,
399 + loaded_weight: torch.Tensor,
400 + shard_id: str,
401 + num_experts: int,
402 + ) -> bool:
403 + param = params_dict[name]
404 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
405 + loaded_local_expert = False
406 + for expert_id in range(num_experts):
407 + curr_expert_weight = loaded_weight[expert_id]
408 + success = weight_loader(
409 + param,
410 + curr_expert_weight,
411 + name,
412 + shard_id,
413 + expert_id,
414 + return_success=True,
415 + )
416 + if success:
417 + loaded_local_expert = True
418 +
419 + return loaded_local_expert
420 +
421 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
422 + # FP8 indexer wk dequant buffer (weight and scale arrive separately).
423 + pending_wk_fp8: dict[str, dict[str, torch.Tensor]] = {}
424 + pp_missing_layer_names = get_pp_missing_layer_names(self)
425 + skip_topk_layers = compute_skip_topk_layers(self.config)
426 +
427 + # Must not be cached on `self`: `process_weights_after_loading` swaps in
428 + # kernel-specific expert weights via `replace_parameter`, and a cache
429 + # outliving this call would pin the pre-shuffle storage (OOM on large
430 + # MoE) and make a later reload target orphaned tensors.
431 + params_dict = dict(self.named_parameters())
432 + # Split per-expert mapping (V3 style): experts.0.gate_proj.weight
433 + split_expert_params_mapping = self.get_expert_mapping()
434 + loaded_params: set[str] = set()
435 +
436 + # Sink weights are sharded like the q/k/v linears.
437 + sink_tp_size = get_tensor_model_parallel_world_size()
438 + sink_tp_rank = get_tensor_model_parallel_rank()
439 + n_local_head = self.config.num_attention_heads // sink_tp_size
440 + head_rank_start = n_local_head * sink_tp_rank
441 + head_rank_end = n_local_head * (sink_tp_rank + 1)
442 + base_layer = (
443 + "base_layer." if any(".base_layer." in name for name in params_dict) else ""
444 + )
445 + # Fused expert mapping: experts.gate_up_proj (all experts in one tensor).
446 + # The packed weights are owned by the RoutedExperts submodule, so the
447 + # targets are experts.routed_experts.[base_layer.]w{13,2}_weight.
448 + fused_expert_prefix = f".experts.routed_experts.{base_layer}"
449 + fused_expert_params_mapping = [
450 + (f"{fused_expert_prefix}w13_weight", ".experts.gate_up_proj", 0, "w1"),
451 + (f"{fused_expert_prefix}w2_weight", ".experts.down_proj", 0, "w2"),
452 + ]
453 + num_experts = getattr(self.config, "num_experts", 0)
454 +
455 + def _should_skip_missing_param(param_name: str) -> bool:
456 + # Sparse checkpoints may contain indexer weights for layers that
457 + # fall back to dense MLA on unsupported GPUs.
458 + return ".indexer." in param_name
459 +
460 + def _is_split_expert_weight(weight_name: str) -> bool:
461 + """Whether this weight is in split (per-expert) format."""
462 + # Split format: mlp.experts.<id>.gate_proj.weight
463 + return bool(re.search(r"\.experts\.\d+\.", weight_name))
464 +
465 + def _is_fused_expert_weight(weight_name: str) -> bool:
466 + """Whether this weight is in fused (all-experts-packed) format."""
467 + return ".experts.gate_up_proj" in weight_name or (
468 + ".experts.down_proj" in weight_name
469 + and not _is_split_expert_weight(weight_name)
470 + )
471 +
472 + for name, loaded_weight in weights:
473 + if is_skip_topk_indexer_weight(name, skip_topk_layers):
474 + continue
475 + if _try_load_fp8_indexer_wk(
476 + name,
477 + loaded_weight,
478 + pending_wk_fp8,
479 + params_dict,
480 + loaded_params,
481 + pp_missing_layer_names,
482 + ):
483 + continue
484 + if self.config.tie_word_embeddings and "lm_head.weight" in name:
485 + continue
486 + # Names and stacked shard ids are normalized upstream by
487 + # AutoWeightsLoader, through quant_config.get_cache_scale_mapper()
488 + # and hf_to_vllm_mapper.
489 + stacked_shard_id = getattr(loaded_weight, "shard_id", None)
490 + if ".indexer.wk." in name:
491 + name = name.replace(".indexer.wk.", ".indexer.wk_weights_proj.")
492 + stacked_shard_id = 0
493 + # Skip loading extra bias for GPTQ models.
494 + if name.endswith(".bias") and name not in params_dict:
495 + continue
496 +
497 + # Determine per-weight whether this is fused or split format.
498 + is_fused_expert = _is_fused_expert_weight(name)
499 + expert_params_mapping = (
500 + fused_expert_params_mapping
501 + if is_fused_expert
502 + else split_expert_params_mapping
503 + )
504 +
505 + is_expert_weight = False
506 + loaded_expert_param_names: set[str] = set()
507 + for mapping in expert_params_mapping:
508 + param_name, weight_name, expert_id, shard_id = mapping
509 + if weight_name not in name:
510 + continue
511 +
512 + # This is an expert weight and must not be attempted as any
513 + # other kind of weight later on.
514 + is_expert_weight = True
515 +
516 + # Do not modify `name`: the loop may continue past this point.
517 + name_mapped = name.replace(weight_name, param_name)
518 + if is_pp_missing_parameter(name_mapped, self):
519 + continue
520 + if is_fused_expert:
521 + if "experts.gate_up_proj" in name:
522 + chunks = loaded_weight.chunk(2, dim=-2)
523 + success_w1 = self.load_fused_expert_weights(
524 + name_mapped, params_dict, chunks[0], "w1", num_experts
525 + )
526 + success_w3 = self.load_fused_expert_weights(
527 + name_mapped, params_dict, chunks[1], "w3", num_experts
528 + )
529 + success = success_w1 and success_w3
530 + else:
531 + success = self.load_fused_expert_weights(
532 + name_mapped,
533 + params_dict,
534 + loaded_weight,
535 + shard_id,
536 + num_experts,
537 + )
538 + if success:
539 + name = name_mapped
540 + break
541 + else:
542 + # Split per-expert format (V3 style).
543 + if name_mapped not in params_dict:
544 + if _should_skip_missing_param(name_mapped):
545 + continue
546 + logger.warning_once(
547 + "Skipping unknown checkpoint weight: %s",
548 + name_mapped,
549 + )
550 + continue
551 + param = params_dict[name_mapped]
552 + # Ask the weight loader whether it succeeded, otherwise we
553 + # may skip experts that have other available replicas.
554 + weight_loader: Callable[..., typing.Any] = param.weight_loader
555 + success = weight_loader(
556 + param,
557 + loaded_weight,
558 + name_mapped,
559 + shard_id=shard_id,
560 + expert_id=expert_id,
561 + return_success=True,
562 + )
563 + if success:
564 + if not is_fused_expert:
565 + loaded_expert_param_names.add(name_mapped)
566 + continue
567 + name = name_mapped
568 + break
569 + else:
570 + if loaded_expert_param_names:
571 + loaded_params.update(loaded_expert_param_names)
572 + continue
573 + if "learnable_sink_param" in name:
574 + if is_pp_missing_parameter(name, self):
575 + continue
576 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
577 + n = narrow_weight.shape[0]
578 + with torch.no_grad():
579 + params_dict[name][:n].copy_(narrow_weight)
580 + else:
581 + if is_expert_weight:
582 + # An expert weight that is not mapped to this rank.
583 + continue
584 + if "gate.e_score_correction_bias" in name:
585 + name = name.replace(
586 + "gate.e_score_correction_bias", "expert_bias"
587 + )
588 + if "router.gate." in name:
589 + name = name.replace("router.", "")
590 + if "hc_fn" in name:
591 + name = name.replace("hc_fn", "hc_fn.weight")
592 + if "hc_head_fn" in name:
593 + name = name.replace("hc_head_fn", "hc_head_fn.weight")
594 +
595 + if is_pp_missing_parameter(name, self):
596 + continue
597 +
598 + if name not in params_dict:
599 + if _should_skip_missing_param(name):
600 + continue
601 + logger.warning_once(
602 + "Skipping unknown checkpoint weight: %s",
603 + name,
604 + )
605 + continue
606 +
607 + param = params_dict[name]
608 + weight_loader = getattr(
609 + param, "weight_loader", default_weight_loader
610 + )
611 + if stacked_shard_id is None:
612 + weight_loader(param, loaded_weight)
613 + else:
614 + weight_loader(param, loaded_weight, stacked_shard_id)
615 + loaded_params.add(name)
616 +
617 + return loaded_params
618 +
619 +
620 +class HYV4ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
621 + hf_to_vllm_mapper = WeightsMapper(
622 + orig_to_new_stacked={
623 + ".q_a_proj": (".fused_qkv_a_proj", 0),
624 + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
625 + ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
626 + ".mlp.up_proj": (".mlp.gate_up_proj", 1),
627 + ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
628 + ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
629 + ".indexer.weights_proj.": (".indexer.wk_weights_proj.", 1),
630 + }
631 + )
632 + packed_modules_mapping = {
633 + "gate_up_proj": ["gate_proj", "up_proj"],
634 + # MLA runs both latent down-projections as one GEMM.
635 + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
636 + # The indexer fuses wk and weights_proj into one GEMM.
637 + "wk_weights_proj": ["wk", "weights_proj"],
638 + }
639 +
640 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
641 + super().__init__()
642 + config = vllm_config.model_config.hf_config
643 + quant_config = vllm_config.quant_config
644 + self.config = config
645 + self.quant_config = quant_config
646 + if quant_config is not None:
647 + quant_config.packed_modules_mapping = self.packed_modules_mapping
648 +
649 + parallel_config = vllm_config.parallel_config
650 + eplb_config = parallel_config.eplb_config
651 + self.num_redundant_experts = eplb_config.num_redundant_experts
652 +
653 + self.model = HYV4Model(
654 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
655 + )
656 + if get_pp_group().is_last_rank:
657 + # The head stays in the model dtype; ``enable_lm_head_fp32`` is
658 + # surfaced as ``head_dtype`` on the config (see HYV4Config), which
659 + # makes LogitsProcessor accumulate the projection straight into
660 + # fp32 instead of materializing an fp32 copy of the weight.
661 + self.lm_head = ParallelLMHead(
662 + config.vocab_size,
663 + config.hidden_size,
664 + quant_config=quant_config,
665 + prefix=maybe_prefix(prefix, "lm_head"),
666 + )
667 + # With tie_word_embeddings, embed_tokens is kept on the last rank
668 + # (see HYV4Model.__init__) so the weight can be shared here.
669 + if self.config.tie_word_embeddings:
670 + self.lm_head.weight = self.model.embed_tokens.weight
671 + else:
672 + self.lm_head = PPMissingLayer()
673 + self.logits_processor = LogitsProcessor(config.vocab_size)
674 + self.make_empty_intermediate_tensors = (
675 + self.model.make_empty_intermediate_tensors
676 + )
677 +
678 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
679 + return self.model.embed_input_ids(input_ids)
680 +
681 + def forward(
682 + self,
683 + input_ids: torch.Tensor,
684 + positions: torch.Tensor,
685 + intermediate_tensors: IntermediateTensors | None = None,
686 + inputs_embeds: torch.Tensor | None = None,
687 + ) -> torch.Tensor | IntermediateTensors:
688 + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
689 +
690 + def compute_logits(
691 + self,
692 + hidden_states: torch.Tensor,
693 + ) -> torch.Tensor | None:
694 + logits = self.logits_processor(self.lm_head, hidden_states)
695 +
696 + if getattr(self.config, "soft_logits_capping", False):
697 + soft_cap = self.config.soft_logits_capping_logits
698 + logits = soft_cap * torch.nn.functional.tanh(logits / soft_cap)
699 +
700 + return logits
701 +
702 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
703 + # MTP layers are loaded by the MTP head, not by the target model. They
704 + # appear as model.mtp_layers.<i>.* and as the layer ids past the backbone.
705 + mtp_start = self.config.num_hidden_layers
706 + drop_prefixes: dict[str, str | None] = {"model.mtp_layers.": None}
707 + for i in range(getattr(self.config, "num_nextn_predict_layers", 0)):
708 + drop_prefixes[f"model.layers.{mtp_start + i}."] = None
709 + if self.config.tie_word_embeddings:
710 + drop_prefixes["lm_head."] = None
711 + drop_mapper = WeightsMapper(orig_to_new_prefix=drop_prefixes)
712 + loader = AutoWeightsLoader(self)
713 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper | drop_mapper)
714 +
715 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
716 + return self.model.get_expert_mapping()
vllm/models/hy_v4/nvidia/moe.py new
+190
@@ -0,0 +1,190 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Dense FFN and MoE blocks for HY V4 (NVIDIA)."""
4 +
5 +import torch
6 +from torch import nn
7 +from transformers import PretrainedConfig
8 +
9 +from vllm.config import get_current_vllm_config
10 +from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.activation import SiluAndMul
13 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear
14 +from vllm.model_executor.layers.linear import (
15 + MergedColumnParallelLinear,
16 + RowParallelLinear,
17 +)
18 +from vllm.model_executor.layers.quantization import QuantizationConfig
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +class HYV4FeedForward(nn.Module):
24 + """Dense SwiGLU feed-forward block.
25 +
26 + Used both for the dense decoder layers and for the MoE shared experts.
27 + The routed-expert SwiGLU clamp does not apply here, matching the reference
28 + implementation.
29 + """
30 +
31 + def __init__(
32 + self,
33 + hidden_size: int,
34 + intermediate_size: int,
35 + hidden_act: str,
36 + quant_config: QuantizationConfig | None = None,
37 + reduce_results: bool = True,
38 + prefix: str = "",
39 + ) -> None:
40 + super().__init__()
41 + if hidden_act != "silu":
42 + raise ValueError(
43 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."
44 + )
45 + self.gate_up_proj = MergedColumnParallelLinear(
46 + hidden_size,
47 + [intermediate_size] * 2,
48 + bias=False,
49 + quant_config=quant_config,
50 + prefix=f"{prefix}.gate_up_proj",
51 + )
52 + self.down_proj = RowParallelLinear(
53 + intermediate_size,
54 + hidden_size,
55 + bias=False,
56 + quant_config=quant_config,
57 + reduce_results=reduce_results,
58 + prefix=f"{prefix}.down_proj",
59 + )
60 + self.act_fn = SiluAndMul()
61 +
62 + def forward(self, x: torch.Tensor) -> torch.Tensor:
63 + gate_up, _ = self.gate_up_proj(x)
64 + out = self.act_fn(gate_up)
65 + out, _ = self.down_proj(out)
66 + return out
67 +
68 +
69 +class HYV4MoEFused(nn.Module):
70 + """HY V4 MoE layer with optional SwiGLU clamp support.
71 +
72 + When ``config.swiglu_limit > 0`` the routed experts use a clamped SwiGLU::
73 +
74 + gate = clamp(gate, max=limit)
75 + up = clamp(up, -limit, limit)
76 + output = silu(gate) * up
77 +
78 + Dense layers and shared experts are NOT clamped. The clamp is forwarded to
79 + ``FusedMoEFactory`` as ``swiglu_limit`` so it travels through the quant
80 + config to every expert backend instead of mutating global state.
81 +
82 + NOTE: The reference implementation additionally supports HPC gate kernels,
83 + TPCP/MLPSP sharded shared experts and a side-stream overlap for the shared
84 + experts. Those depend on infrastructure that is absent here, so this port
85 + keeps only the standard path. TODO: restore once available.
86 + """
87 +
88 + def __init__(
89 + self,
90 + config: PretrainedConfig,
91 + quant_config: QuantizationConfig | None = None,
92 + prefix: str = "",
93 + enable_eplb: bool = False,
94 + ):
95 + super().__init__()
96 + self.tp_size = get_tensor_model_parallel_world_size()
97 + self.ep_group = get_ep_group().device_group
98 + self.ep_rank = get_ep_group().rank_in_group
99 + self.ep_size = self.ep_group.size()
100 + self.n_routed_experts = config.num_experts
101 + if self.tp_size > config.num_experts:
102 + raise ValueError(
103 + f"Tensor parallel size {self.tp_size} is greater than "
104 + f"the number of experts {config.num_experts}."
105 + )
106 + top_k = config.num_experts_per_tok
107 + intermediate_size = config.expert_hidden_dim
108 + router_scaling_factor = getattr(config, "router_scaling_factor", 1.0)
109 + vllm_config = get_current_vllm_config()
110 + eplb_config = vllm_config.parallel_config.eplb_config
111 + self.enable_eplb = enable_eplb
112 +
113 + self.n_logical_experts = self.n_routed_experts
114 + self.n_redundant_experts = eplb_config.num_redundant_experts
115 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
116 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
117 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
118 + self.physical_expert_end = (
119 + self.physical_expert_start + self.n_local_physical_experts
120 + )
121 +
122 + self.gate = GateLinear(
123 + config.hidden_size,
124 + config.num_experts,
125 + bias=False,
126 + out_dtype=torch.float32,
127 + params_dtype=torch.float32,
128 + prefix=f"{prefix}.gate",
129 + )
130 +
131 + self.shared_experts: HYV4FeedForward | None
132 + if config.num_shared_experts > 0:
133 + self.shared_experts = HYV4FeedForward(
134 + hidden_size=config.hidden_size,
135 + intermediate_size=config.expert_hidden_dim * config.num_shared_experts,
136 + hidden_act=config.hidden_act,
137 + quant_config=quant_config,
138 + prefix=f"{prefix}.shared_experts",
139 + reduce_results=False,
140 + )
141 + else:
142 + self.shared_experts = None
143 +
144 + self.expert_bias = nn.Parameter(
145 + torch.empty(config.num_experts, dtype=torch.float32)
146 + )
147 +
148 + raw_swiglu_limit = getattr(config, "swiglu_limit", 0)
149 + use_swiglu_clamp = bool(raw_swiglu_limit) and float(raw_swiglu_limit) > 0
150 + moe_swiglu_limit = float(raw_swiglu_limit) if use_swiglu_clamp else None
151 +
152 + self.experts = FusedMoEFactory(
153 + num_experts=self.n_routed_experts,
154 + top_k=top_k,
155 + hidden_size=config.hidden_size,
156 + intermediate_size=intermediate_size,
157 + renormalize=config.route_norm,
158 + quant_config=quant_config,
159 + prefix=f"{prefix}.experts",
160 + enable_eplb=self.enable_eplb,
161 + num_redundant_experts=self.n_redundant_experts,
162 + scoring_func="sigmoid",
163 + use_grouped_topk=True,
164 + num_expert_group=1,
165 + topk_group=1,
166 + routed_scaling_factor=router_scaling_factor,
167 + e_score_correction_bias=self.expert_bias,
168 + shared_experts=self.shared_experts,
169 + swiglu_limit=moe_swiglu_limit,
170 + )
171 + self.prefix = prefix
172 + if use_swiglu_clamp:
173 + logger.debug_once(
174 + "HYV4MoEFused: swiglu_limit=%.1f enabled for routed experts",
175 + moe_swiglu_limit,
176 + )
177 +
178 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
179 + orig_shape = hidden_states.shape
180 + hidden_dim = hidden_states.shape[-1]
181 + hidden_states = hidden_states.view(-1, hidden_dim)
182 +
183 + # router_logits: (num_tokens, n_experts); the gate stores fp32 weights
184 + # and emits fp32 logits, so no explicit upcast is needed here.
185 + router_logits, _ = self.gate(hidden_states)
186 +
187 + final_hidden_states = self.experts(
188 + hidden_states=hidden_states, router_logits=router_logits
189 + )
190 + return final_hidden_states.view(orig_shape)
vllm/models/hy_v4/nvidia/mtp.py new
+937
@@ -0,0 +1,937 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Multi-token prediction (MTP) head for HY V4 (NVIDIA)."""
4 +
5 +import copy
6 +import typing
7 +from collections.abc import Callable, Iterable
8 +
9 +import torch
10 +from torch import nn
11 +from transformers import PretrainedConfig
12 +
13 +from vllm.config import CacheConfig, ModelConfig, VllmConfig
14 +from vllm.distributed import (
15 + get_tensor_model_parallel_rank,
16 + get_tensor_model_parallel_world_size,
17 +)
18 +from vllm.logger import init_logger
19 +from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
20 +from vllm.model_executor.layers.layernorm import RMSNorm
21 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
22 +from vllm.model_executor.layers.quantization import QuantizationConfig
23 +from vllm.model_executor.layers.vocab_parallel_embedding import (
24 + ParallelLMHead,
25 + VocabParallelEmbedding,
26 +)
27 +from vllm.model_executor.model_loader.weight_utils import (
28 + default_weight_loader,
29 + maybe_remap_kv_scale_name,
30 +)
31 +from vllm.model_executor.models.deepseek_v2 import _try_load_fp8_indexer_wk
32 +from vllm.model_executor.models.utils import (
33 + get_pp_missing_layer_names,
34 + is_pp_missing_parameter,
35 + maybe_prefix,
36 +)
37 +from vllm.sequence import IntermediateTensors
38 +from vllm.v1.outputs import SamplerOutput
39 +from vllm.v1.sample.metadata import SamplingMetadata
40 +from vllm.v1.sample.sampler import Sampler
41 +
42 +from .model import HYV4DecoderLayer, _normalize_hyv4_config
43 +
44 +logger = init_logger(__name__)
45 +
46 +_MTP_SCALE_ONLY_PARAM_SUBSTRINGS = (
47 + "_scale",
48 + "_weight_scale",
49 + "_weight_scale_inv",
50 + "_input_scale",
51 +)
52 +
53 +# Names of the quant-config attribute that stores the unquantized-module
54 +# prefix list, in the order we look them up. Different QuantizationConfig
55 +# subclasses expose the list under different names:
56 +# * ``ignored_layers`` -- Fp8Config, GPTQ / AWQ configs.
57 +# * ``exclude_modules`` -- ModelOpt* configs (Fp8, NvFp4, MxFp8, Mixed).
58 +_MTP_QUANT_EXCLUSION_ATTRS = ("ignored_layers", "exclude_modules")
59 +
60 +
61 +def _get_spec_layer_idx_from_weight_name(
62 + config: PretrainedConfig, weight_name: str
63 +) -> int | None:
64 + """Return the MTP layer index a checkpoint weight belongs to, or None.
65 +
66 + Compatible with``num_nextn_predict_layers`` of 1 and 2 for a single MTP
67 + head, matching how the HY V4 checkpoints are exported.
68 + """
69 + if getattr(config, "num_nextn_predict_layers", 0) < 1:
70 + return None
71 + layer_idx = config.num_hidden_layers
72 + num_mtp_layers = max(config.num_nextn_predict_layers - 1, 1)
73 + for i in range(num_mtp_layers):
74 + if weight_name.startswith(f"model.layers.{layer_idx + i}."):
75 + return layer_idx + i
76 + return None
77 +
78 +
79 +def _should_skip_missing_mtp_scale_param(
80 + quant_config: QuantizationConfig | None,
81 + name: str,
82 +) -> bool:
83 + """Whether an unmatched scale parameter can be silently ignored."""
84 + return quant_config is None and any(
85 + scale_name in name for scale_name in _MTP_SCALE_ONLY_PARAM_SUBSTRINGS
86 + )
87 +
88 +
89 +def _resolve_fused_expert_param(
90 + param_base: str,
91 + ckpt_suffix: str,
92 + params_dict: dict,
93 +) -> str | None:
94 + """Map a fused expert checkpoint suffix onto a draft parameter name.
95 +
96 + Fused checkpoints pack every expert into one tensor and name the companion
97 + scales by suffixing the projection (``experts.gate_up_proj_scale_inv``),
98 + while the draft model owns ``experts.routed_experts.w13_weight`` plus a
99 + separate ``..._scale_inv`` parameter. Dropping the suffix would push the
100 + scale into the weight parameter and leave the scale at its sentinel init.
101 +
102 + Args:
103 + param_base: Draft parameter holding the packed weight.
104 + ckpt_suffix: Checkpoint text following the projection name.
105 + params_dict: The draft model's named parameters.
106 +
107 + Returns:
108 + The matching draft parameter name, or None when there is none.
109 + """
110 + if not ckpt_suffix:
111 + return param_base if param_base in params_dict else None
112 + tail = ckpt_suffix.lstrip(".").removeprefix("weight_")
113 + for candidate in (
114 + f"{param_base}_{tail}",
115 + f"{param_base}_scale_inv",
116 + f"{param_base}_scale",
117 + ):
118 + if candidate in params_dict:
119 + return candidate
120 + return None
121 +
122 +
123 +def _prepare_mtp_fp8_expert_scale(
124 + quant_config: QuantizationConfig | None,
125 + name: str,
126 + loaded_weight: torch.Tensor,
127 +) -> tuple[str, torch.Tensor]:
128 + """Normalize block-wise FP8 expert scale names and dtypes."""
129 + from vllm.model_executor.layers.quantization.fp8 import Fp8Config
130 +
131 + if (
132 + isinstance(quant_config, Fp8Config)
133 + and quant_config.weight_block_size is not None
134 + and ".mlp.experts." in name
135 + and name.endswith(".scale")
136 + ):
137 + name = name[: -len(".scale")] + ".weight_scale_inv"
138 + if (
139 + getattr(quant_config, "is_scale_e8m0", False)
140 + and loaded_weight.dtype == torch.uint8
141 + ):
142 + # UE8M0 scale exponents are stored as raw uint8 bytes: reinterpret
143 + # the bits instead of numerically converting them.
144 + loaded_weight = loaded_weight.view(torch.float8_e8m0fnu)
145 + return name, loaded_weight
146 +
147 +
148 +def _create_mtp_quant_config(
149 + hf_config: PretrainedConfig,
150 + backbone_quant_config: QuantizationConfig | None = None,
151 +) -> QuantizationConfig | None:
152 + """Create the quantization config for the MTP layers.
153 +
154 + The MTP quantization algorithm is given by the ``mtp_quant_algo`` field in
155 + ``config.json``, independently of the backbone's quantization. Supported
156 + values are ``"FP8"``, ``"NONE"`` (inherit the backbone) and ``"BF16"`` /
157 + ``"FP16"`` (unquantized). A missing or ``"NONE"`` value falls back to the
158 + backbone config.
159 +
160 + Args:
161 + hf_config: The draft model's HF config.
162 + backbone_quant_config: The target model's quantization config.
163 +
164 + Returns:
165 + The quantization config to use for the MTP layers, or None when the MTP
166 + layers are unquantized.
167 + """
168 + mtp_quant_algo = getattr(hf_config, "mtp_quant_algo", None)
169 +
170 + if mtp_quant_algo is None or mtp_quant_algo.upper() == "NONE":
171 + logger.info_once(
172 + "No explicit MTP quant algo (mtp_quant_algo=%s). "
173 + "Using the backbone's quant config.",
174 + mtp_quant_algo,
175 + )
176 + return backbone_quant_config
177 +
178 + mtp_quant_algo = mtp_quant_algo.upper()
179 +
180 + if mtp_quant_algo in ("BF16", "FP16"):
181 + logger.info_once(
182 + "MTP layers are unquantized (mtp_quant_algo=%s).", mtp_quant_algo
183 + )
184 + return None
185 +
186 + if mtp_quant_algo == "FP8":
187 + from vllm.model_executor.layers.quantization.fp8 import Fp8Config
188 +
189 + logger.info_once("MTP uses FP8 quantization (mtp_quant_algo=FP8).")
190 + hf_quant_config = getattr(hf_config, "quantization_config", None) or {}
191 + weight_block_size = hf_quant_config.get("weight_block_size")
192 + activation_scheme = hf_quant_config.get("activation_scheme", "dynamic")
193 + if weight_block_size is not None:
194 + activation_scheme = "dynamic"
195 + ignored_layers = hf_quant_config.get("ignored_layers") or hf_quant_config.get(
196 + "modules_to_not_convert"
197 + )
198 +
199 + fp8_config = Fp8Config(
200 + is_checkpoint_fp8_serialized=True,
201 + activation_scheme=activation_scheme,
202 + ignored_layers=ignored_layers or [],
203 + weight_block_size=weight_block_size,
204 + )
205 + if hf_quant_config.get("scale_fmt") == "ue8m0":
206 + # Fp8LinearMethod reads this through getattr, so it is not a
207 + # declared Fp8Config field.
208 + setattr(fp8_config, "is_scale_e8m0", True) # noqa: B010
209 + return fp8_config
210 +
211 + logger.warning(
212 + "Unknown mtp_quant_algo=%s, falling back to the backbone config.",
213 + mtp_quant_algo,
214 + )
215 + return backbone_quant_config
216 +
217 +
218 +def _remap_mtp_quant_exclusions(
219 + quant_config: QuantizationConfig | None,
220 + mtp_start_layer_idx: int,
221 + num_mtp_layers: int,
222 +) -> QuantizationConfig | None:
223 + """Translate checkpoint-named MTP quant exclusions to draft prefixes.
224 +
225 + ``modules_to_not_convert`` / ``exclude_modules`` name MTP modules the way
226 + the checkpoint does (``model.mtp_layers.0.self_attn.linear_gate``), while
227 + the draft model builds them under ``model.layers.<num_hidden_layers + i>``.
228 + ``is_layer_skipped`` compares prefixes for exact equality, so without this
229 + translation an excluded MTP module gets a quant method even though its
230 + checkpoint weight is BF16 with no ``weight_scale`` / ``weight_scale_inv``
231 + companion: the scale then keeps its ``finfo(float32).min`` sentinel and the
232 + layer silently produces garbage.
233 +
234 + The list lives under a different attribute per config class, hence the
235 + lookup over ``_MTP_QUANT_EXCLUSION_ATTRS``.
236 +
237 + Args:
238 + quant_config: The MTP quantization config to widen.
239 + mtp_start_layer_idx: Index of the first MTP layer in the draft model.
240 + num_mtp_layers: Number of MTP layers.
241 +
242 + Returns:
243 + A shallow copy with the translated exclusions, or the input unchanged.
244 + """
245 + if quant_config is None:
246 + return None
247 +
248 + # Different QuantizationConfig subclasses expose the exclusion list under
249 + # different attribute names; only the one that is actually populated on
250 + # the backbone config governs which draft modules stay unquantized. For
251 + # ModelOpt checkpoints (``quant_method: modelopt`` and friends) this is
252 + # ``exclude_modules``; for the legacy fp8 path it is ``ignored_layers``.
253 + exclusion_attr: str | None = None
254 + patterns: list[str] | None = None
255 + for attr in _MTP_QUANT_EXCLUSION_ATTRS:
256 + value = getattr(quant_config, attr, None)
257 + if value:
258 + exclusion_attr = attr
259 + patterns = list(value)
260 + break
261 + if not patterns or exclusion_attr is None:
262 + return quant_config
263 +
264 + extra: list[str] = []
265 + for offset in range(num_mtp_layers):
266 + source_prefix = f"model.mtp_layers.{offset}."
267 + draft_prefix = f"model.layers.{mtp_start_layer_idx + offset}."
268 + for pattern in patterns:
269 + if pattern.startswith(source_prefix):
270 + extra.append(draft_prefix + pattern[len(source_prefix) :])
271 + # ModelOpt exclusion patterns may end with a ``*`` wildcard; the
272 + # remap must preserve that suffix so ``is_layer_skipped`` still
273 + # matches the draft module prefixes.
274 + elif pattern.rstrip("*").startswith(source_prefix):
275 + stripped = pattern.rstrip("*")
276 + trailing_stars = pattern[len(stripped) :]
277 + extra.append(
278 + draft_prefix + stripped[len(source_prefix) :] + trailing_stars
279 + )
280 + if not extra:
281 + return quant_config
282 +
283 + # The backbone shares this object with the target model; copy before
284 + # widening the exclusion list. The attribute is only present on the
285 + # concrete configs that support exclusions, hence the getattr/setattr pair.
286 + quant_config = copy.copy(quant_config)
287 + setattr(quant_config, exclusion_attr, patterns + extra) # noqa: B010
288 + logger.info_once(
289 + "HYV4 MTP: mapped %d checkpoint-named quant exclusions (%s) onto the "
290 + "draft module prefixes: %s",
291 + len(extra),
292 + exclusion_attr,
293 + ", ".join(extra),
294 + )
295 + return quant_config
296 +
297 +
298 +class HYV4SharedHead(nn.Module):
299 + """Holds the draft LM head shared with the target model."""
300 +
301 + def __init__(
302 + self,
303 + config: PretrainedConfig,
304 + quant_config: QuantizationConfig | None = None,
305 + ) -> None:
306 + super().__init__()
307 + # Use the same ``lm_head`` prefix as the backbone so that quant-config
308 + # exclusion lists (e.g. ModelOpt's ``exclude_modules: ["lm_head", ...]``)
309 + # also apply here: an unquantized checkpoint lm_head must not gain a
310 + # spurious ``weight_scale`` parameter on the draft side, which would
311 + # otherwise stay at its FP-min sentinel until the proposer swaps the
312 + # whole ParallelLMHead for the target model's ``lm_head``.
313 + self.head = ParallelLMHead(
314 + config.vocab_size,
315 + config.hidden_size,
316 + quant_config=quant_config,
317 + prefix="lm_head",
318 + )
319 +
320 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
321 + return hidden_states
322 +
323 +
324 +class HYV4MultiTokenPredictorLayer(nn.Module):
325 + """A single MTP draft block."""
326 +
327 + def __init__(
328 + self,
329 + config: PretrainedConfig,
330 + prefix: str,
331 + vllm_config: VllmConfig,
332 + model_config: ModelConfig,
333 + cache_config: CacheConfig | None = None,
334 + quant_config: QuantizationConfig | None = None,
335 + topk_indices_buffer: torch.Tensor | None = None,
336 + ) -> None:
337 + super().__init__()
338 + del model_config # kept for signature parity with the other MTP heads
339 + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
340 + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
341 + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
342 + self.shared_head = HYV4SharedHead(config=config, quant_config=quant_config)
343 +
344 + # HYV4DecoderLayer indexes layer_types/mlp_layer_types by the numeric
345 + # layer id parsed from the prefix. MTP layers live after the backbone
346 + # layers, so extend the per-layer config lists to cover those ids.
347 + layer_idx = int(prefix.rsplit(".", 1)[-1])
348 + mtp_config = copy.deepcopy(config)
349 + # MTP checkpoints do not carry iHC pre/post weights for mtp_layers, so
350 + # the draft blocks are built without iHC to match their structure.
351 + mtp_config.enable_ihc = False
352 + mtp_config.layer_types = _extend_layer_types(
353 + getattr(mtp_config, "layer_types", None), layer_idx, "full_attention"
354 + )
355 + mtp_config.mlp_layer_types = _extend_layer_types(
356 + getattr(mtp_config, "mlp_layer_types", None), layer_idx, "sparse"
357 + )
358 +
359 + self.mtp_block = HYV4DecoderLayer(
360 + config=mtp_config,
361 + vllm_config=vllm_config,
362 + cache_config=cache_config,
363 + quant_config=quant_config,
364 + prefix=prefix,
365 + topk_indices_buffer=topk_indices_buffer,
366 + )
367 + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
368 +
369 + def forward(
370 + self,
371 + input_ids: torch.Tensor,
372 + positions: torch.Tensor,
373 + previous_hidden_states: torch.Tensor,
374 + inputs_embeds: torch.Tensor | None = None,
375 + ) -> torch.Tensor:
376 + del input_ids # embeddings are resolved by the caller
377 + assert inputs_embeds is not None
378 + inputs_embeds = self.enorm(inputs_embeds)
379 + previous_hidden_states = self.hnorm(previous_hidden_states)
380 + hidden_states = self.eh_proj(
381 + torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
382 + )
383 + hidden_states, residual = self.mtp_block(
384 + positions=positions,
385 + hidden_states=hidden_states,
386 + residual=None,
387 + )
388 + hidden_states, _ = self.final_layernorm(hidden_states, residual)
389 + return hidden_states
390 +
391 +
392 +def _extend_layer_types(
393 + layer_types: list[str] | None, layer_idx: int, fallback: str
394 +) -> list[str] | None:
395 + """Pad a per-layer config list so``layer_idx`` is addressable."""
396 + if layer_types is None or len(layer_types) > layer_idx:
397 + return layer_types
398 + fill = layer_types[-1] if layer_types else fallback
399 + return list(layer_types) + [fill] * (layer_idx + 1 - len(layer_types))
400 +
401 +
402 +class HYV4MultiTokenPredictor(nn.Module):
403 + """Owns the MTP draft blocks and their shared embedding / logits path."""
404 +
405 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
406 + super().__init__()
407 + target_config = vllm_config.model_config.hf_config
408 + speculative_config = vllm_config.speculative_config
409 + assert speculative_config is not None, (
410 + "HYV4MTP can only be built as a speculative draft model"
411 + )
412 + draft_model_config = speculative_config.draft_model_config
413 + assert draft_model_config is not None
414 + config = _normalize_hyv4_config(draft_model_config.hf_config)
415 + self.mtp_start_layer_idx = target_config.num_hidden_layers
416 + self.num_mtp_layers = max(
417 + getattr(target_config, "num_nextn_predict_layers", 1), 1
418 + )
419 + self.quant_config = _remap_mtp_quant_exclusions(
420 + _create_mtp_quant_config(config, vllm_config.quant_config),
421 + self.mtp_start_layer_idx,
422 + self.num_mtp_layers,
423 + )
424 + if self.quant_config is not None:
425 + self.quant_config.packed_modules_mapping = HYV4MTP.packed_modules_mapping
426 +
427 + # Sparse DSA draft blocks must share the target model's top-k buffer.
428 + # The proposer injects it through HYV4MTP.set_topk_indices_buffer()
429 + # before any draft forward.
430 + self.topk_indices_buffer: torch.Tensor | None = None
431 + if hasattr(config, "index_topk"):
432 + from vllm.platforms import current_platform
433 +
434 + self.topk_indices_buffer = torch.empty(
435 + vllm_config.scheduler_config.max_num_batched_tokens,
436 + config.index_topk,
437 + dtype=torch.int32,
438 + device=current_platform.device_type,
439 + )
440 +
441 + self.layers = nn.ModuleDict(
442 + {
443 + str(idx): HYV4MultiTokenPredictorLayer(
444 + config,
445 + f"{prefix}.layers.{idx}",
446 + vllm_config=vllm_config,
447 + model_config=draft_model_config,
448 + cache_config=vllm_config.cache_config,
449 + quant_config=self.quant_config,
450 + topk_indices_buffer=self.topk_indices_buffer,
451 + )
452 + for idx in range(
453 + self.mtp_start_layer_idx,
454 + self.mtp_start_layer_idx + self.num_mtp_layers,
455 + )
456 + }
457 + )
458 + self.requires_topk_indices_buffer = any(
459 + layer.mtp_block.self_attn.is_sparse for layer in self.layers.values()
460 + )
461 +
462 + self.embed_tokens = VocabParallelEmbedding(
463 + config.vocab_size,
464 + config.hidden_size,
465 + )
466 + self.logits_processor = LogitsProcessor(config.vocab_size)
467 + self.spec_step_idx: int = 0
468 +
469 + def set_skip_topk(self, skip: bool) -> None:
470 + """Toggle the draft indexer for ``index_share_for_mtp_iteration``.
471 +
472 + The proposer clears the flag for draft step 0 so the MTP layer builds
473 + its own top-k indices, then sets it for steps 1+ so they reuse what
474 + step 0 wrote into the shared buffer instead of re-running the indexer.
475 + """
476 + for layer in self.layers.values():
477 + layer.mtp_block.self_attn.skip_topk = skip
478 +
479 + def compact_topk_indices(self, slot_ids: torch.Tensor) -> None:
480 + """Move the top-k rows at ``slot_ids`` to the front of the buffer.
481 +
482 + Step 0 writes one row per query token of the multi-token batch, while
483 + steps 1+ decode a single token per request and index the buffer from 0.
484 + Without this gather they would read another token's rows.
485 + """
486 + num_slots = slot_ids.numel()
487 + if self.topk_indices_buffer is None or num_slots == 0:
488 + return
489 + self.topk_indices_buffer[:num_slots] = self.topk_indices_buffer[slot_ids]
490 +
491 + def forward(
492 + self,
493 + input_ids: torch.Tensor,
494 + positions: torch.Tensor,
495 + previous_hidden_states: torch.Tensor,
496 + inputs_embeds: torch.Tensor | None = None,
497 + ) -> torch.Tensor:
498 + if inputs_embeds is None:
499 + inputs_embeds = self.embed_tokens(input_ids)
500 + inputs_embeds = torch.where((positions == 0).unsqueeze(-1), 0, inputs_embeds)
501 +
502 + current_step_idx = self.spec_step_idx % self.num_mtp_layers
503 + return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
504 + input_ids,
505 + positions,
506 + previous_hidden_states,
507 + inputs_embeds,
508 + )
509 +
510 + def compute_logits(
511 + self,
512 + hidden_states: torch.Tensor,
513 + ) -> torch.Tensor:
514 + current_step_idx = self.spec_step_idx % self.num_mtp_layers
515 + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
516 + lm_head = mtp_layer.shared_head.head
517 + proj_input = mtp_layer.shared_head(hidden_states)
518 + return self.logits_processor(lm_head, proj_input)
519 +
520 +
521 +class HYV4MTP(nn.Module):
522 + """HY V4 MTP draft head.
523 +
524 + Not a pipeline-parallel stage: the draft head always runs on a single rank,
525 + matching `DeepseekV32MTP` / `KimiK3MTP` / `HYV3MTP`.
526 + """
527 +
528 + packed_modules_mapping = {
529 + "gate_up_proj": ["gate_proj", "up_proj"],
530 + # MLA runs both latent down-projections as one GEMM.
531 + "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
532 + # The indexer fuses wk and weights_proj into one GEMM.
533 + "wk_weights_proj": ["wk", "weights_proj"],
534 + }
535 +
536 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
537 + super().__init__()
538 + self.config = vllm_config.model_config.hf_config
539 + self.model = HYV4MultiTokenPredictor(
540 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
541 + )
542 + self.quant_config = self.model.quant_config
543 + self.sampler = Sampler()
544 +
545 + def set_topk_indices_buffer(self, topk_indices_buffer: torch.Tensor) -> None:
546 + """Share the target sparse-index buffer with every draft consumer.
547 +
548 + Proposers that walk ``named_modules()`` instead of calling this reach
549 + the same consumers via ``HYV4MLAAttention.topk_indices_buffer``.
550 + """
551 + self.model.topk_indices_buffer = topk_indices_buffer
552 + for layer in self.model.layers.values():
553 + self_attn = layer.mtp_block.self_attn
554 + if not self_attn.is_sparse:
555 + continue
556 +
557 + indexer = self_attn.indexer
558 + assert indexer is not None, "Sparse HYV4 MTP attention requires an indexer"
559 + indexer.topk_indices_buffer = topk_indices_buffer
560 + indexer.indexer_op.topk_indices_buffer = topk_indices_buffer
561 +
562 + attn_impl = self_attn.mla_attn.impl
563 + assert hasattr(attn_impl, "topk_indices_buffer"), (
564 + "Sparse HYV4 MTP attention backend requires a top-k indices buffer"
565 + )
566 + attn_impl.topk_indices_buffer = topk_indices_buffer
567 +
568 + def forward(
569 + self,
570 + input_ids: torch.Tensor,
571 + positions: torch.Tensor,
572 + hidden_states: torch.Tensor,
573 + intermediate_tensors: IntermediateTensors | None = None,
574 + inputs_embeds: torch.Tensor | None = None,
575 + spec_step_idx: int = 0,
576 + ) -> torch.Tensor:
577 + del intermediate_tensors # the MTP head is single-stage
578 + if (
579 + self.model.requires_topk_indices_buffer
580 + and self.model.topk_indices_buffer is None
581 + ):
582 + raise RuntimeError(
583 + "HYV4 sparse MTP requires the target model's top-k indices buffer. "
584 + "The proposer must call HYV4MTP.set_topk_indices_buffer() before "
585 + "the first draft forward."
586 + )
587 + self.model.spec_step_idx = spec_step_idx
588 + return self.model(input_ids, positions, hidden_states, inputs_embeds)
589 +
590 + def compute_logits(
591 + self,
592 + hidden_states: torch.Tensor,
593 + spec_step_idx: int = 0,
594 + ) -> torch.Tensor | None:
595 + self.model.spec_step_idx = spec_step_idx
596 + return self.model.compute_logits(hidden_states)
597 +
598 + def sample(
599 + self,
600 + logits: torch.Tensor,
601 + sampling_metadata: SamplingMetadata,
602 + ) -> SamplerOutput | None:
603 + return self.sampler(logits, sampling_metadata)
604 +
605 + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
606 + if f"model.layers.{spec_layer}.embed_tokens" in name:
607 + return "__skip__"
608 + if f"model.layers.{spec_layer}.shared_head" in name:
609 + return "__skip__"
610 +
611 + spec_layer_weight_names = ["enorm", "hnorm", "eh_proj", "final_layernorm"]
612 + spec_layer_weight = any(
613 + weight_name in name for weight_name in spec_layer_weight_names
614 + )
615 + if not spec_layer_weight:
616 + name = name.replace(
617 + f"model.layers.{spec_layer}.",
618 + f"model.layers.{spec_layer}.mtp_block.",
619 + )
620 + return name
621 +
622 + def _load_fused_expert_weights(
623 + self,
624 + name: str,
625 + params_dict: dict,
626 + loaded_weight: torch.Tensor,
627 + shard_id: str,
628 + num_experts: int,
629 + ) -> bool:
630 + if name not in params_dict:
631 + return False
632 + param = params_dict[name]
633 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
634 + loaded_local_expert = False
635 + for expert_id in range(num_experts):
636 + curr_expert_weight = loaded_weight[expert_id]
637 + success = weight_loader(
638 + param,
639 + curr_expert_weight,
640 + name,
641 + shard_id,
642 + expert_id,
643 + return_success=True,
644 + )
645 + if success:
646 + loaded_local_expert = True
647 + return loaded_local_expert
648 +
649 + def _load_expert_weight(
650 + self,
651 + name: str,
652 + loaded_weight: torch.Tensor,
653 + params_dict: dict,
654 + loaded_params: set[str],
655 + split_expert_params_mapping: list[tuple[str, str, int, str]],
656 + fused_expert_param_names: dict[tuple[str, str], str],
657 + num_experts: int,
658 + ) -> bool:
659 + """Load one routed-expert weight in either checkpoint layout.
660 +
661 + Args:
662 + name: Weight name already rewritten to draft-module naming.
663 + loaded_weight: The checkpoint tensor.
664 + params_dict: The draft model's named parameters.
665 + loaded_params: Set updated with the parameters that received a value.
666 + split_expert_params_mapping: Mapping for the per-expert layout.
667 + fused_expert_param_names: ``(mlp_prefix, tag) -> param name`` for the
668 + all-experts-packed layout.
669 + num_experts: Total number of routed experts.
670 +
671 + Returns:
672 + True when the weight was consumed (even if this rank holds none of
673 + the addressed experts).
674 + """
675 + base = name.split(".experts.")[0]
676 + for ckpt_proj, tag in (
677 + (".experts.gate_up_proj", "w13_weight"),
678 + (".experts.down_proj", "w2_weight"),
679 + ):
680 + if ckpt_proj not in name:
681 + continue
682 + param_base = fused_expert_param_names.get((base, tag))
683 + if param_base is None:
684 + return False
685 + # Keep the checkpoint suffix (e.g. `_scale_inv`) so block-scale
686 + # tensors land in the scale parameter, not in the weight.
687 + target = _resolve_fused_expert_param(
688 + param_base, name.split(ckpt_proj, 1)[1], params_dict
689 + )
690 + if target is None:
691 + return False
692 + if tag == "w13_weight":
693 + chunks = loaded_weight.chunk(2, dim=-2)
694 + loaded_w1 = self._load_fused_expert_weights(
695 + target, params_dict, chunks[0], "w1", num_experts
696 + )
697 + loaded_w3 = self._load_fused_expert_weights(
698 + target, params_dict, chunks[1], "w3", num_experts
699 + )
700 + loaded = loaded_w1 and loaded_w3
701 + else:
702 + loaded = self._load_fused_expert_weights(
703 + target, params_dict, loaded_weight, "w2", num_experts
704 + )
705 + if loaded:
706 + loaded_params.add(target)
707 + # The weight belongs to the experts either way; never fall through.
708 + return True
709 +
710 + consumed = False
711 + for param_name, weight_name, expert_id, shard_id in split_expert_params_mapping:
712 + if weight_name not in name:
713 + continue
714 + consumed = True
715 + name_mapped = name.replace(weight_name, param_name)
716 + if name_mapped not in params_dict:
717 + continue
718 + param = params_dict[name_mapped]
719 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
720 + if weight_loader(
721 + param,
722 + loaded_weight,
723 + name_mapped,
724 + shard_id=shard_id,
725 + expert_id=expert_id,
726 + return_success=True,
727 + ):
728 + loaded_params.add(name_mapped)
729 + return consumed
730 +
731 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
732 + params_dict = dict(self.named_parameters())
733 + pp_missing_layer_names = get_pp_missing_layer_names(self)
734 + loaded_params: set[str] = set()
735 +
736 + mtp_start = self.config.num_hidden_layers
737 + shared_weights = {
738 + "model.embed_tokens.weight": "model.embed_tokens.weight",
739 + "lm_head.weight": f"model.layers.{mtp_start}.shared_head.head.weight",
740 + }
741 +
742 + num_experts = getattr(self.config, "n_routed_experts", 0)
743 + sink_tp_size = get_tensor_model_parallel_world_size()
744 + sink_tp_rank = get_tensor_model_parallel_rank()
745 + n_local_head = self.config.num_attention_heads // sink_tp_size
746 + head_rank_start = n_local_head * sink_tp_rank
747 + head_rank_end = n_local_head * (sink_tp_rank + 1)
748 +
749 + # Routed-expert weights come in two checkpoint layouts:
750 + # split: mlp.experts.<id>.gate_proj.weight
751 + # fused: mlp.experts.gate_up_proj (all experts in one tensor)
752 + # The split layout is resolved through the shared EPLB helper so the
753 + # target param names stay in sync with the RoutedExperts module layout;
754 + # the fused layout is resolved from the live parameter names.
755 + split_expert_params_mapping = fused_moe_make_expert_params_mapping(
756 + self,
757 + ckpt_gate_proj_name="gate_proj",
758 + ckpt_down_proj_name="down_proj",
759 + ckpt_up_proj_name="up_proj",
760 + num_experts=num_experts,
761 + )
762 + fused_expert_param_names: dict[tuple[str, str], str] = {}
763 + for param_name in params_dict:
764 + for tag in ("w13_weight", "w2_weight"):
765 + if param_name.endswith(tag) and ".experts." in param_name:
766 + base = param_name.split(".experts.")[0]
767 + fused_expert_param_names[base, tag] = param_name
768 + stacked_mapping = [
769 + (".gate_up_proj", ".gate_proj", 0),
770 + (".gate_up_proj", ".up_proj", 1),
771 + # MLA runs both latent down-projections as one GEMM.
772 + (".fused_qkv_a_proj", ".q_a_proj", 0),
773 + (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
774 + ]
775 + # Sparse DSA draft blocks build an Indexer whose wk / weights_proj are
776 + # fused into a single MergedColumnParallelLinear (wk_weights_proj).
777 + indexer_stacked_mapping = [
778 + (".wk_weights_proj", ".wk", 0),
779 + (".wk_weights_proj", ".weights_proj", 1),
780 + ]
781 + # FP8 indexer wk dequant buffer (weight and scale arrive separately).
782 + pending_wk_fp8: dict[str, dict[str, torch.Tensor]] = {}
783 +
784 + for name, loaded_weight in weights:
785 + if name in shared_weights:
786 + target_name = shared_weights[name]
787 + if target_name in params_dict:
788 + param = params_dict[target_name]
789 + weight_loader = getattr(
790 + param, "weight_loader", default_weight_loader
791 + )
792 + weight_loader(param, loaded_weight)
793 + loaded_params.add(target_name)
794 + continue
795 +
796 + spec_layer = None
797 + if name.startswith("model.mtp_layers."):
798 + parts = name.split(".")
799 + if len(parts) > 3 and parts[2].isdigit():
800 + spec_layer = mtp_start + int(parts[2])
801 + name = name.replace(
802 + f"model.mtp_layers.{parts[2]}.",
803 + f"model.layers.{spec_layer}.",
804 + )
805 + else:
806 + spec_layer = _get_spec_layer_idx_from_weight_name(self.config, name)
807 +
808 + if spec_layer is None:
809 + continue
810 +
811 + name = self._rewrite_spec_layer_name(spec_layer, name)
812 + if name == "__skip__":
813 + continue
814 + name, loaded_weight = _prepare_mtp_fp8_expert_scale(
815 + self.quant_config, name, loaded_weight
816 + )
817 +
818 + if "mlp.gate.e_score_correction_bias" in name:
819 + name = name.replace("gate.e_score_correction_bias", "expert_bias")
820 +
821 + is_loaded = False
822 + for param_name, weight_name, shard_id in stacked_mapping:
823 + if weight_name not in name or ".experts." in name:
824 + continue
825 + name_mapped = name.replace(weight_name, param_name)
826 + if name_mapped not in params_dict:
827 + if is_pp_missing_parameter(name_mapped, self):
828 + is_loaded = True
829 + break
830 + param = params_dict[name_mapped]
831 + param.weight_loader(param, loaded_weight, shard_id)
832 + loaded_params.add(name_mapped)
833 + is_loaded = True
834 + break
835 + if is_loaded:
836 + continue
837 +
838 + # FP8 indexer wk: dequantize to BF16 and load into the fused
839 + # wk_weights_proj. PP-aware (skips layers not held by this rank).
840 + if _try_load_fp8_indexer_wk(
841 + name,
842 + loaded_weight,
843 + pending_wk_fp8,
844 + params_dict,
845 + loaded_params,
846 + pp_missing_layer_names,
847 + ):
848 + continue
849 +
850 + # BF16 indexer wk / weights_proj: merge into the fused param.
851 + is_loaded = False
852 + for param_name, weight_name, shard_id in indexer_stacked_mapping:
853 + if weight_name not in name or "wk_weights" in name:
854 + continue
855 + name_mapped = name.replace(weight_name, param_name)
856 + if name_mapped not in params_dict:
857 + if is_pp_missing_parameter(name_mapped, self):
858 + is_loaded = True
859 + break
860 + param = params_dict[name_mapped]
861 + param.weight_loader(param, loaded_weight, shard_id)
862 + loaded_params.add(name_mapped)
863 + is_loaded = True
864 + break
865 + if is_loaded:
866 + continue
867 +
868 + is_loaded = False
869 + if ".experts." in name:
870 + is_loaded = self._load_expert_weight(
871 + name,
872 + loaded_weight,
873 + params_dict,
874 + loaded_params,
875 + split_expert_params_mapping,
876 + fused_expert_param_names,
877 + num_experts,
878 + )
879 + if is_loaded:
880 + continue
881 +
882 + if "learnable_sink_param" in name:
883 + if name in params_dict:
884 + narrow_weight = loaded_weight[head_rank_start:head_rank_end]
885 + n = narrow_weight.shape[0]
886 + with torch.no_grad():
887 + params_dict[name][:n].copy_(narrow_weight)
888 + loaded_params.add(name)
889 + continue
890 +
891 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)
892 + if remapped_name is None:
893 + continue
894 + name = remapped_name
895 + if name not in params_dict:
896 + if is_pp_missing_parameter(name, self):
897 + continue
898 + if _should_skip_missing_mtp_scale_param(self.quant_config, name):
899 + continue
900 + logger.warning_once("Skipping unknown MTP weight: %s", name)
901 + continue
902 +
903 + param = params_dict[name]
904 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
905 + weight_loader(param, loaded_weight)
906 + loaded_params.add(name)
907 +
908 + logger.info_once("HYV4 MTP draft model loaded: %d params", len(loaded_params))
909 + unassigned_all = sorted(set(params_dict) - loaded_params)
910 + # KVCacheScaleParameter (``k_scale`` / ``v_scale`` / ``q_scale`` /
911 + # ``prob_scale``) is created by ``BaseKVCacheMethod.create_weights`` for
912 + # every Attention layer that runs on an fp8-family quant method, and
913 + # its -1.0 sentinel is replaced by the runtime default (1.0) in
914 + # ``process_weights_after_loading`` -- which runs *after* this method.
915 + # HY V4 checkpoints intentionally omit these (``kv_cache_quant_algo:
916 + # null``); silence them only when the parameter is actually the
917 + # sentinel type so a genuine mismatch on a same-named tensor still
918 + # reaches the warning.
919 + from vllm.model_executor.layers.quantization.kv_cache import (
920 + KVCacheScaleParameter,
921 + )
922 +
923 + unassigned = [
924 + name
925 + for name in unassigned_all
926 + if not isinstance(params_dict.get(name), KVCacheScaleParameter)
927 + ]
928 + if unassigned:
929 + # A draft parameter with no checkpoint source keeps its sentinel
930 + # init value (FP8 scales start at finfo(float32).min), which
931 + # silently destroys draft quality instead of failing the load.
932 + logger.warning(
933 + "HYV4 MTP draft model: %d parameters received no checkpoint value: %s",
934 + len(unassigned),
935 + ", ".join(unassigned),
936 + )
937 + return loaded_params
vllm/parser/harmony.py
+2 -1
@@ -451,9 +451,10 @@ def get_harmony_structural_tag(
451 builtin_tools: list[BuiltinToolParam],
452 tool_choice: SimplifiedToolChoice,
453 reasoning: bool,
454 + token_suffix: str = "",
455 ) -> StructuralTag:
456 # reasoning always enabled for Harmony
456 - del reasoning
457 + del reasoning, token_suffix
458
459 if builtin_tools:
460 # Fallback for built-in tools
vllm/reasoning/__init__.py
+4
@@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = {
84 "hy_v3_reasoning_parser",
85 "HYV3ReasoningParser",
86 ),
87 + "hy_v4": (
88 + "hy_v4_reasoning_parser",
89 + "HYV4ReasoningParser",
90 + ),
91 "kimi_k2": (
92 "kimi_k2_reasoning_parser",
93 "KimiK2ReasoningParser",
vllm/reasoning/hy_v4_reasoning_parser.py new
+324
@@ -0,0 +1,324 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +from collections.abc import Iterable, Sequence
7 +from typing import TypedDict
8 +
9 +import regex as re
10 +
11 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
12 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage
13 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
14 +from vllm.logger import init_logger
15 +from vllm.reasoning import ReasoningParser
16 +from vllm.tokenizers import TokenizerLike
17 +
18 +logger = init_logger(__name__)
19 +
20 +# Must match the chat template's own default for ``reasoning_effort``.
21 +DEFAULT_REASONING_EFFORT = "high"
22 +
23 +
24 +class ReasoningDelta(TypedDict):
25 + """One streaming reasoning delta. Either/both fields may be None."""
26 +
27 + reasoning: str | None
28 + content: str | None
29 +
30 +
31 +# NOTE: mirrored in ``vllm.tool_parsers.hy_v4_tool_parser`` (same pattern as
32 +# ``gemma4_utils``) so neither package depends on the other.
33 +def detect_token_suffix(tokenizer: TokenizerLike) -> str:
34 + """Detect the per-checkpoint suffix used by Hunyuan structural tokens.
35 +
36 + Args:
37 + tokenizer: Tokenizer of the served checkpoint.
38 +
39 + Returns:
40 + The suffix including its leading colon (e.g. ``":6124c78e"``), or ``""``
41 + when the checkpoint uses unsuffixed tokens.
42 +
43 + Raises:
44 + RuntimeError: The tokenizer declares the structural tokens through
45 + ``model_specific_special_tokens``, which transformers 5 no longer
46 + round-trips.
47 + """
48 +
49 + import transformers
50 +
51 + if int(transformers.__version__.split(".")[0]) >= 5:
52 + init_kwargs = getattr(tokenizer, "init_kwargs", None) or {}
53 + think_begin_as_special = init_kwargs.get(
54 + "model_specific_special_tokens", {}
55 + ).get("think_begin_token", "")
56 + if think_begin_as_special:
57 + raise RuntimeError(
58 + "This checkpoint declares HYV4 structural tokens (think_begin_token"
59 + "/toolcalls_begin_token/argkey_begin_token) in "
60 + "tokenizer_config.json, which transformers 5 no longer supports. "
61 + "Remove those fields and keep the tokens in the tokenizer's own "
62 + "token definitions so the suffix can be read from the vocab."
63 + )
64 +
65 + structural_token_re = re.compile(
66 + r"<(?:think|tool_calls|tool_call|arg_key|arg_value)(:[^\s>]+)?>"
67 + )
68 + for token in tokenizer.get_vocab():
69 + match = structural_token_re.fullmatch(token)
70 + if match:
71 + return match.group(1) or ""
72 +
73 + return ""
74 +
75 +
76 +class HYV4ReasoningExtractor:
77 + """Reasoning extraction for HYV4, on plain data.
78 +
79 + ``thinking`` selects the mode:
80 + - ``True``: parse ``<think>...</think>``; the start token is injected at the
81 + end of the prompt and is therefore usually absent from the output.
82 + - ``False``: ``no_think`` -- the whole output is content, no reasoning.
83 + """
84 +
85 + def __init__(self, vocab: dict[str, int], token_suffix: str, thinking: bool = True):
86 + self.thinking = thinking
87 + self.start_token: str = f"<think{token_suffix}>"
88 + self.end_token: str = f"</think{token_suffix}>"
89 +
90 + start_token_id = vocab.get(self.start_token)
91 + end_token_id = vocab.get(self.end_token)
92 + if thinking and (start_token_id is None or end_token_id is None):
93 + raise RuntimeError(
94 + "HYV4 reasoning extractor could not locate think "
95 + "start/end tokens in the tokenizer!"
96 + )
97 + # In no_think mode the ids are unused (kept for symmetry / debugging).
98 + self.start_token_id: int | None = start_token_id
99 + self.end_token_id: int | None = end_token_id
100 +
101 + def has_reasoning_ended(self, token_ids: Sequence[int]) -> bool:
102 + """Idempotent STATE query: has ``</think>`` been emitted yet?
103 +
104 + Monotonic -- once the end token is present this stays True. Use for
105 + "has reasoning finished, given everything so far". (no_think is always
106 + ended.) Contrast with ``reasoning_ended_in_delta`` (edge detector).
107 + """
108 + if not self.thinking:
109 + return True
110 + for i in range(len(token_ids) - 1, -1, -1):
111 + if token_ids[i] == self.start_token_id:
112 + return False
113 + if token_ids[i] == self.end_token_id:
114 + return True
115 + return False
116 +
117 + def reasoning_ended_in_delta(
118 + self, token_ids: Sequence[int], delta_ids: Iterable[int]
119 + ) -> bool:
120 + """Edge detector (NOT idempotent): did reasoning end *in this* step?
121 +
122 + True only on the step whose ``delta_ids`` contains ``</think>``. Callers
123 + use it once behind their own "already ended" latch, so it need not be
124 + monotonic; it only checks the new delta (cheaper than scanning the whole
125 + sequence). Contrast with ``has_reasoning_ended`` (cumulative state).
126 + """
127 + if not self.thinking:
128 + return True
129 + end_token_id = self.end_token_id
130 + if end_token_id is None:
131 + return False
132 + return end_token_id in delta_ids
133 +
134 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
135 + """Return the token ids after ``</think>`` (all of them in no_think)."""
136 + if not self.thinking:
137 + return input_ids
138 + head = input_ids[: max(0, len(input_ids) - 1)]
139 + if self.end_token_id not in head:
140 + return []
141 + return input_ids[input_ids.index(self.end_token_id) + 1 :]
142 +
143 + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
144 + """Count reasoning tokens (0 in no_think).
145 +
146 + The ``<think>`` start token is injected at the END of the prompt, so it
147 + is normally absent from ``token_ids``. Only skip a leading start token
148 + when it actually appears first (legacy / no prompt-injection); otherwise
149 + count every token up to ``</think>``.
150 + """
151 + if not self.thinking:
152 + return 0
153 + start = 1 if (token_ids and token_ids[0] == self.start_token_id) else 0
154 + try:
155 + return token_ids.index(self.end_token_id, start) - start
156 + except ValueError:
157 + # No end token -- reasoning was truncated; everything counts.
158 + return len(token_ids) - start
159 +
160 + def extract_reasoning(self, model_output: str) -> tuple[str | None, str | None]:
161 + """Split a complete output into ``(reasoning, content)``.
162 +
163 + In no_think mode reasoning is None and the whole output is content. The
164 + start token may be absent (injected in the prompt), so reasoning is
165 + assumed to begin at the start of the output.
166 + """
167 + if not self.thinking:
168 + return None, model_output
169 +
170 + parts = model_output.partition(self.start_token)
171 + model_output = parts[2] if parts[1] else parts[0]
172 +
173 + if self.end_token not in model_output:
174 + return model_output, None
175 + reasoning, _, content = model_output.partition(self.end_token)
176 + return reasoning, content or None
177 +
178 + def extract_reasoning_streaming(
179 + self,
180 + previous_text: str,
181 + current_text: str,
182 + delta_text: str,
183 + previous_token_ids: Sequence[int],
184 + current_token_ids: Sequence[int],
185 + delta_token_ids: Sequence[int],
186 + ) -> ReasoningDelta | None:
187 + """Classify a streaming delta as reasoning and/or content.
188 +
189 + In no_think mode every delta is content. Otherwise: same logic as vLLM's
190 + generic thinking parser, except the final branch (no start token seen
191 + anywhere) treats output as reasoning-then-content, because HYV4 injects
192 + ``<think>`` in the prompt (absent from the output).
193 + """
194 + if not self.thinking:
195 + if not delta_text:
196 + return None
197 + return ReasoningDelta(reasoning=None, content=delta_text)
198 +
199 + # Skip a delta that is a single lone start/end special token.
200 + if len(delta_token_ids) == 1 and delta_token_ids[0] in (
201 + self.start_token_id,
202 + self.end_token_id,
203 + ):
204 + return None
205 +
206 + if self.start_token_id in previous_token_ids:
207 + if self.end_token_id in delta_token_ids:
208 + end_index = delta_text.find(self.end_token)
209 + reasoning = delta_text[:end_index]
210 + content = delta_text[end_index + len(self.end_token) :]
211 + return ReasoningDelta(reasoning=reasoning, content=content or None)
212 + elif self.end_token_id in previous_token_ids:
213 + return ReasoningDelta(reasoning=None, content=delta_text)
214 + else:
215 + return ReasoningDelta(reasoning=delta_text, content=None)
216 + elif self.start_token_id in delta_token_ids:
217 + if self.end_token_id in delta_token_ids:
218 + start_index = delta_text.find(self.start_token)
219 + end_index = delta_text.find(self.end_token)
220 + reasoning = delta_text[start_index + len(self.start_token) : end_index]
221 + content = delta_text[end_index + len(self.end_token) :]
222 + return ReasoningDelta(reasoning=reasoning, content=content or None)
223 + else:
224 + return ReasoningDelta(reasoning=delta_text, content=None)
225 + else:
226 + # HYV4: <think> was injected at the end of the prompt, so it never
227 + # appears in the output. Treat the stream as reasoning until
228 + # </think>, then content.
229 + if self.end_token_id in delta_token_ids:
230 + end_index = delta_text.find(self.end_token)
231 + reasoning = delta_text[:end_index]
232 + content = delta_text[end_index + len(self.end_token) :]
233 + return ReasoningDelta(reasoning=reasoning, content=content or None)
234 + elif self.end_token_id in previous_token_ids:
235 + return ReasoningDelta(reasoning=None, content=delta_text)
236 + else:
237 + return ReasoningDelta(reasoning=delta_text, content=None)
238 +
239 +
240 +def build_reasoning_extractor(
241 + tokenizer: TokenizerLike, *, thinking: bool = True
242 +) -> HYV4ReasoningExtractor:
243 + return HYV4ReasoningExtractor(
244 + tokenizer.get_vocab(), detect_token_suffix(tokenizer), thinking
245 + )
246 +
247 +
248 +class HYV4ReasoningParser(ReasoningParser):
249 + """vLLM adapter around :class:`HYV4ReasoningExtractor`.
250 +
251 + Decides ``thinking`` from the request's ``reasoning_effort`` and converts
252 + the extractor's ``ReasoningDelta`` to ``DeltaMessage``.
253 + """
254 +
255 + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
256 + super().__init__(tokenizer, *args, **kwargs)
257 +
258 + # Decide thinking vs no_think from the request; the extractor handles
259 + # both modes internally. Prefer chat_template_kwargs.reasoning_effort,
260 + # then the top-level reasoning_effort. The default must match the chat
261 + # template, which sets reasoning_effort='high' when the request omits
262 + # it -- defaulting to no_think here would return the whole
263 + # chain-of-thought (and a stray ``</think:SUF>``) as content.
264 + chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
265 + reasoning_effort = chat_kwargs.get("reasoning_effort", None)
266 + if not reasoning_effort:
267 + reasoning_effort = (
268 + kwargs.get("reasoning_effort") or DEFAULT_REASONING_EFFORT
269 + )
270 + logger.debug("reasoning_effort for choosing parser: %s", reasoning_effort)
271 + thinking = reasoning_effort != "no_think"
272 +
273 + self._extractor = build_reasoning_extractor(tokenizer, thinking=thinking)
274 +
275 + # Reported regardless of the per-request thinking mode: ``ReasoningConfig``
276 + # is engine-wide and needs the checkpoint's delimiters to derive the token
277 + # ids used for thinking budget enforcement.
278 + @property
279 + def reasoning_start_str(self) -> str:
280 + return self._extractor.start_token
281 +
282 + @property
283 + def reasoning_end_str(self) -> str:
284 + return self._extractor.end_token
285 +
286 + def is_reasoning_end(self, input_ids) -> bool:
287 + return self._extractor.has_reasoning_ended(input_ids)
288 +
289 + def is_reasoning_end_streaming(self, input_ids, delta_ids) -> bool:
290 + return self._extractor.reasoning_ended_in_delta(input_ids, delta_ids)
291 +
292 + def extract_content_ids(self, input_ids) -> list[int]:
293 + return self._extractor.extract_content_ids(input_ids)
294 +
295 + def count_reasoning_tokens(self, token_ids) -> int:
296 + return self._extractor.count_reasoning_tokens(token_ids)
297 +
298 + def extract_reasoning(
299 + self,
300 + model_output: str,
301 + request: ChatCompletionRequest | ResponsesRequest,
302 + ) -> tuple[str | None, str | None]:
303 + return self._extractor.extract_reasoning(model_output)
304 +
305 + def extract_reasoning_streaming(
306 + self,
307 + previous_text: str,
308 + current_text: str,
309 + delta_text: str,
310 + previous_token_ids,
311 + current_token_ids,
312 + delta_token_ids,
313 + ) -> DeltaMessage | None:
314 + delta = self._extractor.extract_reasoning_streaming(
315 + previous_text,
316 + current_text,
317 + delta_text,
318 + previous_token_ids,
319 + current_token_ids,
320 + delta_token_ids,
321 + )
322 + if delta is None:
323 + return None
324 + return DeltaMessage(reasoning=delta["reasoning"], content=delta["content"])
vllm/tool_parsers/__init__.py
+4
@@ -94,6 +94,10 @@ _TOOL_PARSERS_TO_REGISTER = {
94 "hy_v3_tool_parser",
95 "HYV3ToolParser",
96 ),
97 + "hy_v4": (
98 + "hy_v4_tool_parser",
99 + "HYV4ToolParser",
100 + ),
101 "internlm": (
102 "internlm2_tool_parser",
103 "Internlm2ToolParser",
vllm/tool_parsers/hy_v4_tool_parser.py new
+1166
@@ -0,0 +1,1166 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import ast
7 +import json
8 +from typing import Any, TypedDict
9 +
10 +import regex as re
11 +
12 +import vllm.envs as envs
13 +from vllm.entrypoints.chat_utils import make_tool_call_id
14 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
15 +from vllm.entrypoints.openai.engine.protocol import (
16 + DeltaFunctionCall,
17 + DeltaMessage,
18 + DeltaToolCall,
19 + ExtractedToolCallInformation,
20 + FunctionCall,
21 + ToolCall,
22 +)
23 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24 +from vllm.logger import init_logger
25 +from vllm.tokenizers import TokenizerLike
26 +from vllm.tool_parsers.abstract_tool_parser import ToolParser
27 +from vllm.tool_parsers.utils import partial_tag_overlap
28 +
29 +logger = init_logger(__name__)
30 +
31 +# =============================================================================
32 +# Plain-dict contract with the vLLM-side wrapper
33 +# =============================================================================
34 +
35 +
36 +class ToolSchema(TypedDict):
37 + """A tool as seen by this parser (built by the wrapper from ``request.tools``)."""
38 +
39 + name: str
40 + parameters: dict[str, Any] | None # JSON Schema, or None
41 +
42 +
43 +class ToolCallDict(TypedDict):
44 + """One extracted tool call. ``arguments`` is a JSON-encoded string."""
45 +
46 + name: str
47 + arguments: str
48 +
49 +
50 +class ExtractResult(TypedDict):
51 + """Return of ``extract_tool_calls`` (non-streaming)."""
52 +
53 + tools_called: bool
54 + content: str | None
55 + tool_calls: list[ToolCallDict]
56 +
57 +
58 +class StreamToolCall(TypedDict):
59 + """One streaming tool-call delta.
60 +
61 + ``new`` is True only on the chunk that first emits this tool's name -- the
62 + wrapper mints an id + sets type on that chunk only. ``arguments`` is the
63 + incremental JSON diff (or None).
64 + """
65 +
66 + index: int
67 + new: bool
68 + name: str | None
69 + arguments: str | None
70 +
71 +
72 +class StreamDelta(TypedDict):
73 + """Return of ``extract_tool_calls_streaming``.
74 +
75 + A single engine delta can carry content plus every tool call that was
76 + drained from the buffer (batched decode), so ``tool_calls`` is not limited
77 + to one entry.
78 + """
79 +
80 + content: str | None
81 + tool_calls: list[StreamToolCall]
82 +
83 +
84 +# NOTE: mirrored in ``vllm.reasoning.hy_v4_reasoning_parser`` (same pattern as
85 +# ``gemma4_utils``) so neither package depends on the other.
86 +def detect_token_suffix(tokenizer: TokenizerLike) -> str:
87 + """Detect the per-checkpoint suffix used by Hunyuan structural tokens.
88 +
89 + Args:
90 + tokenizer: Tokenizer of the served checkpoint.
91 +
92 + Returns:
93 + The suffix including its leading colon (e.g. ``":6124c78e"``), or ``""``
94 + when the checkpoint uses unsuffixed tokens.
95 +
96 + Raises:
97 + RuntimeError: The tokenizer declares the structural tokens through
98 + ``model_specific_special_tokens``, which transformers 5 no longer
99 + round-trips.
100 + """
101 +
102 + import transformers
103 +
104 + if int(transformers.__version__.split(".")[0]) >= 5:
105 + init_kwargs = getattr(tokenizer, "init_kwargs", None) or {}
106 + think_begin_as_special = init_kwargs.get(
107 + "model_specific_special_tokens", {}
108 + ).get("think_begin_token", "")
109 + if think_begin_as_special:
110 + raise RuntimeError(
111 + "This checkpoint declares HYV4 structural tokens (think_begin_token"
112 + "/toolcalls_begin_token/argkey_begin_token) in "
113 + "tokenizer_config.json, which transformers 5 no longer supports. "
114 + "Remove those fields and keep the tokens in the tokenizer's own "
115 + "token definitions so the suffix can be read from the vocab."
116 + )
117 +
118 + structural_token_re = re.compile(
119 + r"<(?:think|tool_calls|tool_call|arg_key|arg_value)(:[^\s>]+)?>"
120 + )
121 + for token in tokenizer.get_vocab():
122 + match = structural_token_re.fullmatch(token)
123 + if match:
124 + return match.group(1) or ""
125 +
126 + return ""
127 +
128 +
129 +# =============================================================================
130 +# Argument-value parsing utilities (stateless, pure functions)
131 +#
132 +# These map the model's raw ``<arg_value>`` text to typed Python values using
133 +# the tool's JSON Schema. They do not touch any parser state, so they live at
134 +# module level instead of tangling up HYV4ToolExtractor.
135 +# =============================================================================
136 +
137 +_TYPE_ALIASES: dict[str, str] = {
138 + "str": "string",
139 + "text": "string",
140 + "varchar": "string",
141 + "char": "string",
142 + "enum": "string",
143 + "bool": "boolean",
144 + "binary": "boolean",
145 + "int": "integer",
146 + "float": "number",
147 + "double": "number",
148 + "list": "array",
149 + "dict": "object",
150 + "map": "object",
151 +}
152 +
153 +# Prefix-based wildcard matching for non-standard type names.
154 +# Following the same approach as qwen3coder_tool_parser._convert_param_value
155 +# which uses param_type.startswith("int"), startswith("uint"), etc.
156 +_INTEGER_PREFIXES: tuple[str, ...] = ("int", "uint", "long", "short", "unsigned")
157 +_NUMBER_PREFIXES: tuple[str, ...] = ("num", "float")
158 +
159 +
160 +def _normalize_type(raw_type: str) -> str:
161 + """Map non-standard type aliases to JSON Schema standard names.
162 +
163 + First performs exact lookup in _TYPE_ALIASES. On miss, falls back to
164 + prefix-based matching using startswith()
165 + - int*/uint*/long*/short*/unsigned* → "integer"
166 + - num*/float* → "number"
167 + """
168 + exact = _TYPE_ALIASES.get(raw_type)
169 + if exact is not None:
170 + return exact
171 + lower = raw_type.lower()
172 + if any(lower.startswith(p) for p in _INTEGER_PREFIXES):
173 + return "integer"
174 + if any(lower.startswith(p) for p in _NUMBER_PREFIXES):
175 + return "number"
176 + return raw_type
177 +
178 +
179 +def _get_arg_schema(
180 + function_name: str,
181 + arg_key: str,
182 + tools: list[ToolSchema] | None,
183 +) -> dict[str, Any]:
184 + """Look up a specific argument's property schema from the tools list."""
185 + if tools is None:
186 + return {}
187 + for tool in tools:
188 + if tool.get("name") == function_name:
189 + parameters = tool.get("parameters")
190 + if parameters is None:
191 + return {}
192 + return parameters.get("properties", {}).get(arg_key, {})
193 + logger.warning_once("No tool named '%s'.", function_name)
194 + return {}
195 +
196 +
197 +def _get_schema_options(arg_schema: dict[str, Any]) -> list[dict[str, Any]]:
198 + """Normalize any property schema into a list of sub-schemas.
199 + - has type (single type) → return [arg_schema]
200 + - anyOf → return the anyOf list
201 + - oneOf → return the oneOf list
202 + - fallback → [{"type": "string"}]
203 +
204 + Note: single ``type`` has the highest priority.
205 + """
206 + if "type" in arg_schema:
207 + type_val = arg_schema["type"]
208 + # JSON Schema allows "type" to be an array to represent union types,
209 + # e.g. "type": ["string", "object"].
210 + # Expand it into an anyOf-equivalent format:
211 + # [{"type": "string"}, {"type": "object"}]
212 + # so that _get_types / _parse_value can handle it uniformly later.
213 + if isinstance(type_val, list):
214 + return [{"type": t} for t in type_val]
215 + return [arg_schema]
216 + if "anyOf" in arg_schema:
217 + return arg_schema["anyOf"]
218 + if "oneOf" in arg_schema:
219 + return arg_schema["oneOf"]
220 +
221 + return [{"type": "string"}]
222 +
223 +
224 +def _get_types(arg_schema: dict[str, Any]) -> set[str]:
225 + """Extract normalized, non-null type set from a property schema."""
226 + schemas = _get_schema_options(arg_schema)
227 + return {_normalize_type(s.get("type", "string")) for s in schemas} - {"null"}
228 +
229 +
230 +def _is_only_string_type(
231 + function_name: str,
232 + arg_key: str,
233 + tools: list[ToolSchema] | None,
234 +) -> bool:
235 + """Return True if the parameter's type set is exactly {"string"}.
236 +
237 + Only pure string types get partial value streaming; compound types like
238 + anyOf(string | array) do not, since the partial value might end up being a
239 + JSON array or object.
240 + """
241 + types = _get_types(_get_arg_schema(function_name, arg_key, tools))
242 + return types == {"string"}
243 +
244 +
245 +def _try_parse_bool(value: str) -> bool | None:
246 + """Try to parse a string as bool; return None on failure."""
247 + lower = value.lower()
248 + if lower == "true":
249 + return True
250 + elif lower == "false":
251 + return False
252 + return None
253 +
254 +
255 +def _try_parse_int(value: str) -> int | None:
256 + """Try to parse a string as int; return None on failure."""
257 + try:
258 + return int(value)
259 + except (ValueError, TypeError):
260 + return None
261 +
262 +
263 +def _try_parse_wildcard_number(value: str) -> int | float | None:
264 + """Try to parse a string as a number (int or float).
265 +
266 + Decision rule: if the string contains '.' or 'e'/'E' (scientific notation),
267 + parse as float; otherwise parse as int.
268 +
269 + Examples: "5" → 5, "5.0" → 5.0, "1e3" → 1000.0, "-3" → -3.
270 + Return None on failure.
271 + """
272 + try:
273 + if "." in value or "e" in value or "E" in value:
274 + return float(value)
275 + return int(value)
276 + except (ValueError, TypeError):
277 + return None
278 +
279 +
280 +def _deserialize(value: str) -> Any:
281 + """Deserialize a string value using json.loads then ast.literal_eval."""
282 + try:
283 + return json.loads(value)
284 + except Exception:
285 + pass
286 + try:
287 + return ast.literal_eval(value)
288 + except Exception:
289 + pass
290 + return value
291 +
292 +
293 +def _parse_value(
294 + value: str,
295 + function_name: str,
296 + arg_key: str,
297 + tools: list[ToolSchema] | None,
298 +) -> Any:
299 + """Unified argument value parser with anyOf/oneOf support.
300 +
301 + Fallthrough chain:
302 + bool → int → number(wildcard_number)
303 + → json.loads for array/object → string → _deserialize
304 + """
305 + types = _get_types(_get_arg_schema(function_name, arg_key, tools))
306 +
307 + # 1. Try bool
308 + if "boolean" in types:
309 + result_bool = _try_parse_bool(value)
310 + if result_bool is not None:
311 + return result_bool
312 +
313 + # 2. Try int
314 + if "integer" in types:
315 + result_int = _try_parse_int(value)
316 + if result_int is not None:
317 + return result_int
318 +
319 + # 3. Try number (wildcard_number: int if no '.'/e/E, float otherwise)
320 + if "number" in types:
321 + result_number = _try_parse_wildcard_number(value)
322 + if result_number is not None:
323 + return result_number
324 +
325 + # 4. Try json.loads (covers array/object and other unlisted types)
326 + if types - {"string", "boolean", "integer", "number"}:
327 + try:
328 + return json.loads(value)
329 + except (json.JSONDecodeError, ValueError):
330 + pass
331 +
332 + # 5. String fallback
333 + if "string" in types:
334 + return value
335 +
336 + # 6. Final fallback
337 + return _deserialize(value)
338 +
339 +
340 +# =============================================================================
341 +# Stateful tool-call parser
342 +# =============================================================================
343 +
344 +
345 +class HYV4ToolExtractor:
346 + """Pure tool-call parsing logic for HYV4. Returns plain dicts; no vLLM types.
347 +
348 + Holds only parser state: the structural token strings/ids, the compiled
349 + regexes, and the streaming incremental state. Stateless argument parsing is
350 + delegated to the module-level utilities above.
351 +
352 + Streaming has two entry paths sharing one incremental argument parser,
353 + selected by the ``guided`` flag of :meth:`extract_tool_calls_streaming`:
354 + normal / auto tool choice detects markers on atomic token ids so ordinary
355 + ``<`` in content is never withheld, while structural-tag guided decoding
356 + detects them on the string level because the grammar may split a marker
357 + across sub-word tokens.
358 + """
359 +
360 + def __init__(self, vocab: dict[str, int], token_suffix: str, strict: bool):
361 + self._strict = strict
362 +
363 + # Kept so the wrapper can build a matching structural tag for
364 + # constrained decoding without re-scanning the tokenizer vocab.
365 + self.token_suffix: str = token_suffix
366 +
367 + # Running state read by the vLLM serving layer (mirrored by the wrapper).
368 + self.prev_tool_call_arr: list[dict] = []
369 + self.current_tool_id: int = -1
370 + self.streamed_args_for_tool: list[str] = []
371 +
372 + # Streaming state: send tool name first, then return arguments at once
373 + self._streaming_tool_name: str | None = None # tool name being streamed
374 +
375 + # State fields for incremental argument streaming
376 + self._completed_args: dict = {} # closed {key: parsed_value}
377 + self._current_arg_key: str | None = None # key being collected
378 + self._current_arg_is_string: bool = False # is current arg pure string?
379 + self._streamed_json_len: int = 0 # bytes of JSON already sent
380 +
381 + self.tool_calls_start_token: str = f"<tool_calls{token_suffix}>"
382 + self.tool_calls_end_token: str = f"</tool_calls{token_suffix}>"
383 +
384 + self.tool_call_start_token: str = f"<tool_call{token_suffix}>"
385 + self.tool_call_end_token: str = f"</tool_call{token_suffix}>"
386 +
387 + self.arg_key_start_token: str = f"<arg_key{token_suffix}>"
388 + self.arg_key_end_token: str = f"</arg_key{token_suffix}>"
389 +
390 + self.arg_value_start_token: str = f"<arg_value{token_suffix}>"
391 + self.arg_value_end_token: str = f"</arg_value{token_suffix}>"
392 +
393 + logger.debug(
394 + "HYV4ToolExtractor structural tokens: %s ... %s, suffix=%r, strict=%s",
395 + self.tool_calls_start_token,
396 + self.tool_calls_end_token,
397 + token_suffix,
398 + strict,
399 + )
400 +
401 + # The suffix comes from the tokenizer vocab, so it must not be treated
402 + # as a pattern.
403 + self.tool_call_regex = re.compile(
404 + rf"{re.escape(self.tool_call_start_token)}(.*?)"
405 + rf"{re.escape(self.tool_call_end_token)}",
406 + re.DOTALL,
407 + )
408 +
409 + self.func_args_regex = re.compile(
410 + rf"{re.escape(self.arg_key_start_token)}(.*?)"
411 + rf"{re.escape(self.arg_key_end_token)}"
412 + rf"{re.escape(self.arg_value_start_token)}(.*?)"
413 + rf"{re.escape(self.arg_value_end_token)}",
414 + re.DOTALL,
415 + )
416 +
417 + self.tool_calls_start_token_id = vocab.get(self.tool_calls_start_token)
418 + self.tool_calls_end_token_id = vocab.get(self.tool_calls_end_token)
419 +
420 + self.tool_call_start_token_id = vocab.get(self.tool_call_start_token)
421 + self.tool_call_end_token_id = vocab.get(self.tool_call_end_token)
422 + self._buffer = ""
423 +
424 + if (
425 + self.tool_calls_start_token_id is None
426 + or self.tool_calls_end_token_id is None
427 + ):
428 + raise RuntimeError(
429 + "HYV4 tool extractor could not locate tool call "
430 + "start/end tokens in the tokenizer!"
431 + )
432 +
433 + def _extract_tool_calls(
434 + self,
435 + model_output: str,
436 + tools: list[ToolSchema] | None,
437 + ) -> list[ToolCallDict]:
438 + try:
439 + tool_calls: list[ToolCallDict] = []
440 + function_calls = self.tool_call_regex.findall(model_output)
441 + for tool_call_body in function_calls:
442 + arg_start = tool_call_body.find(self.arg_key_start_token)
443 + if arg_start == -1:
444 + function_name = tool_call_body.strip()
445 + function_args = ""
446 + else:
447 + function_name = tool_call_body[:arg_start].strip()
448 + function_args = tool_call_body[arg_start:].strip()
449 +
450 + arg_pairs = self.func_args_regex.findall(function_args)
451 + arg_dict = {}
452 + for key, value in arg_pairs:
453 + arg_dict[key] = _parse_value(value, function_name, key, tools)
454 + tool_calls.append(
455 + ToolCallDict(
456 + name=function_name,
457 + arguments=json.dumps(arg_dict, ensure_ascii=False),
458 + )
459 + )
460 + return tool_calls
461 + except Exception:
462 + logger.exception("Error in extracting tool call from response.")
463 + return []
464 +
465 + def _extract_tool_calls_strict(
466 + self,
467 + model_output: str,
468 + tools: list[ToolSchema] | None,
469 + ) -> list[ToolCallDict]:
470 + start_idx = model_output.find(self.tool_calls_start_token)
471 + end_idx = model_output.find(self.tool_calls_end_token)
472 + if start_idx == -1 or end_idx == -1 or end_idx < start_idx:
473 + raise ValueError("Malformed tool block wrapper.")
474 +
475 + tool_block = model_output[
476 + start_idx + len(self.tool_calls_start_token) : end_idx
477 + ]
478 + if tool_block.count(self.tool_call_start_token) != tool_block.count(
479 + self.tool_call_end_token
480 + ):
481 + raise ValueError("Malformed tool_call wrapper count.")
482 +
483 + tool_calls: list[ToolCallDict] = []
484 + function_calls = self.tool_call_regex.findall(tool_block)
485 + if tool_block.count(self.tool_call_start_token) != len(function_calls):
486 + raise ValueError("Malformed tool_call body.")
487 +
488 + for tool_call_body in function_calls:
489 + arg_start = tool_call_body.find(self.arg_key_start_token)
490 + if arg_start == -1:
491 + function_name = tool_call_body.strip()
492 + function_args = ""
493 + else:
494 + function_name = tool_call_body[:arg_start].strip()
495 + function_args = tool_call_body[arg_start:].strip()
496 +
497 + if not function_name:
498 + raise ValueError("Empty function name in tool call.")
499 +
500 + if function_args.count(self.arg_key_start_token) != function_args.count(
501 + self.arg_key_end_token
502 + ) or function_args.count(self.arg_value_start_token) != function_args.count(
503 + self.arg_value_end_token
504 + ):
505 + raise ValueError("Malformed argument tag count.")
506 +
507 + arg_pairs = self.func_args_regex.findall(function_args)
508 + remainder = self.func_args_regex.sub("", function_args)
509 + if remainder.strip():
510 + raise ValueError("Unparsed argument payload remains.")
511 +
512 + arg_dict = {}
513 + for key, value in arg_pairs:
514 + arg_dict[key] = _parse_value(value, function_name, key, tools)
515 +
516 + tool_calls.append(
517 + ToolCallDict(
518 + name=function_name,
519 + arguments=json.dumps(arg_dict, ensure_ascii=False),
520 + )
521 + )
522 +
523 + if not tool_calls:
524 + raise ValueError("No valid tool calls extracted.")
525 +
526 + return tool_calls
527 +
528 + def extract_tool_calls(
529 + self,
530 + model_output: str,
531 + tools: list[ToolSchema] | None,
532 + ) -> ExtractResult:
533 + # sanity check; avoid unnecessary processing
534 + if self.tool_calls_start_token not in model_output:
535 + return ExtractResult(
536 + tools_called=False, content=model_output, tool_calls=[]
537 + )
538 + else:
539 + try:
540 + if self._strict:
541 + tool_calls = self._extract_tool_calls_strict(model_output, tools)
542 + else:
543 + tool_calls = self._extract_tool_calls(model_output, tools)
544 + if not tool_calls:
545 + raise ValueError("No valid tool calls extracted.")
546 +
547 + s_index = model_output.find(self.tool_calls_start_token)
548 + content = model_output[:s_index] if s_index != -1 else model_output
549 + return ExtractResult(
550 + tools_called=True,
551 + content=content if content else None,
552 + tool_calls=tool_calls,
553 + )
554 +
555 + except Exception:
556 + logger.warning(
557 + "HYV4ToolExtractor detected malformed tool output; "
558 + "falling back to raw content."
559 + )
560 + return ExtractResult(
561 + tools_called=False,
562 + content=model_output,
563 + tool_calls=[],
564 + )
565 +
566 + def _reset_streaming_tool_state(self):
567 + """Reset the streaming state for a single tool call."""
568 + self._streaming_tool_name = None
569 + self._completed_args = {}
570 + self._current_arg_key = None
571 + self._current_arg_is_string = False
572 + self._streamed_json_len = 0
573 +
574 + def extract_tool_calls_streaming(
575 + self,
576 + previous_text: str,
577 + current_text: str,
578 + delta_text: str,
579 + previous_token_ids,
580 + current_token_ids,
581 + delta_token_ids,
582 + tools: list[ToolSchema] | None,
583 + *,
584 + guided: bool = False,
585 + ) -> StreamDelta | None:
586 + """Consume one engine delta and emit content / tool-call deltas.
587 +
588 + Two entry paths share the same incremental argument parser:
589 + - ``guided=False`` (normal / auto tool choice): the structural markers
590 + arrive as atomic tokenizer tokens, so detection is done on token ids
591 + and ordinary ``<`` in model content is never withheld.
592 + - ``guided=True`` (structural-tag constrained decoding): the grammar
593 + only constrains strings, so ``<tool_calls...>`` may be split across
594 + sub-word tokens and must be detected on the string level.
595 +
596 + Args:
597 + previous_text: Text decoded before this delta.
598 + current_text: Full text decoded so far.
599 + delta_text: Text decoded in this delta.
600 + previous_token_ids: Token ids before this delta.
601 + current_token_ids: Token ids decoded so far.
602 + delta_token_ids: Token ids in this delta.
603 + tools: Plain tool schemas used for argument typing.
604 + guided: Whether guided decoding may split markers into sub-word
605 + tokens.
606 +
607 + Returns:
608 + A streaming delta carrying content and/or the tool calls drained
609 + from the buffer, or None when nothing can be emitted yet.
610 + """
611 + content_delta: str | None = None
612 + tool_calls: list[StreamToolCall] = []
613 +
614 + if guided:
615 + # Guided path: a structural-tag grammar constrains strings, not
616 + # token ids, so the atomic ``tool_calls_start_token_id`` never
617 + # appears in the token stream and an id-based check would
618 + # misclassify every tool delta as plain content.
619 + marker = self.tool_calls_start_token
620 + marker_pos = current_text.find(marker)
621 +
622 + if marker_pos == -1:
623 + # Marker not seen yet. The tail of ``current_text`` might still
624 + # be a partial prefix of the marker split across deltas; hold
625 + # that tail back and only stream the safe portion as content.
626 + safe_end = len(current_text) - partial_tag_overlap(current_text, marker)
627 + # Emit only the newly-arrived, safe content (delta relative to
628 + # what was previously streamed).
629 + content = current_text[len(previous_text) : safe_end]
630 + return StreamDelta(
631 + content=content if content else None,
632 + tool_calls=[],
633 + )
634 +
635 + # Marker present. Stream any content that precedes it (only the
636 + # part we have not already emitted), then buffer the tool payload.
637 + #
638 + # The payload is everything after the marker in the *full* text. We
639 + # keep the buffer in sync with it by only appending the portion that
640 + # arrived in this delta, mirroring the incremental ``+=`` semantics
641 + # while remaining correct when the marker itself was split across
642 + # sub-word tokens.
643 + content_end = marker_pos
644 + content = current_text[len(previous_text) : content_end]
645 +
646 + payload_start = marker_pos + len(marker)
647 + # How much of the payload had already been buffered before this
648 + # delta.
649 + prev_payload_len = max(len(previous_text) - payload_start, 0)
650 + new_payload = current_text[payload_start + prev_payload_len :]
651 + self._buffer += new_payload
652 +
653 + if content:
654 + content_delta = content
655 + else:
656 + # Atomic-token path: normal / auto generation emits the HYV4
657 + # structural markers as single tokenizer tokens. Pass ordinary
658 + # content through exactly as decoded, including comparison
659 + # operators such as ``<``.
660 + if self.tool_calls_start_token_id not in current_token_ids:
661 + return StreamDelta(content=delta_text, tool_calls=[])
662 +
663 + if self.tool_calls_start_token in delta_text:
664 + text_parts = delta_text.split(self.tool_calls_start_token)
665 + self._buffer += text_parts[-1]
666 + if text_parts[0]:
667 + content_delta = text_parts[0]
668 + else:
669 + self._buffer += delta_text
670 +
671 + # Encountered finish, extract valid arguments
672 + if (
673 + current_text.find(self.tool_call_end_token + self.tool_calls_end_token)
674 + != -1
675 + and self._buffer.find(self.tool_call_end_token) == -1
676 + ):
677 + self._buffer += self.tool_call_end_token + self.tool_calls_end_token
678 +
679 + # Drain every complete tool call that is already buffered so a single
680 + # engine delta can carry content + N tool calls (batched decode).
681 + while True:
682 + cur_text = self._buffer
683 +
684 + # Haven't encountered tool_call start tag yet; drop the inter-tag
685 + # padding but keep any tail that may be a partial tool_call start
686 + # tag (guided decoding can split it across deltas).
687 + start_idx = cur_text.find(self.tool_call_start_token)
688 + if start_idx == -1 and self._streaming_tool_name is None:
689 + keep = partial_tag_overlap(cur_text, self.tool_call_start_token)
690 + self._buffer = cur_text[len(cur_text) - keep :] if keep else ""
691 + break
692 +
693 + # === Phase 1: Detect tool name (send when first arg_key is seen) =
694 + # ``name_new`` mirrors the old ``name_delta is not None``: it marks
695 + # the chunk that first emits this tool's name, which is the only
696 + # chunk the wrapper mints an id + type for.
697 + name_new = False
698 + pending_name: str | None = None
699 + if self._streaming_tool_name is None:
700 + arg_idx = cur_text.find(self.arg_key_start_token, start_idx)
701 + end_idx = cur_text.find(self.tool_call_end_token, start_idx)
702 + if arg_idx == -1 and end_idx == -1:
703 + # tool name not yet complete; keep buffering from
704 + # tool_call_start
705 + self._buffer = cur_text[start_idx:]
706 + break
707 +
708 + name_start = start_idx + len(self.tool_call_start_token)
709 + name_end = arg_idx if arg_idx != -1 else end_idx
710 +
711 + tool_name = cur_text[name_start:name_end].strip()
712 + self._streaming_tool_name = tool_name
713 +
714 + if arg_idx != -1:
715 + self._buffer = cur_text[arg_idx:]
716 + else:
717 + self._buffer = cur_text[end_idx:]
718 +
719 + # Increment tool_id and mark that a name chunk should be emitted
720 + self.current_tool_id += 1
721 + name_new = True
722 + pending_name = tool_name
723 +
724 + # Check if buffer already has complete arguments
725 + # (all-in-one-delta); otherwise only the name is ready now.
726 + if self.tool_call_end_token not in self._buffer:
727 + tool_calls.append(
728 + StreamToolCall(
729 + index=self.current_tool_id,
730 + new=True,
731 + name=tool_name,
732 + arguments=None,
733 + )
734 + )
735 + break
736 + # Buffer already has a complete tool call; fall through to
737 + # phase 2 below.
738 +
739 + # === Phase 2: Incremental argument streaming ===
740 + result = self._extract_streaming_incremental(name_new, pending_name, tools)
741 + if result:
742 + tool_calls.extend(result["tool_calls"])
743 +
744 + # Continue only after the current tool call has completed and the
745 + # remaining structural payload starts with another tool call.
746 + # A literal ``<tool_call>`` may legitimately appear inside an open
747 + # string argument; in that case the current tool name is still set
748 + # and the buffer has not been consumed, so continuing would spin on
749 + # the same buffer forever.
750 + if self._streaming_tool_name is None and self._buffer.startswith(
751 + self.tool_call_start_token
752 + ):
753 + continue
754 + break
755 +
756 + if content_delta is not None or tool_calls:
757 + return StreamDelta(content=content_delta, tool_calls=tool_calls)
758 + return None
759 +
760 + def _make_args_delta(self, argument_diff: str) -> StreamDelta:
761 + """Build an args-only streaming delta (no id/type -> new=False)."""
762 + return StreamDelta(
763 + content=None,
764 + tool_calls=[
765 + StreamToolCall(
766 + index=self.current_tool_id,
767 + new=False,
768 + name=None,
769 + arguments=argument_diff,
770 + )
771 + ],
772 + )
773 +
774 + def _extract_streaming_incremental(
775 + self,
776 + name_new: bool,
777 + pending_name: str | None,
778 + tools: list[ToolSchema] | None,
779 + ) -> StreamDelta | None:
780 + """Incremental phase-2: scan tags in buffer, emit JSON diffs.
781 +
782 + Strategy:
783 + - Track completed args and emit each one as a JSON fragment.
784 + - For string-typed args, stream the value character-by-character.
785 + - Withhold the closing ``}`` until ``</tool_call>`` is seen.
786 +
787 + We build JSON manually via fragments rather than using json.dumps
788 + with a cursor, because json.dumps of partial-vs-full string values
789 + produces incompatible prefixes (e.g. ``""}`` vs ``"Hello"}``).
790 + """
791 + buf = self._buffer
792 + is_complete = self.tool_call_end_token in buf
793 +
794 + if is_complete:
795 + end_idx = buf.find(self.tool_call_end_token)
796 + args_text = buf[:end_idx]
797 + remaining = buf[end_idx + len(self.tool_call_end_token) :]
798 + else:
799 + args_text = buf
800 + remaining = ""
801 +
802 + # --- scan all fully closed kv pairs ---
803 + arg_pairs = self.func_args_regex.findall(args_text)
804 + for key, value in arg_pairs:
805 + key = key.strip()
806 + if key not in self._completed_args:
807 + self._completed_args[key] = _parse_value(
808 + value, self._streaming_tool_name or "", key, tools
809 + )
810 +
811 + # --- detect partial (unclosed) kv at the tail ---
812 + last_closed_end = 0
813 + for m in self.func_args_regex.finditer(args_text):
814 + last_closed_end = m.end()
815 + tail = args_text[last_closed_end:]
816 +
817 + partial_key: str | None = None
818 + partial_value: str | None = None
819 +
820 + ak_start = tail.find(self.arg_key_start_token)
821 + if ak_start != -1:
822 + ak_end = tail.find(
823 + self.arg_key_end_token,
824 + ak_start + len(self.arg_key_start_token),
825 + )
826 + if ak_end != -1:
827 + partial_key = tail[
828 + ak_start + len(self.arg_key_start_token) : ak_end
829 + ].strip()
830 + self._current_arg_key = partial_key
831 + self._current_arg_is_string = _is_only_string_type(
832 + self._streaming_tool_name or "",
833 + partial_key,
834 + tools,
835 + )
836 +
837 + av_start = tail.find(self.arg_value_start_token, ak_end)
838 + if av_start != -1:
839 + val_content_start = av_start + len(self.arg_value_start_token)
840 + if self._current_arg_is_string:
841 + raw_value = tail[val_content_start:]
842 + # Guided decoding can split ``</arg_value>`` across
843 + # deltas; hold back a partial closing tag so it is not
844 + # streamed as part of the value.
845 + keep = len(raw_value) - partial_tag_overlap(
846 + raw_value, self.arg_value_end_token
847 + )
848 + partial_value = raw_value[:keep]
849 + else:
850 + # key not yet closed
851 + self._current_arg_key = None
852 + self._current_arg_is_string = False
853 +
854 + # --- build the current JSON snapshot as a string ---
855 + # We construct JSON manually so we can precisely control
856 + # what gets sent incrementally.
857 + snapshot_parts: list[str] = []
858 + for k, v in self._completed_args.items():
859 + k_json = json.dumps(k, ensure_ascii=False)
860 + v_json = json.dumps(v, ensure_ascii=False)
861 + snapshot_parts.append(f"{k_json}: {v_json}")
862 +
863 + if partial_key is not None and partial_value is not None:
864 + k_json = json.dumps(partial_key, ensure_ascii=False)
865 + # For string partial value, we build the JSON string
866 + # WITHOUT the closing quote, so the prefix stays stable
867 + # as the value grows. The closing `"` and `}` will be
868 + # sent when the value or tool_call closes.
869 + escaped_val = json.dumps(partial_value, ensure_ascii=False)[1:-1]
870 + # Note: no closing " here - it's appended only on close
871 + snapshot_parts.append(f'{k_json}: "{escaped_val}')
872 +
873 + snapshot = "{" + ", ".join(snapshot_parts) + "}"
874 +
875 + # --- compute diff ---
876 + argument_diff: str | None = None
877 +
878 + if is_complete:
879 + # Tool call finished - send everything remaining.
880 + # Build final snapshot with proper JSON (all values closed).
881 + final_args = dict(self._completed_args)
882 + final_json = json.dumps(final_args, ensure_ascii=False)
883 + if self._streamed_json_len < len(final_json):
884 + argument_diff = final_json[self._streamed_json_len :]
885 + self._streamed_json_len = len(final_json)
886 +
887 + # Record into prev_tool_call_arr
888 + self.prev_tool_call_arr.append(
889 + {
890 + "name": self._streaming_tool_name,
891 + "arguments": final_args,
892 + }
893 + )
894 + self.streamed_args_for_tool.append(final_json)
895 +
896 + self._reset_streaming_tool_state()
897 + self._buffer = remaining
898 + else:
899 + # Still in progress - withhold the tail.
900 + # For open strings: snapshot ends with ...partial_val}
901 + # we withhold "}" (1 char) - the missing closing " will
902 + # be sent when the value closes.
903 + # For no open string: snapshot ends with ...value"}
904 + # we withhold "}" (1 char).
905 + end = len(snapshot) - 1 # exclude trailing "}"
906 + if end > self._streamed_json_len:
907 + argument_diff = snapshot[self._streamed_json_len : end]
908 + self._streamed_json_len = end
909 +
910 + # --- construct return dict ---
911 + if name_new and argument_diff:
912 + return StreamDelta(
913 + content=None,
914 + tool_calls=[
915 + StreamToolCall(
916 + index=self.current_tool_id,
917 + new=True,
918 + name=pending_name,
919 + arguments=argument_diff,
920 + )
921 + ],
922 + )
923 + elif name_new:
924 + return StreamDelta(
925 + content=None,
926 + tool_calls=[
927 + StreamToolCall(
928 + index=self.current_tool_id,
929 + new=True,
930 + name=pending_name,
931 + arguments=None,
932 + )
933 + ],
934 + )
935 + elif argument_diff:
936 + return self._make_args_delta(argument_diff)
937 + else:
938 + return None
939 +
940 +
941 +def build_tool_extractor(
942 + tokenizer: TokenizerLike, *, strict: bool
943 +) -> HYV4ToolExtractor:
944 + return HYV4ToolExtractor(
945 + tokenizer.get_vocab(), detect_token_suffix(tokenizer), strict
946 + )
947 +
948 +
949 +def _tools_to_plain(tools) -> list[ToolSchema] | None:
950 + """Convert tool-schema objects into plain dicts for the extractor."""
951 + if not tools:
952 + return None
953 + return [
954 + ToolSchema(name=t.function.name, parameters=t.function.parameters)
955 + for t in tools
956 + ]
957 +
958 +
959 +class HYV4ToolParser(ToolParser):
960 + """vLLM adapter around :class:`HYV4ToolExtractor`.
961 +
962 + Converts the extractor's plain dicts to vLLM protocol objects, mints
963 + tool-call ids, and mirrors ``prev_tool_call_arr`` /
964 + ``streamed_args_for_tool`` so the serving layer sees live state.
965 + """
966 +
967 + # HYV4 emits its own <tool_calls>/<tool_call>/<arg_key>/<arg_value>
968 + # structural tokens, not a plain JSON array. The generic required/named
969 + # JSON streaming path (streaming.extract_required_tool_call_streaming)
970 + # cannot parse this format and produces misaligned multi-tool deltas in
971 + # streaming mode. Fall back to our own (auto-path) parser for
972 + # required/named too, so streaming and non-streaming stay consistent.
973 + supports_required_and_named: bool = False
974 + structural_tag_model: str = "hy_v4"
975 +
976 + def __init__(self, tokenizer: TokenizerLike, tools=None):
977 + super().__init__(tokenizer, tools)
978 +
979 + # Tools are supplied at construction (``tool_parser(tokenizer,
980 + # request.tools)``); keep a plain-dict copy for the extractor.
981 + self._plain_tools = _tools_to_plain(tools)
982 + # Malformed tag structure is rejected and returned as raw content
983 + # instead of a half-parsed tool call.
984 + self._extractor = build_tool_extractor(tokenizer, strict=True)
985 + self._mirror_state()
986 +
987 + def get_structural_tag(
988 + self,
989 + request: ChatCompletionRequest | ResponsesRequest,
990 + *,
991 + reasoning: bool = False,
992 + ):
993 + """Build a structural tag matching HYV4's tool tokens.
994 +
995 + Overridden only to pass the per-checkpoint ``token_suffix`` through to
996 + the builder; everything else follows the base implementation.
997 +
998 + Named (forced) tool choice is constrained with the same structural tag:
999 + because ``supports_required_and_named`` is False, required *and* named
1000 + requests are parsed by the HYV4 extractor, which only recognizes the HYV4
1001 + structural tokens -- the default JSON guided-decoding path would make
1002 + the model emit JSON the extractor silently drops.
1003 +
1004 + Args:
1005 + request: The request being adjusted.
1006 + reasoning: Whether the grammar also covers the reasoning phase.
1007 +
1008 + Returns:
1009 + The structural tag, or None when structural tagging does not apply.
1010 + """
1011 + if not envs.VLLM_ENFORCE_STRICT_TOOL_CALLING:
1012 + return None
1013 +
1014 + # Only constrain required / forced tool choice with the structural tag.
1015 + # For "auto", tool calling is optional and most requests never emit a
1016 + # tool call, so applying the grammar to every auto request would make
1017 + # the whole stream pay the output inflation and the per-token
1018 + # bitmask/accept overhead for a minority benefit. Skip it here so auto
1019 + # falls through to free generation; the HYV4 parser still reads back
1020 + # any tool calls the model emits.
1021 + tool_choice = request.tool_choice
1022 + if (
1023 + tool_choice is None
1024 + or tool_choice == "auto"
1025 + or getattr(tool_choice, "mode", None) == "auto"
1026 + ):
1027 + return None
1028 +
1029 + from vllm.tool_parsers.structural_tag_registry import get_model_structural_tag
1030 +
1031 + try:
1032 + return get_model_structural_tag(
1033 + model=self.structural_tag_model,
1034 + tools=request.tools,
1035 + tool_choice=request.tool_choice,
1036 + reasoning=reasoning,
1037 + token_suffix=self._extractor.token_suffix,
1038 + )
1039 + except Exception:
1040 + logger.warning(
1041 + "HYV4ToolParser failed to build a structural tag; falling back "
1042 + "to the default guided-decoding path.",
1043 + exc_info=True,
1044 + )
1045 + return None
1046 +
1047 + def _mirror_state(self) -> None:
1048 + # Alias the SAME list objects the extractor appends to, so the serving
1049 + # layer reads live state. ``current_tool_id`` is an int, so re-read it.
1050 + self.prev_tool_call_arr = self._extractor.prev_tool_call_arr
1051 + self.streamed_args_for_tool = self._extractor.streamed_args_for_tool
1052 + self.current_tool_id = self._extractor.current_tool_id
1053 +
1054 + @staticmethod
1055 + def _is_guided(request: ChatCompletionRequest) -> bool:
1056 + """Whether a structural tag constrains this request's decoding.
1057 +
1058 + ``get_structural_tag`` only applies to required / named tool choice and
1059 + the serving layer records the result on the request. A structural tag
1060 + constrains strings rather than token ids, so the ``<tool_calls>`` marker
1061 + can arrive split across sub-word tokens and must be detected on the
1062 + string level.
1063 +
1064 + Args:
1065 + request: The request being streamed.
1066 +
1067 + Returns:
1068 + True when the streaming parser must use the string-marker path.
1069 + """
1070 + structured_outputs = getattr(request, "structured_outputs", None)
1071 + return (
1072 + structured_outputs is not None
1073 + and structured_outputs.structural_tag is not None
1074 + )
1075 +
1076 + def _tools(self, request: ChatCompletionRequest) -> list[ToolSchema] | None:
1077 + # Prefer the tools passed at construction; fall back to the request.
1078 + if self._plain_tools is not None:
1079 + return self._plain_tools
1080 + return _tools_to_plain(getattr(request, "tools", None))
1081 +
1082 + def extract_tool_calls(
1083 + self,
1084 + model_output: str,
1085 + request: ChatCompletionRequest,
1086 + ) -> ExtractedToolCallInformation:
1087 + result = self._extractor.extract_tool_calls(model_output, self._tools(request))
1088 + self._mirror_state()
1089 + return ExtractedToolCallInformation(
1090 + tools_called=result["tools_called"],
1091 + content=result["content"],
1092 + tool_calls=[
1093 + ToolCall(
1094 + type="function",
1095 + function=FunctionCall(
1096 + name=tc["name"],
1097 + arguments=tc["arguments"],
1098 + ),
1099 + )
1100 + for tc in result["tool_calls"]
1101 + ],
1102 + )
1103 +
1104 + def extract_tool_calls_streaming(
1105 + self,
1106 + previous_text: str,
1107 + current_text: str,
1108 + delta_text: str,
1109 + previous_token_ids,
1110 + current_token_ids,
1111 + delta_token_ids,
1112 + request: ChatCompletionRequest,
1113 + ) -> DeltaMessage | None:
1114 + delta = self._extractor.extract_tool_calls_streaming(
1115 + previous_text,
1116 + current_text,
1117 + delta_text,
1118 + previous_token_ids,
1119 + current_token_ids,
1120 + delta_token_ids,
1121 + self._tools(request),
1122 + guided=self._is_guided(request),
1123 + )
1124 + self._mirror_state()
1125 + if delta is None:
1126 + return None
1127 +
1128 + tool_calls: list[DeltaToolCall] = []
1129 + for tc in delta.get("tool_calls", []):
1130 + # Only set the fields that are present, leaving the rest UNSET (not
1131 + # None). Streaming serialization uses exclude_unset, so an explicit
1132 + # ``arguments=None`` would emit ``"arguments": null`` and break
1133 + # clients that concatenate argument deltas. Matches hy_v3.
1134 + fn_kwargs: dict[str, str] = {}
1135 + name = tc.get("name")
1136 + if name is not None:
1137 + fn_kwargs["name"] = name
1138 + arguments = tc.get("arguments")
1139 + if arguments is not None:
1140 + fn_kwargs["arguments"] = arguments
1141 + function = DeltaFunctionCall(**fn_kwargs)
1142 + if tc.get("new"):
1143 + # First chunk for this tool: mint an id + set type here only.
1144 + tool_calls.append(
1145 + DeltaToolCall(
1146 + index=tc["index"],
1147 + id=make_tool_call_id(),
1148 + type="function",
1149 + function=function,
1150 + )
1151 + )
1152 + else:
1153 + tool_calls.append(DeltaToolCall(index=tc["index"], function=function))
1154 +
1155 + # A single engine delta can carry content plus N drained tool calls.
1156 + # Only set the fields that are present so streaming serialization
1157 + # (exclude_unset) never emits ``content: null`` on a tool-call chunk.
1158 + content = delta.get("content")
1159 + kwargs: dict[str, Any] = {}
1160 + if content is not None:
1161 + kwargs["content"] = content
1162 + if tool_calls:
1163 + kwargs["tool_calls"] = tool_calls
1164 + if not kwargs:
1165 + return None
1166 + return DeltaMessage(**kwargs)
vllm/tool_parsers/structural_tag_registry.py
+158 -5
@@ -49,6 +49,7 @@ StructuralTagBuilder: TypeAlias = Callable[
49 list[BuiltinToolParam],
50 SimplifiedToolChoice,
51 bool,
52 + str,
53 ],
54 StructuralTag,
55 ]
@@ -70,7 +71,7 @@ XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset(
71 "deepseek_v4",
72 }
73 )
73 -VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "kimi_k3"})
74 +VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "hy_v4", "kimi_k3"})
75 SUPPORTED_STRUCTURAL_TAG_MODELS = (
76 XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
77 )
@@ -78,7 +79,9 @@ SUPPORTED_STRUCTURAL_TAG_MODELS = (
79 _VLLM_STRUCTURAL_TAG_REGISTRY: dict[str, StructuralTagBuilder] = {}
80
81
81 -def register_vllm_structural_tag(model: str):
82 +def register_vllm_structural_tag(
83 + model: str,
84 +) -> Callable[[StructuralTagBuilder], StructuralTagBuilder]:
85 """Register a vLLM-owned structural tag builder."""
86
87 def decorator(func: StructuralTagBuilder) -> StructuralTagBuilder:
@@ -104,6 +107,7 @@ def get_model_structural_tag(
107 tools: Sequence[ChatCompletionToolsParam | ResponsesTool] | None,
108 tool_choice: ToolChoice,
109 reasoning: bool,
110 + token_suffix: str = "",
111 ) -> StructuralTag | None:
112 """Build a structural tag with xgrammar's builtin model templates."""
113
@@ -126,12 +130,19 @@ def get_model_structural_tag(
130 builtin_tools,
131 simplified_tool_choice,
132 reasoning,
133 + token_suffix,
134 )
135
136 if model not in XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS:
137 supported = sorted(SUPPORTED_STRUCTURAL_TAG_MODELS)
138 raise ValueError(f"Unknown format type: {model}, supported types: {supported}")
139
140 + if token_suffix:
141 + raise ValueError(
142 + f"Structural tag model {model!r} is an xgrammar builtin with fixed "
143 + f"tokens and cannot apply token_suffix={token_suffix!r}"
144 + )
145 +
146 return get_xgrammar_model_structural_tag(
147 model=model,
148 tools=dumped_tools,
@@ -244,8 +255,9 @@ def get_hermes_structural_tag(
255 builtin_tools: list[BuiltinToolParam],
256 tool_choice: SimplifiedToolChoice,
257 reasoning: bool,
258 + token_suffix: str = "",
259 ) -> StructuralTag:
248 - del builtin_tools, reasoning
260 + del builtin_tools, reasoning, token_suffix
261
262 tool_call_trigger = "<tool_call>"
263
@@ -293,8 +305,9 @@ def get_minimax_structural_tag(
305 builtin_tools: list[BuiltinToolParam],
306 tool_choice: SimplifiedToolChoice,
307 reasoning: bool,
308 + token_suffix: str = "",
309 ) -> StructuralTag:
297 - del builtin_tools, reasoning
310 + del builtin_tools, reasoning, token_suffix
311
312 tool_call_begin = "<minimax:tool_call>\n"
313 tool_call_end = "</minimax:tool_call>"
@@ -589,8 +602,9 @@ def get_kimi_k3_structural_tag(
602 builtin_tools: list[BuiltinToolParam],
603 tool_choice: SimplifiedToolChoice,
604 reasoning: bool,
605 + token_suffix: str = "",
606 ) -> StructuralTag:
593 - del builtin_tools, reasoning
607 + del builtin_tools, reasoning, token_suffix
608
609 trailer = OptionalFormat(content=ConstStringFormat(value=_K3_MESSAGE_CLOSE))
610
@@ -611,3 +625,142 @@ def get_kimi_k3_structural_tag(
625 return StructuralTag(
626 format=SequenceFormat(elements=[*_k3_response_prefix(), tools_part, trailer])
627 )
628 +
629 +
630 +# ---------------------------------------------------------------------------
631 +# HYV4 (<tool_calls>/<tool_call>/<arg_key>/<arg_value> structural tokens)
632 +# ---------------------------------------------------------------------------
633 +# HYV4 assistant output after the reasoning gate (``</think:SUF>``):
634 +# <tool_calls:SUF>
635 +# <tool_call:SUF>NAME
636 +# <arg_key:SUF>K</arg_key:SUF><arg_value:SUF>V</arg_value:SUF>
637 +# </tool_call:SUF>
638 +# </tool_calls:SUF>
639 +# Adjacent <tool_call> blocks have no separator, and argument values are emitted
640 +# verbatim (no surrounding quotes), so the value body cannot be described by
641 +# JSONSchemaFormat. Only the skeleton (tool names, argument keys, tag order) is
642 +# constrained; each <arg_value> body stays free text.
643 +
644 +
645 +def _hy_v4_argument_keys(parameters: dict[str, Any]) -> tuple[list[str], list[str]]:
646 + """Split a tool's argument keys into ``(required, optional)``.
647 +
648 + Both lists keep declaration order. Required keys missing from
649 + ``properties`` are still treated as required.
650 + """
651 + properties = parameters.get("properties") or {}
652 + required = parameters.get("required") or []
653 + required_set = set(required)
654 + required_keys = [k for k in properties if k in required_set]
655 + required_keys += [k for k in required if k not in properties]
656 + optional_keys = [k for k in properties if k not in required_set]
657 + return required_keys, optional_keys
658 +
659 +
660 +@register_vllm_structural_tag("hy_v4")
661 +def get_hy_v4_structural_tag(
662 + tools: list[FunctionToolParam],
663 + builtin_tools: list[BuiltinToolParam],
664 + tool_choice: SimplifiedToolChoice,
665 + reasoning: bool,
666 + token_suffix: str = "",
667 +) -> StructuralTag:
668 + """Build HYV4 <tool_calls>/<tool_call>/<arg_key>/<arg_value> structural tags.
669 +
670 + Args:
671 + tools: Normalized function tools the model may call.
672 + builtin_tools: Unused; HYV4 has no builtin tools.
673 + tool_choice: Simplified tool choice.
674 + reasoning: Whether the grammar also covers the reasoning phase.
675 + token_suffix: Per-checkpoint structural-token suffix including the
676 + leading colon (e.g. ``":6124c78e"``), or ``""`` when the checkpoint
677 + uses unsuffixed tokens. The HYV4 tool parser reads it off the
678 + tokenizer vocab and passes it to ``get_model_structural_tag``.
679 + """
680 + del builtin_tools
681 +
682 + think_end = f"</think{token_suffix}>"
683 + tool_calls_begin = f"<tool_calls{token_suffix}>"
684 + tool_calls_end = f"</tool_calls{token_suffix}>"
685 + tool_call_begin = f"<tool_call{token_suffix}>"
686 + tool_call_end = f"</tool_call{token_suffix}>"
687 + arg_key_begin = f"<arg_key{token_suffix}>"
688 + arg_key_end = f"</arg_key{token_suffix}>"
689 + arg_value_begin = f"<arg_value{token_suffix}>"
690 + arg_value_end = f"</arg_value{token_suffix}>"
691 +
692 + def arg_pair(key: str) -> SequenceFormat:
693 + return SequenceFormat(
694 + elements=[
695 + ConstStringFormat(value=arg_key_begin),
696 + ConstStringFormat(value=key),
697 + ConstStringFormat(value=arg_key_end),
698 + ConstStringFormat(value=arg_value_begin),
699 + AnyTextFormat(),
700 + ConstStringFormat(value=arg_value_end),
701 + ]
702 + )
703 +
704 + def single_tool_call(tool: FunctionToolParam) -> TagFormat:
705 + function = tool.function
706 + parameters = (
707 + function.parameters if isinstance(function.parameters, dict) else {}
708 + )
709 + required_keys, optional_keys = _hy_v4_argument_keys(parameters)
710 + elements: list[Any] = [arg_pair(k) for k in required_keys]
711 + if optional_keys:
712 + # Optional keys may appear in any order, or not at all.
713 + optional_pairs = [arg_pair(k) for k in optional_keys]
714 + elements.append(
715 + StarFormat(
716 + content=optional_pairs[0]
717 + if len(optional_pairs) == 1
718 + else OrFormat(elements=optional_pairs)
719 + )
720 + )
721 + content: Any = (
722 + SequenceFormat(elements=elements) if elements else AnyTextFormat()
723 + )
724 + return TagFormat(
725 + begin=tool_call_begin + function.name,
726 + content=content,
727 + end=tool_call_end,
728 + )
729 +
730 + tags = [single_tool_call(tool) for tool in tools]
731 +
732 + if tool_choice == "auto":
733 + block = TagFormat(
734 + begin=tool_calls_begin,
735 + content=TagsWithSeparatorFormat(
736 + tags=tags,
737 + separator="",
738 + at_least_one=True,
739 + ),
740 + end=tool_calls_end,
741 + )
742 + suffix_tag: Any = TriggeredTagsFormat(triggers=[tool_calls_begin], tags=[block])
743 + else:
744 + # HYV4 places tool_call blocks back-to-back with no separator. "forced"
745 + # has already been filtered down to the single named tool and must emit
746 + # exactly one block.
747 + suffix_tag = SequenceFormat(
748 + elements=[
749 + ConstStringFormat(value=tool_calls_begin),
750 + TagsWithSeparatorFormat(
751 + tags=tags,
752 + separator="",
753 + at_least_one=True,
754 + stop_after_first=tool_choice == "forced",
755 + ),
756 + ConstStringFormat(value=tool_calls_end),
757 + ]
758 + )
759 +
760 + if not reasoning:
761 + return StructuralTag(format=suffix_tag)
762 +
763 + # The reasoning phase is constrained too, so the tag must explicitly skip
764 + # the ``<think>...</think:SUF>`` prefix.
765 + prefix_tag = TagFormat(begin="", content=AnyTextFormat(), end=think_end)
766 + return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag]))
vllm/transformers_utils/config.py
+1
@@ -91,6 +91,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
91 granite4_vision="Granite4VisionConfig",
92 hyperclovax="HyperCLOVAXConfig",
93 hy_v3="HYV3Config",
94 + hy_v4="HYV4Config",
95 isaac="IsaacConfig",
96 kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3
97 kimi_linear="KimiLinearConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -41,6 +41,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
41 "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat",
42 "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision",
43 "HYV3Config": "vllm.transformers_utils.configs.hy_v3",
44 + "HYV4Config": "vllm.transformers_utils.configs.hy_v4",
45 "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax",
46 "IsaacConfig": "vllm.transformers_utils.configs.isaac",
47 # RWConfig is for the original tiiuae/falcon-40b(-instruct) and
@@ -130,6 +131,7 @@ __all__ = [
131 "FunAudioChatAudioEncoderConfig",
132 "Granite4VisionConfig",
133 "HYV3Config",
134 + "HYV4Config",
135 "HyperCLOVAXConfig",
136 "IsaacConfig",
137 "RWConfig",
vllm/transformers_utils/configs/hy_v4.py new
+172
@@ -0,0 +1,172 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from huggingface_hub.dataclasses import strict
4 +from transformers.configuration_utils import PreTrainedConfig
5 +from transformers.modeling_rope_utils import RopeParameters
6 +
7 +
8 +@strict
9 +class HYV4Config(PreTrainedConfig):
10 + r"""Configuration class for the HY V4 (`HYV4ForCausalLM`) model.
11 +
12 + Kept field-for-field in sync with
13 + `transformers.models.hy_v4.configuration_hy_v4.HYV4Config` so that a
14 + checkpoint resolves to the same values through either class. Registering it
15 + in vLLM's `_CONFIG_REGISTRY` lets the architecture load without
16 + `trust_remote_code` and without an `auto_map` module next to the weights.
17 +
18 + HYV4 is a mixture-of-experts causal language model using Multi-head Latent
19 + Attention (MLA), DeepSeek-style sparse attention (DSA), gated MLA,
20 + learnable attention sinks, and independent Hyper-Connections (iHC).
21 +
22 + The `layer_types`, `num_nextn_predict_layers` and `mtp_loss_factor` fields
23 + below are vLLM-only additions: upstream leaves them to `**kwargs`, but
24 + vLLM reads them directly (sparse-attention selection and MTP speculative
25 + decoding), so they are declared to give checkpoints that omit them a
26 + working default instead of a `None`.
27 +
28 + Args:
29 + mlp_layer_types: Per-layer MLP kind, `"dense"` or `"sparse"`. Defaults
30 + to one dense layer followed by sparse MoE layers.
31 + layer_types: Per-layer attention kind, e.g. `"full_attention"` or
32 + `"deepseek_sparse_attention"`.
33 + indexer_types: Per-layer DSA indexer kind, `"full"` or `"shared"`. A
34 + shared layer reuses the most recent full indexer in the same
35 + forward request.
36 + index_topk: Maximum number of key positions selected by each DSA query.
37 + index_head_dim: Hidden dimension of each DSA indexer head.
38 + index_n_heads: Number of DSA indexer heads.
39 + enable_lm_head_fp32: Whether the language-model head emits float32
40 + logits. Surfaced as ``head_dtype`` so `LogitsProcessor` accumulates
41 + the projection into fp32 while the weight stays in the model dtype.
42 + enable_ihc: Whether independent Hyper-Connections are enabled.
43 + hc_mult: Number of hidden-state channels maintained by iHC.
44 + hc_magnitude: Scale applied to the iHC post-gating branch.
45 + hc_eps: Numerical epsilon added to iHC sigmoid gates.
46 + gated_mla: Whether to gate the MLA output.
47 + gating_type: MLA gate granularity, `"elementwise"` or `"headwise"`.
48 + learnable_sink: Whether to add a learned per-head attention sink.
49 + learnable_sink_init: Initial value of each learned attention-sink
50 + logit.
51 + swiglu_limit: Magnitude of the routed-expert SwiGLU clamp. Values at or
52 + below zero disable the clamp.
53 + num_nextn_predict_layers: Number of MTP layers in the checkpoint.
54 + mtp_loss_factor: Training-time MTP loss weight (unused at inference).
55 + """
56 +
57 + model_type = "hy_v4"
58 + keys_to_ignore_at_inference = ["past_key_values"]
59 + attribute_map = {
60 + "num_local_experts": "n_routed_experts",
61 + }
62 + base_model_tp_plan = {
63 + "layers.*.self_attn.q_b_proj": "colwise",
64 + "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
65 + "layers.*.self_attn.kv_b_proj": "colwise",
66 + "layers.*.self_attn.o_proj": "rowwise",
67 + "layers.*.self_attn.linear_gate": "colwise",
68 + "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
69 + "layers.*.mlp.experts.down_proj": "rowwise",
70 + "layers.*.mlp.experts": "moe_tp_experts",
71 + "layers.*.mlp.shared_experts.gate_proj": "colwise",
72 + "layers.*.mlp.shared_experts.up_proj": "colwise",
73 + "layers.*.mlp.shared_experts.down_proj": "rowwise",
74 + "layers.*.mlp.gate_proj": "colwise",
75 + "layers.*.mlp.up_proj": "colwise",
76 + "layers.*.mlp.down_proj": "rowwise",
77 + }
78 + base_model_pp_plan = {
79 + "embed_tokens": (["input_ids"], ["inputs_embeds"]),
80 + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
81 + "norm": (["hidden_states"], ["hidden_states"]),
82 + }
83 +
84 + vocab_size: int = 120832
85 + hidden_size: int = 2816
86 + intermediate_size: int = 6912
87 + moe_intermediate_size: int = 768
88 + num_hidden_layers: int = 34
89 + num_attention_heads: int = 32
90 + num_key_value_heads: int = 32
91 + head_dim: int = 256
92 + hidden_act: str = "silu"
93 + max_position_embeddings: int = 262144
94 + initializer_range: float = 0.006
95 + rms_norm_eps: float = 1e-5
96 + use_cache: bool = True
97 + pad_token_id: int | None = 120002
98 + bos_token_id: int | None = 120000
99 + eos_token_id: int | list[int] | None = 120025
100 + tie_word_embeddings: bool = False
101 + attention_bias: bool = False
102 + attention_dropout: float = 0.0
103 + n_routed_experts: int = 256
104 + n_shared_experts: int = 1
105 + num_experts_per_tok: int = 8
106 + routed_scaling_factor: float = 2.827
107 + norm_topk_prob: bool = True
108 + q_lora_rank: int = 1536
109 + kv_lora_rank: int = 512
110 + qk_nope_head_dim: int = 192
111 + qk_rope_head_dim: int = 64
112 + v_head_dim: int = 256
113 + mlp_layer_types: list[str] | None = None
114 + layer_types: list[str] | None = None
115 + index_topk: int = 2048
116 + index_head_dim: int = 128
117 + index_n_heads: int = 16
118 + indexer_types: list[str] | None = None
119 + enable_lm_head_fp32: bool = True
120 + enable_ihc: bool = True
121 + hc_mult: int = 4
122 + hc_magnitude: float = 2.0
123 + hc_eps: float = 1e-6
124 + gated_mla: bool = True
125 + gating_type: str = "elementwise"
126 + learnable_sink: bool = True
127 + learnable_sink_init: float = 0.0
128 + swiglu_limit: float = 10.0
129 + rope_parameters: RopeParameters | dict | None = None
130 + num_nextn_predict_layers: int = 1
131 + mtp_loss_factor: float = 0.1
132 +
133 + def __post_init__(self, **kwargs):
134 + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
135 + # RoPE applies only to the rope slice, so point `head_dim` at it: the
136 + # inherited rotary embedding reads `config.head_dim` and then computes
137 + # the right frequencies with no override.
138 + self.head_dim = self.qk_rope_head_dim
139 +
140 + if self.mlp_layer_types is None:
141 + self.mlp_layer_types = ["dense"] * min(1, self.num_hidden_layers) + [
142 + "sparse"
143 + ] * max(self.num_hidden_layers - 1, 0)
144 + if self.indexer_types is None:
145 + self.indexer_types = [
146 + "full" if layer_idx == 0 or (layer_idx - 1) % 4 == 0 else "shared"
147 + for layer_idx in range(self.num_hidden_layers)
148 + ]
149 + # `PreTrainedConfig` validates `layer_types` against a fixed allow-list
150 + # in which "deepseek_sparse_attention" is the only accepted spelling of
151 + # a DSA layer. The model treats "sparse_attention" and "sparse" as
152 + # equivalent, so fold them in before the base class validates.
153 + if self.layer_types is not None:
154 + self.layer_types = [
155 + "deepseek_sparse_attention"
156 + if lt in ("sparse_attention", "sparse")
157 + else lt
158 + for lt in self.layer_types
159 + ]
160 +
161 + # `ModelConfig.head_dtype` reads this off the HF config, so an fp32 head
162 + # needs no model-side dtype juggling: LogitsProcessor picks the
163 + # `torch.mm(out_dtype=float32)` path and leaves the weight in the model
164 + # dtype. Only fill it in when the checkpoint (or --hf-overrides) has not
165 + # already pinned a value.
166 + if self.enable_lm_head_fp32 and getattr(self, "head_dtype", None) is None:
167 + self.head_dtype = "float32"
168 +
169 + super().__post_init__(**kwargs)
170 +
171 +
172 +__all__ = ["HYV4Config"]
vllm/transformers_utils/model_arch_config_convertor.py
+2
@@ -332,6 +332,8 @@ class ModelArchConfigConvertorBase:
332 "bailing_hybrid",
333 "bailing_hybrid_mtp",
334 "bailing_hybrid_v3_mtp",
335 + "hy_v4",
336 + "hy_v4_mtp",
337 ):
338 # check is deepseek_v4 model
339 if hasattr(self.hf_text_config, "compress_ratios"):
vllm/v1/worker/gpu/spec_decode/speculator.py
+23 -5
@@ -30,6 +30,19 @@ from vllm.v1.worker.utils import AttentionGroup
30 logger = init_logger(__name__)
31
32
33 +def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
34 + """Whether the target replaces the drafter's input with its HC residual.
35 +
36 + Keyed on the same hook the model runner calls to perform that swap. It is
37 + resolved from the target model class because speculators are built before
38 + the target model is instantiated.
39 + """
40 + from vllm.model_executor.model_loader.utils import get_model_cls
41 +
42 + target_cls = get_model_cls(vllm_config.model_config)
43 + return hasattr(target_cls, "get_mtp_target_hidden_states")
44 +
45 +
46 class BaseSpeculator(ABC):
47 @abstractmethod
48 def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
@@ -90,11 +103,16 @@ class DraftModelSpeculator(BaseSpeculator):
103 # the draft model's hidden size can be different from the target model's
104 # hidden size (e.g., Llama 3.3 70B).
105 self.hidden_size = self.draft_model_config.get_hidden_size()
93 - # Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP
94 - # draft the target's pre-hc_head (T, hc_mult * hidden_size) residual).
95 - # Non-HC models default to hc_mult=1 and are unaffected.
96 - hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
97 - self.hidden_size = self.hidden_size * hc_mult
106 + # Widen for HC-multiplexed residuals: a target that implements
107 + # get_mtp_target_hidden_states() (e.g. DeepSeek V4) hands the drafter
108 + # its pre-hc_head (T, hc_mult * hidden_size) residual in place of the
109 + # collapsed hidden states, so the drafter's buffers must match. Key off
110 + # that hook rather than hc_mult alone -- HY V4 runs iHC in its backbone
111 + # (hc_mult=4) but its MTP head consumes the collapsed states, so
112 + # widening it feeds propose() a 4x-too-wide buffer.
113 + if _target_feeds_hc_residual(vllm_config):
114 + hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
115 + self.hidden_size = self.hidden_size * hc_mult
116 self.vocab_size = self.draft_model_config.get_vocab_size()
117 self.dtype = vllm_config.model_config.dtype
118 self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel