@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6adad08767

Add Muse Glimmer model support (#51655)

Signed-off-by: Tiezhen Wang <tiezhen@inferact.ai> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Beto de Paola <betodepaola@meta.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: zixi-qi <zixi@inferact.ai>

Tiezhen WANG committed Aug 14, 2026 at 12:28 UTC 6adad08767583f52eb4d2122111af0bf638ed5e6
21 files changed +4333 -16
examples/tool_chat_template_muse_glimmer.jinja new
+1
@@ -0,0 +1 @@
1 +{%- macro render_content(content) -%}{%- if content is string -%}{{- content -}}{%- elif content is not none -%}{%- for part in content -%}{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}{%- elif part['type'] == 'text' -%}{{- part['text'] -}}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endmacro -%}{%- macro render_atem(tc) -%}{%- set args = tc.function.arguments -%}{%- if args is not mapping -%}{{- raise_exception('MuseGlimmer ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.') -}}{%- endif -%}{{- '<atem:function_calls>\n<atem:invoke name="' + tc.function.name + '">\n' -}}{%- for k, v in args.items() -%}{{- '<atem:parameter name="' + k + '">' -}}{%- if v is boolean -%}{%- if v -%}true{%- else -%}false{%- endif -%}{%- elif v is none -%}null{%- elif v is mapping or (v is iterable and v is not string) -%}{{- v | tojson -}}{%- else -%}{{- v -}}{%- endif -%}{{- '</atem:parameter>\n' -}}{%- endfor -%}{{- '</atem:invoke>\n</atem:function_calls>' -}}{%- endmacro -%}{%- macro render_tool_defs(tools) -%}{{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}}{{- 'You can invoke a function by writing a "<atem:function_calls>" block like the following:\n' -}}{{- '<atem:function_calls>\n<atem:invoke name="$FUNCTION_NAME">\n<atem:parameter name="$PARAMETER_NAME">$PARAMETER_VALUE</atem:parameter>\n...\n</atem:invoke>\n</atem:function_calls>\n\n' -}}{{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}}{{- 'Here are the functions available in JSONSchema format:\n' -}}{{- '// Tool metadata\n' -}}{%- set nsns = namespace(seen=[]) -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in nsns.seen -%}{%- set nsns.seen = nsns.seen + [tns] -%}{%- endif -%}{%- endfor -%}{%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%}{%- for tns in nsns.seen -%}{{- '{"name": ' + (tns | tojson) + ', "description": ' + ((nd[tns] if tns in nd else '') | tojson) + '}\n' -}}{%- endfor -%}{{- '// Function schemas' -}}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{{- '\n{"name": ' + (fn.name | tojson) + ', "description": ' + (fn.description | tojson) + ', "parameters": ' + (fn.parameters | tojson) + '}' -}}{%- endfor -%}{{- '\n\nHere\'s an example of how to call a function in the tool set. To make parallel tool calls, emit each call as its own message (one tool call per message) in consecutive messages:\n' -}}{{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than `example_tool_name.example_function_name`)\n\n' -}}{{- 'to=example_tool_name.example_function_name\n\n' -}}{{- '<atem:function_calls>\n<atem:invoke name="example_tool_name.example_function_name">\n' -}}{{- '<atem:parameter name="example_parameter_1">value_1</atem:parameter>\n' -}}{{- '<atem:parameter name="example_parameter_2">This is the value for the second parameter\nthat can span\n"multiple" lines\n</atem:parameter>\n' -}}{{- '</atem:invoke>\n</atem:function_calls>' -}}{%- endmacro -%}{%- macro render_reasoning() -%}{%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else 'high' -%}{{- 'Reasoning strength: ' + rs + '.' -}}{%- endmacro -%}{%- macro render_system_meta(tools) -%}{%- set rns = namespace(recipients=['"self"'], nslist=[]) -%}{%- if tools -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in rns.nslist -%}{%- set rns.nslist = rns.nslist + [tns] -%}{%- endif -%}{%- endfor -%}{%- for tns in rns.nslist -%}{%- set rns.recipients = rns.recipients + ['"' + tns + '.*"'] -%}{%- endfor -%}{%- endif -%}{%- set rns.recipients = rns.recipients + ['"user"'] -%}{{- '# Valid recipients: ' + rns.recipients | join(', ') + '.' -}}{%- endmacro -%}{{- bos_token -}}{%- set ns = namespace(has_system=false) -%}{%- for m in messages -%}{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}{%- endfor -%}{%- if not ns.has_system and (add_generation_prompt or tools) -%}{{- '<|start|>system<|message|>You are a helpful AI assistant.' -}}{%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else '2026-01-04' -%}{{- '\nKnowledge cutoff: ' + kc + '.' -}}{%- if current_date is defined and current_date -%}{{- '\nCurrent date: ' + current_date + '.' -}}{%- elif strftime_now is defined -%}{{- '\nCurrent date: ' + strftime_now('%Y-%m-%d') + '.' -}}{%- endif -%}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- endif -%}{%- for message in messages -%}{%- set role = message['role'] -%}{%- set end_token = '<|eom|>' if (not loop.last and messages[loop.index0 + 1]['role'] == role) else '<|eot|>' -%}{%- if role == 'system' -%}{{- '<|start|>system<|message|>' -}}{{- render_content(message['content']) -}}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- elif role == 'user' -%}{{- '<|start|>user<|message|>' -}}{{- render_content(message['content']) -}}{{- '<|eot|>' -}}{%- elif role == 'tool' -%}{%- set tname = message.get('name') -%}{%- if not tname -%}{%- set tcid = message.get('tool_call_id') -%}{%- set rns = namespace(name=tcid if tcid else '') -%}{%- for m in messages -%}{%- if m.get('tool_calls') -%}{%- for tc in m['tool_calls'] -%}{%- if tcid is not none and tc.id is defined and tc.id == tcid -%}{%- set rns.name = tc.function.name -%}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}{%- set tname = rns.name -%}{%- endif -%}{{- '<|start|>tool ' + tname + '<|message|><tool_output name="' + tname + '">\n' -}}{{- render_content(message['content']) -}}{{- '\n</tool_output><|eot|>' -}}{%- elif role == 'assistant' -%}{%- if message.get('reasoning_content') -%}{{- '<|start|>assistant to=self<|message|>' + message['reasoning_content'] + '<|eom|>' -}}{%- endif -%}{%- if message.get('tool_calls') -%}{%- for tc in message['tool_calls'] -%}{{- '<|start|>assistant to=' + tc.function.name + '<|message|>' -}}{{- render_atem(tc) -}}{%- if loop.last -%}{{- end_token -}}{%- else -%}{{- '<|eom|>' -}}{%- endif -%}{%- endfor -%}{%- else -%}{%- set recipient = message.get('recipient') or 'user' -%}{%- set end_turn = message.get('end_turn') -%}{%- if end_turn is none -%}{%- set end_turn = not (recipient and recipient != 'user') -%}{%- endif -%}{{- '<|start|>assistant' -}}{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}{{- '<|message|>' -}}{{- render_content(message['content']) -}}{{- ('<|eot|>' if end_turn else '<|eom|>') -}}{%- endif -%}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}
tests/models/registry.py
+18
@@ -1192,6 +1192,12 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1192 # required by current PrefixLM implementation
1193 max_num_batched_tokens=31872,
1194 ),
1195 + "MuseGlimmerForConditionalGeneration": _HfExamplesInfo(
1196 + "meta-models/Muse-Glimmer-30B",
1197 + ),
1198 + "MuseGlimmerForCausalLM": _HfExamplesInfo(
1199 + "meta-models/Muse-Glimmer-30B",
1200 + ),
1201 "NVLM_D": _HfExamplesInfo("nvidia/NVLM-D-72B", trust_remote_code=True),
1202 "Llama_Nemotron_Nano_VL": _HfExamplesInfo(
1203 "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1",
@@ -1489,6 +1495,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1495 max_num_seqs=32,
1496 min_transformers_version="4.56.3", # Required for Qwen3Next
1497 ),
1498 + "MuseGlimmerAssistantModel": _HfExamplesInfo(
1499 + "meta-models/Muse-Glimmer-30B",
1500 + speculative_model="meta-models/Muse-Glimmer-30B-assistant",
1501 + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1502 + max_num_seqs=32,
1503 + ),
1504 + "DFlashMuseGlimmerAssistantModel": _HfExamplesInfo(
1505 + "meta-models/Muse-Glimmer-30B",
1506 + speculative_model="meta-models/Muse-Glimmer-30B-assistant",
1507 + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1508 + max_num_seqs=32,
1509 + ),
1510 # [DSpark]
1511 "DSparkDraftModel": _HfExamplesInfo(
1512 "deepseek-ai/DeepSeek-V4-Pro-DSpark",
tests/models/utils.py
+7 -2
@@ -564,15 +564,20 @@ def dummy_hf_overrides(
564 )
565
566 if hasattr(hf_config, "vision_config"):
567 - hf_config.vision_config.update(
567 + vision_config = hf_config.vision_config
568 + vision_config.update(
569 {
570 "num_layers": 1,
571 "num_hidden_layers": 1,
572 }
573 )
574
575 + # Keep per-layer metadata consistent with the reduced layer count.
576 + if layer_types := getattr(vision_config, "layer_types", None):
577 + vision_config.update({"layer_types": layer_types[:1]})
578 +
579 if model_arch in ("Moondream3ForCausalLM", "HfMoondream"):
575 - hf_config.vision_config.update({"enc_n_layers": 1})
580 + vision_config.update({"enc_n_layers": 1})
581
582 # e.g.: ibm-granite/granite-speech-3.3-2b
583 if hasattr(hf_config, "encoder_config"):
tests/tool_use/test_muse_glimmer.py new
+386
@@ -0,0 +1,386 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Unit tests for the MuseGlimmer ATEM tool parser and reasoning parser.
5 +
6 +MuseGlimmer writes every turn as a sequence of channel-scoped messages rather
7 +than JSON, so the two parsers are tested together: the reasoning parser strips
8 +the reasoning span and forwards the remaining channels as content, and the tool
9 +parser reads ATEM markup out of those channels.
10 +
11 +Four areas, in order:
12 +
13 + 1. non-streaming tool-call extraction, including channel scoping (an
14 + ``<atem:invoke>`` echoed inside reasoning must never become a call);
15 + 2. the reasoning -> tool-parser handoff, which regressed once by returning
16 + ``content=None`` and starving the tool parser;
17 + 3. streaming, where markers routinely straddle chunk boundaries, plus
18 + truncation isolation for an unterminated ``to=self`` block;
19 + 4. tool-name normalization against the tools registered on the request.
20 +
21 +These drive the parsers directly and need no checkpoint. The tests that require
22 +a real tokenizer live in ``test_muse_glimmer_parse_delta.py``.
23 +"""
24 +
25 +import json
26 +from types import SimpleNamespace
27 +
28 +import pytest
29 +
30 +from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser
31 +from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser
32 +
33 +R: MuseGlimmerReasoningParser
34 +T: MuseGlimmerToolParser
35 +
36 +
37 +@pytest.fixture(autouse=True)
38 +def _fresh_parsers():
39 + """Give each test request-scoped parser state through the real constructors."""
40 + global R, T
41 + R = MuseGlimmerReasoningParser(object())
42 + T = MuseGlimmerToolParser(object())
43 +
44 +
45 +# Any framing token that must NEVER appear in surfaced reasoning/content.
46 +_FRAMING = [
47 + "<|start|>",
48 + "<|message|>",
49 + "<|eom|>",
50 + "<|eot|>",
51 + "to=self",
52 + "to=user",
53 + "to=read.read",
54 + "<atem:",
55 +]
56 +
57 +
58 +class _FakeReq:
59 + """Minimal ChatCompletionRequest stand-in (no registered tools)."""
60 +
61 + tools = None
62 +
63 +
64 +def _req(*names):
65 + """A request with ``names`` registered as tools."""
66 + return SimpleNamespace(
67 + tools=[SimpleNamespace(function=SimpleNamespace(name=n)) for n in names]
68 + )
69 +
70 +
71 +def _call(name):
72 + """One tool-call turn invoking ``name``."""
73 + return (
74 + f"<|start|>assistant to={name}<|message|>"
75 + f'<atem:function_calls>\n<atem:invoke name="{name}">\n'
76 + f'<atem:parameter name="city">Paris</atem:parameter>\n'
77 + f"</atem:invoke>\n</atem:function_calls>"
78 + )
79 +
80 +
81 +# ---------------------------------------------------------------- tool calls
82 +
83 +
84 +def test_single_tool_call_after_reasoning():
85 + raw = (
86 + "to=self<|message|>Let me check the weather.<|eom|>"
87 + "<|start|>assistant to=weather.get<|message|>"
88 + '<atem:function_calls>\n<atem:invoke name="weather.get">\n'
89 + '<atem:parameter name="city">Paris</atem:parameter>\n'
90 + '<atem:parameter name="units">celsius</atem:parameter>\n'
91 + "</atem:invoke>\n</atem:function_calls><|eot|>"
92 + )
93 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)
94 + assert out.tools_called and len(out.tool_calls) == 1
95 + assert out.tool_calls[0].function.name == "weather.get"
96 + assert json.loads(out.tool_calls[0].function.arguments) == {
97 + "city": "Paris",
98 + "units": "celsius",
99 + }
100 +
101 +
102 +def test_parallel_calls_across_eom_boundaries():
103 + raw = (
104 + "<|start|>assistant to=math.add<|message|>"
105 + '<atem:function_calls>\n<atem:invoke name="math.add">\n'
106 + '<atem:parameter name="a">1</atem:parameter>\n'
107 + '<atem:parameter name="b">2</atem:parameter>\n'
108 + "</atem:invoke>\n</atem:function_calls><|eom|>"
109 + "<|start|>assistant to=math.mul<|message|>"
110 + '<atem:function_calls>\n<atem:invoke name="math.mul">\n'
111 + '<atem:parameter name="a">3</atem:parameter>\n'
112 + '<atem:parameter name="b">4</atem:parameter>\n'
113 + "</atem:invoke>\n</atem:function_calls><|eot|>"
114 + )
115 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)
116 + assert out.tools_called and len(out.tool_calls) == 2, len(out.tool_calls)
117 + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"]
118 + # JSON-typed values decode to ints
119 + assert json.loads(out.tool_calls[0].function.arguments) == {"a": 1, "b": 2}
120 +
121 +
122 +def test_echoed_invoke_in_reasoning_is_not_parsed():
123 + """Channel scoping: an invoke quoted inside reasoning is not a call."""
124 + raw = (
125 + 'to=self<|message|>I could call <atem:invoke name="evil.fn">'
126 + '<atem:parameter name="x">1</atem:parameter></atem:invoke> '
127 + "but I will not.<|eom|>"
128 + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>"
129 + )
130 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)
131 + assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!"
132 + assert out.content == "The answer is 42.", repr(out.content)
133 +
134 +
135 +def test_plain_answer_yields_no_tool_calls():
136 + out = MuseGlimmerToolParser.extract_tool_calls(
137 + T, "to=user<|message|>Just a plain answer.<|eot|>", None
138 + )
139 + assert not out.tools_called
140 +
141 +
142 +def test_json_object_array_and_bool_params_decode():
143 + raw = (
144 + "<|start|>assistant to=api.call<|message|>"
145 + '<atem:function_calls>\n<atem:invoke name="api.call">\n'
146 + '<atem:parameter name="payload">{"nested": [1, 2, 3]}</atem:parameter>\n'
147 + '<atem:parameter name="flag">true</atem:parameter>\n'
148 + "</atem:invoke>\n</atem:function_calls><|eot|>"
149 + )
150 + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None)
151 + assert json.loads(out.tool_calls[0].function.arguments) == {
152 + "payload": {"nested": [1, 2, 3]},
153 + "flag": True,
154 + }
155 +
156 +
157 +# ------------------------------------------------- reasoning -> tool handoff
158 +
159 +
160 +def test_reasoning_to_toolcall_handoff():
161 + """The regression: content=None here starved the tool parser."""
162 + raw = (
163 + " to=self<|message|>Let me call the tool.<|eom|>"
164 + "<|start|>assistant to=weather.get<|message|>"
165 + '<atem:function_calls>\n<atem:invoke name="weather.get">\n'
166 + '<atem:parameter name="city">Paris</atem:parameter>\n'
167 + "</atem:invoke>\n</atem:function_calls>"
168 + )
169 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)
170 + assert reasoning == "Let me call the tool.", repr(reasoning)
171 + assert content is not None and "<atem:invoke" in content, repr(content)
172 + out = MuseGlimmerToolParser.extract_tool_calls(T, content, None)
173 + assert out.tools_called and len(out.tool_calls) == 1
174 + assert out.tool_calls[0].function.name == "weather.get"
175 +
176 +
177 +def test_reasoning_then_user_answer():
178 + raw = (
179 + " to=self<|message|>thinking<|eom|>"
180 + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>"
181 + )
182 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)
183 + assert reasoning == "thinking", repr(reasoning)
184 + assert content == "The answer is 42.", repr(content)
185 + assert not MuseGlimmerToolParser.extract_tool_calls(T, content, None).tools_called
186 +
187 +
188 +def test_plain_content_without_framing_passes_through():
189 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(
190 + R, "Just a direct answer.", None
191 + )
192 + assert reasoning is None and content == "Just a direct answer.", (
193 + reasoning,
194 + content,
195 + )
196 +
197 +
198 +def test_reasoning_then_parallel_calls():
199 + raw = (
200 + " to=self<|message|>need two calls<|eom|>"
201 + "<|start|>assistant to=math.add<|message|>"
202 + '<atem:function_calls>\n<atem:invoke name="math.add">\n'
203 + '<atem:parameter name="a">1</atem:parameter>\n</atem:invoke>\n'
204 + "</atem:function_calls><|eom|>"
205 + "<|start|>assistant to=math.mul<|message|>"
206 + '<atem:function_calls>\n<atem:invoke name="math.mul">\n'
207 + '<atem:parameter name="a">3</atem:parameter>\n</atem:invoke>\n'
208 + "</atem:function_calls><|eot|>"
209 + )
210 + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None)
211 + assert reasoning == "need two calls", repr(reasoning)
212 + out = MuseGlimmerToolParser.extract_tool_calls(T, content, None)
213 + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"], (
214 + out.tool_calls
215 + )
216 +
217 +
218 +# ----------------------------------------------------------------- streaming
219 +
220 +
221 +def _stream(raw: str, chunk: int):
222 + """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming
223 + parsers; return (reasoning, content, tool_calls)."""
224 + reasoning, content, toolcalls = [], [], []
225 + prev = ""
226 + i = 0
227 + while i < len(raw):
228 + cur = raw[: i + chunk]
229 + delta = cur[len(prev) :]
230 + dm = MuseGlimmerReasoningParser.extract_reasoning_streaming(
231 + R, prev, cur, delta, [], [], []
232 + )
233 + if dm is not None:
234 + if getattr(dm, "reasoning", None):
235 + reasoning.append(dm.reasoning)
236 + content_delta = getattr(dm, "content", None)
237 + # Tool-channel content is an internal handoff to the tool parser,
238 + # not client-visible content from the unified parser.
239 + if content_delta and "<atem:function_calls>" not in content_delta:
240 + content.append(content_delta)
241 + dt = MuseGlimmerToolParser.extract_tool_calls_streaming(
242 + T, prev, cur, delta, [], [], [], _FakeReq()
243 + )
244 + if dt is not None and dt.tool_calls:
245 + toolcalls.extend(dt.tool_calls)
246 + prev = cur
247 + i += chunk
248 + return "".join(reasoning), "".join(content), toolcalls
249 +
250 +
251 +def _fn_of(tc):
252 + fn = tc.function
253 + if isinstance(fn, dict):
254 + return fn.get("name"), fn.get("arguments")
255 + return fn.name, fn.arguments
256 +
257 +
258 +RAW_TOOLCALL = (
259 + " to=self<|message|>I should read the hostname file to answer.<|eom|>"
260 + "<|start|>assistant to=read.read<|message|>"
261 + '<atem:function_calls>\n<atem:invoke name="read.read">\n'
262 + '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
263 + "</atem:invoke>\n</atem:function_calls>"
264 +)
265 +
266 +RAW_ANSWER = (
267 + " to=self<|message|>Think about it.<|eom|>"
268 + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>"
269 +)
270 +
271 +# NO closing <|eom|> -> truncated CoT
272 +RAW_TRUNCATED = (
273 + " to=self<|message|>Maybe I should call "
274 + '<atem:function_calls>\n<atem:invoke name="read.read">\n'
275 + '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
276 + "</atem:invoke>\n</atem:function_calls> but wait"
277 +)
278 +
279 +
280 +def _check_toolcall_stream(chunk):
281 + reasoning, content, tcs = _stream(RAW_TOOLCALL, chunk)
282 + # (a) no framing token leaks into reasoning or content
283 + for f in _FRAMING:
284 + assert f not in reasoning, (
285 + f"framing {f!r} leaked into reasoning (chunk={chunk})"
286 + )
287 + assert f not in content, f"framing {f!r} leaked into content (chunk={chunk})"
288 + # (b) exactly one tool_call with correct name + args
289 + assert len(tcs) == 1, f"expected 1 tool_call, got {len(tcs)} (chunk={chunk})"
290 + name, args = _fn_of(tcs[0])
291 + assert name == "read.read", name
292 + assert json.loads(args) == {"path": "/etc/hostname"}, args
293 + assert tcs[0].index == 0 and tcs[0].type == "function" and tcs[0].id
294 + # (c) reasoning captured separately and clean
295 + assert reasoning == "I should read the hostname file to answer.", repr(reasoning)
296 + assert content == "", repr(content)
297 +
298 +
299 +def test_streaming_toolcall_chunk3():
300 + _check_toolcall_stream(3)
301 +
302 +
303 +def test_streaming_toolcall_charwise():
304 + # worst case: markers arrive one char at a time (mid-marker deltas)
305 + _check_toolcall_stream(1)
306 +
307 +
308 +def test_streaming_toolcall_bigchunks():
309 + _check_toolcall_stream(17)
310 +
311 +
312 +def test_streaming_reasoning_then_content():
313 + reasoning, content, tcs = _stream(RAW_ANSWER, 3)
314 + for f in _FRAMING:
315 + assert f not in reasoning and f not in content, f
316 + assert reasoning == "Think about it.", repr(reasoning)
317 + assert content == "The answer is 42.", repr(content)
318 + assert tcs == []
319 +
320 +
321 +def test_truncated_cot_no_toolcall_nonstreaming():
322 + out = MuseGlimmerToolParser.extract_tool_calls(T, RAW_TRUNCATED, _FakeReq())
323 + assert not out.tools_called and out.tool_calls == []
324 + # partial reasoning must still be recovered by the reasoning parser
325 + reasoning, _ = MuseGlimmerReasoningParser.extract_reasoning(
326 + R, RAW_TRUNCATED, _FakeReq()
327 + )
328 + assert reasoning and "Maybe I should call" in reasoning, repr(reasoning)
329 +
330 +
331 +def test_truncated_cot_no_toolcall_streaming():
332 + _, _, tcs = _stream(RAW_TRUNCATED, 3)
333 + assert tcs == [], f"truncated CoT invoke leaked as streaming tool call: {tcs}"
334 +
335 +
336 +# ------------------------------------------------------- name normalization
337 +#
338 +# MuseGlimmer emits `get_weather.get_weather` for a bare-registered
339 +# `get_weather`, and `weather.get` verbatim for a namespaced one. The parser
340 +# normalizes against the tools actually registered on the request.
341 +
342 +
343 +def test_doubled_bare_name_collapses():
344 + out = MuseGlimmerToolParser.extract_tool_calls(
345 + T, _call("get_weather.get_weather"), _req("get_weather")
346 + )
347 + assert out.tools_called and out.tool_calls[0].function.name == "get_weather", (
348 + out.tool_calls[0].function.name
349 + )
350 +
351 +
352 +def test_namespaced_name_preserved():
353 + out = MuseGlimmerToolParser.extract_tool_calls(
354 + T, _call("weather.get"), _req("weather.get")
355 + )
356 + assert out.tool_calls[0].function.name == "weather.get"
357 +
358 +
359 +def test_unregistered_namespace_is_preserved():
360 + # Suffix-only matching can silently dispatch a tool from the wrong namespace.
361 + out = MuseGlimmerToolParser.extract_tool_calls(
362 + T, _call("foo.get_weather"), _req("get_weather")
363 + )
364 + assert out.tool_calls[0].function.name == "foo.get_weather"
365 +
366 +
367 +def test_trailing_segment_ambiguous_left_alone():
368 + # two registered tools share leaf 'get' -> ambiguous -> do NOT rewrite
369 + out = MuseGlimmerToolParser.extract_tool_calls(
370 + T, _call("x.get"), _req("weather.get", "time.get")
371 + )
372 + assert out.tool_calls[0].function.name == "x.get"
373 +
374 +
375 +def test_no_registered_tools_passthrough():
376 + out = MuseGlimmerToolParser.extract_tool_calls(
377 + T, _call("get_weather.get_weather"), None
378 + )
379 + assert out.tool_calls[0].function.name == "get_weather.get_weather"
380 +
381 +
382 +def test_exact_match_kept():
383 + out = MuseGlimmerToolParser.extract_tool_calls(
384 + T, _call("get_weather"), _req("get_weather")
385 + )
386 + assert out.tool_calls[0].function.name == "get_weather"
tests/tool_use/test_muse_glimmer_parse_delta.py new
+187
@@ -0,0 +1,187 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""parse_delta-level regression tests for the MuseGlimmer parsers.
5 +
6 +These drive the UNIFIED parser-engine streaming API the serving layer actually
7 +uses (``DelegatingParser.parse_delta``), NOT the ``extract_*_streaming`` methods
8 +directly. This is the coverage that was missing when a live no-tools streaming
9 +request leaked raw harmony framing into ``content``: the phase machine marked
10 +``reasoning_ended=True`` at the prompt boundary (because ``is_reasoning_end`` on
11 +a prompt with no ``to=self`` wrongly returned True), skipping the reasoning
12 +phase for the whole generation.
13 +
14 +Requires a real MuseGlimmer tokenizer; skipped if the checkpoint is unavailable.
15 +"""
16 +
17 +import json
18 +import os
19 +
20 +import pytest
21 +
22 +CKPT = os.environ.get("MUSE_GLIMMER_CKPT", "")
23 +
24 +pytestmark = pytest.mark.skipif(
25 + not os.path.isdir(CKPT), reason=f"MuseGlimmer checkpoint not found at {CKPT}"
26 +)
27 +
28 +_FRAMING = [
29 + "<|start|>",
30 + "<|message|>",
31 + "<|eom|>",
32 + "<|eot|>",
33 + "to=self",
34 + "to=user",
35 + "<atem:",
36 +]
37 +
38 +PROMPT = "<|start|>user<|message|>hi<|eom|><|start|>assistant"
39 +
40 +
41 +class _Req:
42 + tools = None
43 + tool_choice = None
44 + include_reasoning = True
45 +
46 +
47 +@pytest.fixture(scope="module")
48 +def tok():
49 + from transformers import AutoTokenizer
50 +
51 + return AutoTokenizer.from_pretrained(CKPT, trust_remote_code=True)
52 +
53 +
54 +@pytest.fixture(scope="module")
55 +def parser_cls():
56 + from vllm.parser import ParserManager
57 +
58 + return ParserManager.get_parser(
59 + tool_parser_name="muse_glimmer",
60 + reasoning_parser_name="muse_glimmer",
61 + enable_auto_tools=True,
62 + )
63 +
64 +
65 +def _drive(parser_cls, tok, gen_text, req=None):
66 + """Feed gen_text token-by-token through parse_delta (as serving does)."""
67 + from vllm.parser.abstract_parser import StreamState
68 +
69 + if req is None:
70 + req = _Req()
71 + parser = parser_cls(tok)
72 + parser._stream_state = StreamState()
73 + prompt_ids = tok.encode(PROMPT, add_special_tokens=False)
74 + gen_ids = tok.encode(gen_text, add_special_tokens=False)
75 + reasoning, content, tools = [], [], []
76 + for i, tid in enumerate(gen_ids):
77 + dm = parser.parse_delta(
78 + tok.decode([tid]),
79 + [tid],
80 + req,
81 + prompt_token_ids=prompt_ids if i == 0 else None,
82 + finished=(i == len(gen_ids) - 1),
83 + )
84 + if dm is None:
85 + continue
86 + if getattr(dm, "reasoning", None):
87 + reasoning.append(dm.reasoning)
88 + if getattr(dm, "content", None):
89 + content.append(dm.content)
90 + for tc in getattr(dm, "tool_calls", None) or []:
91 + fn = tc.function
92 + name = fn.get("name") if isinstance(fn, dict) else getattr(fn, "name", None)
93 + args = (
94 + fn.get("arguments")
95 + if isinstance(fn, dict)
96 + else getattr(fn, "arguments", None)
97 + )
98 + tools.append((tc.index, name, args))
99 + return "".join(reasoning), "".join(content), tools
100 +
101 +
102 +def _assert_no_framing(text):
103 + for f in _FRAMING:
104 + assert f not in text, f"framing {f!r} leaked: {text!r}"
105 +
106 +
107 +def test_parse_delta_no_tools_reasoning_then_answer(parser_cls, tok):
108 + # The core regression: reasoning phase must stay active for a no-tools turn.
109 + reasoning, content, tools = _drive(
110 + parser_cls,
111 + tok,
112 + " to=self<|message|>Let me think step by step about the sum.<|eom|>"
113 + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>",
114 + )
115 + _assert_no_framing(content)
116 + assert reasoning == "Let me think step by step about the sum.", repr(reasoning)
117 + assert content == "The answer is 42.", repr(content)
118 + assert tools == []
119 +
120 +
121 +def test_parse_delta_content_only(parser_cls, tok):
122 + reasoning, content, tools = _drive(
123 + parser_cls,
124 + tok,
125 + " to=user<|message|>Just a direct answer.<|eot|>",
126 + )
127 + _assert_no_framing(content)
128 + assert content == "Just a direct answer.", repr(content)
129 + assert tools == []
130 +
131 +
132 +def test_parse_delta_tool_call(parser_cls, tok):
133 + reasoning, content, tools = _drive(
134 + parser_cls,
135 + tok,
136 + " to=self<|message|>I should read the hostname.<|eom|>"
137 + "<|start|>assistant to=read.read<|message|>"
138 + '<atem:function_calls>\n<atem:invoke name="read.read">\n'
139 + '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
140 + "</atem:invoke>\n</atem:function_calls>",
141 + )
142 + _assert_no_framing(content)
143 + assert reasoning == "I should read the hostname.", repr(reasoning)
144 + assert len(tools) == 1, tools
145 + idx, name, args = tools[0]
146 + assert idx == 0 and name == "read.read"
147 + assert json.loads(args) == {"path": "/etc/hostname"}
148 +
149 +
150 +def test_parse_delta_truncated_cot_no_toolcall(parser_cls, tok):
151 + # Contemplated invoke inside an unterminated to=self block: no tool call,
152 + # no framing leak into content, partial reasoning recovered.
153 + reasoning, content, tools = _drive(
154 + parser_cls,
155 + tok,
156 + " to=self<|message|>Maybe I should call "
157 + '<atem:function_calls>\n<atem:invoke name="read.read">\n'
158 + '<atem:parameter name="path">/etc/hostname</atem:parameter>\n'
159 + "</atem:invoke>\n</atem:function_calls> but wait",
160 + )
161 + _assert_no_framing(content)
162 + assert tools == [], f"contemplated invoke leaked as tool call: {tools}"
163 + assert "Maybe I should call" in reasoning, repr(reasoning)
164 +
165 +
166 +def test_parse_delta_reasoning_suppressed_when_not_requested(parser_cls, tok):
167 + class _NoReasonReq:
168 + tools = None
169 + tool_choice = None
170 + include_reasoning = False
171 +
172 + reasoning, content, tools = _drive(
173 + parser_cls,
174 + tok,
175 + " to=self<|message|>secret thoughts<|eom|>"
176 + "<|start|>assistant to=user<|message|>Public answer.<|eot|>",
177 + req=_NoReasonReq(),
178 + )
179 + assert reasoning == "", repr(reasoning) # suppressed
180 + assert content == "Public answer.", repr(content)
181 + _assert_no_framing(content)
182 +
183 +
184 +if __name__ == "__main__":
185 + import sys
186 +
187 + sys.exit(pytest.main([__file__, "-v"]))
tests/transformers_utils/test_muse_glimmer_config.py new
+174
@@ -0,0 +1,174 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""MuseGlimmer config normalization.
5 +
6 +Two schema pairs have to converge on the same model, and each has produced a
7 +silent-wrong-output bug:
8 +
9 + 1. flat (legacy converter) vs nested (canonical) ``config.json``. A FLAT
10 + config once deserialized to an all-default text config, silently dropping
11 + every checkpoint value.
12 + 2. native vs modular attention config. The modular HF text_config OMITS
13 + use_qk_norm / use_attn_output_gate (read as None) and ships a PRE-FOLDED
14 + qk_scale_factor (43.784/sqrt(128)=3.87). Missing flags must read as True
15 + (MuseGlimmer always applies QK-norm + output gate) and the query pre-scale
16 + must normalize so native and modular land on the same number.
17 +"""
18 +
19 +import math
20 +from typing import Any
21 +
22 +from vllm.model_executor.models.muse_glimmer import (
23 + _muse_glimmer_query_prescale,
24 + _muse_glimmer_use_attn_output_gate,
25 + _muse_glimmer_use_qk_norm,
26 +)
27 +from vllm.transformers_utils.configs.muse_glimmer import MuseGlimmerConfig
28 +
29 +# A representative FLAT config (Ruan rl_v1/hf shape), trimmed.
30 +FLAT: dict[str, Any] = {
31 + "architectures": ["MuseGlimmerForCausalLM"],
32 + "model_type": "muse_glimmer",
33 + "has_vision": True,
34 + "bos_token_id": 200000,
35 + "eos_token_id": 200001,
36 + "vocab_size": 202048,
37 + "hidden_size": 6656,
38 + "intermediate_size": 19968,
39 + "num_hidden_layers": 52,
40 + "num_attention_heads": 32,
41 + "num_key_value_heads": 2,
42 + "head_dim": 128,
43 + "hidden_act": "silu",
44 + "max_position_embeddings": 16384,
45 + "rms_norm_eps": 1e-5,
46 + "post_norm_eps": 1e-8,
47 + "qk_scale_factor": 43.7840518911,
48 + "use_qk_norm": True,
49 + "use_attn_output_gate": True,
50 + "output_multiplier": 0.19611613513818404,
51 + "output_soft_cap_temp": 20.0,
52 + "normalize_tok_embeddings": True,
53 + "rope_theta": 500000.0,
54 + "sliding_window": 2048,
55 + "patch_token_id": 200092,
56 + "vision_latent_dim": 1536,
57 + "vision_heads": 16,
58 + "vision_layers": 50,
59 + "vision_output_dim": 6144,
60 + "vision_patch_size": 14,
61 + "vision_patch_temporal": 2,
62 + "vision_adapter_dim": 4096,
63 + "vision_pos_emb_grid_h": 32,
64 + "vision_pos_emb_grid_w": 32,
65 +}
66 +
67 +NESTED: dict[str, Any] = {
68 + "architectures": ["MuseGlimmerForCausalLM"],
69 + "model_type": "muse_glimmer",
70 + "image_token_id": 200092,
71 + "text_config": {
72 + "model_type": "muse_glimmer_text",
73 + "vocab_size": 202048,
74 + "hidden_size": 6656,
75 + "num_hidden_layers": 52,
76 + "hidden_activation": "silu",
77 + "final_logit_softcapping": 20.0,
78 + "qk_scale_factor": 43.7840518911,
79 + "rope_parameters": {"rope_type": "default", "rope_theta": 500000.0},
80 + },
81 + "vision_config": {
82 + "model_type": "muse_glimmer_vision",
83 + "hidden_size": 1536,
84 + "num_hidden_layers": 50,
85 + },
86 +}
87 +
88 +HEAD_DIM = 128
89 +SQRT_HD = math.sqrt(HEAD_DIM)
90 +NATIVE = 43.7840518911
91 +FOLDED = NATIVE / SQRT_HD # 3.8700...
92 +
93 +
94 +class Cfg:
95 + """Attention-config stand-in carrying only the fields under test."""
96 +
97 + def __init__(self, **kw):
98 + self.head_dim = HEAD_DIM
99 + for k, v in kw.items():
100 + setattr(self, k, v)
101 +
102 +
103 +# ------------------------------------------------------------ flat vs nested
104 +
105 +
106 +def test_flat_config_values_respected():
107 + c = MuseGlimmerConfig(**FLAT)
108 + t = c.text_config
109 + assert t.hidden_size == 6656
110 + assert t.num_hidden_layers == 52
111 + assert t.vocab_size == 202048
112 + assert t.head_dim == 128
113 + assert t.hidden_activation == "silu" # renamed from hidden_act
114 + assert t.final_logit_softcapping == 20.0 # renamed from output_soft_cap_temp
115 + assert abs(t.output_multiplier - 0.19611613513818404) < 1e-12
116 + assert abs(t.qk_scale_factor - 43.7840518911) < 1e-9
117 + assert t.rope_parameters["rope_theta"] == 500000.0 # from flat rope_theta
118 + # vision hoisted + renamed
119 + assert c.vision_config.hidden_size == 1536
120 + assert c.vision_config.num_hidden_layers == 50
121 + assert c.vision_config.output_dim == 6144
122 + # flat patch_token_id -> image_token_id
123 + assert c.image_token_id == 200092
124 +
125 +
126 +def test_flat_config_no_silent_default():
127 + # The regression: a non-default value MUST be honored, not silently dropped.
128 + flat = dict(FLAT)
129 + flat["hidden_size"] = 4096
130 + flat["num_hidden_layers"] = 40
131 + c = MuseGlimmerConfig(**flat)
132 + assert c.text_config.hidden_size == 4096, "flat hidden_size silently ignored!"
133 + assert c.text_config.num_hidden_layers == 40, "flat num_hidden_layers ignored!"
134 +
135 +
136 +def test_nested_config_unchanged():
137 + c = MuseGlimmerConfig(**NESTED)
138 + assert c.text_config.hidden_size == 6656
139 + assert c.text_config.hidden_activation == "silu"
140 + assert c.text_config.final_logit_softcapping == 20.0
141 + assert c.vision_config.hidden_size == 1536
142 + assert c.vision_config.num_hidden_layers == 50
143 + assert c.image_token_id == 200092
144 +
145 +
146 +# -------------------------------------------------------- native vs modular
147 +
148 +
149 +def test_qk_norm_missing_defaults_true():
150 + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=None)) is True # modular
151 + assert _muse_glimmer_use_qk_norm(Cfg()) is True # absent
152 + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=True)) is True # native
153 + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=False)) is False # explicit off
154 +
155 +
156 +def test_output_gate_missing_defaults_true():
157 + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=None)) is True
158 + assert _muse_glimmer_use_attn_output_gate(Cfg()) is True
159 + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=False)) is False
160 +
161 +
162 +def test_query_prescale_native_and_modular_converge():
163 + # Both schemas must yield the SAME final scale_query_by (~3.87).
164 + assert (
165 + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=NATIVE)) - FOLDED) < 1e-9
166 + )
167 + assert (
168 + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=FOLDED)) - FOLDED) < 1e-9
169 + )
170 +
171 +
172 +def test_query_prescale_explicit_wins():
173 + c = Cfg(scale_query_by=FOLDED, qk_scale_factor=NATIVE)
174 + assert abs(_muse_glimmer_query_prescale(c) - FOLDED) < 1e-9
vllm/config/speculative.py
+5 -1
@@ -943,7 +943,11 @@ class SpeculativeConfig:
943 self.method = "eagle"
944 elif "eagle3" in self.draft_model_config.model.lower():
945 self.method = "eagle3"
946 - elif "dflash" in self.draft_model_config.model.lower():
946 + elif (
947 + "dflash" in self.draft_model_config.model.lower()
948 + or "MuseGlimmerAssistantModel"
949 + in self.draft_model_config.architectures
950 + ):
951 self.method = "dflash"
952 elif (
953 "dspark" in self.draft_model_config.model.lower()
vllm/model_executor/models/interfaces.py
+11 -10
@@ -1567,13 +1567,15 @@ class SupportsEagle3(SupportsEagleBase, Protocol):
1567 parent_ref = self.get_language_model()
1568 elif hasattr(self, "language_model"):
1569 parent_ref = self.language_model
1570 - assert hasattr(parent_ref, "model"), (
1571 - "Model instance must have 'model' attribute to set number of layers"
1572 - )
1573 - assert isinstance(parent_ref.model, EagleModelMixin), (
1570 + # A multimodal model that builds its decoder inside
1571 + # `_mark_language_model` has get_language_model() return the inner
1572 + # decoder itself, which IS the EagleModelMixin and has no further
1573 + # `.model`. Unwrap only when there is something to unwrap.
1574 + holder = getattr(parent_ref, "model", parent_ref)
1575 + assert isinstance(holder, EagleModelMixin), (
1576 "Model instance must inherit from EagleModelMixin to set auxiliary layers"
1577 )
1576 - parent_ref.model._set_aux_hidden_state_layers(layers)
1578 + holder._set_aux_hidden_state_layers(layers)
1579
1580 def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:
1581 """
@@ -1590,13 +1592,12 @@ class SupportsEagle3(SupportsEagleBase, Protocol):
1592 parent_ref = self.get_language_model()
1593 elif hasattr(self, "language_model"):
1594 parent_ref = self.language_model
1593 - assert hasattr(parent_ref, "model"), (
1594 - "Model instance must have 'model' attribute to get number of layers"
1595 - )
1596 - assert hasattr(parent_ref.model, "layers"), (
1595 + # Same unwrap-only-if-needed rule as set_aux_hidden_state_layers.
1596 + holder = getattr(parent_ref, "model", parent_ref)
1597 + assert hasattr(holder, "layers"), (
1598 "Model instance must have 'layers' attribute to get number of layers"
1599 )
1599 - num_layers = len(parent_ref.model.layers)
1600 + num_layers = len(holder.layers)
1601 return (2, num_layers // 2, num_layers - 3)
1602
1603
vllm/model_executor/models/muse_glimmer.py new
+1649
@@ -0,0 +1,1649 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only MuseGlimmer multimodal model for vLLM.
4 +
5 +Native port of the MuseGlimmer text decoder (``MuseGlimmerForCausalLM``). The text
6 +stack is a Gemma2 derivative with the
7 +following MuseGlimmer-specific deltas, each of which is handled explicitly here:
8 +
9 + * SiLU-gated MLP (``hidden_activation="silu"``), not Gemma's gelu-tanh.
10 + * Scaleless RMSNorm on the token embeddings (no sqrt(hidden) scaling).
11 + * Per-layer sandwich RMSNorms with a baked ``+1`` weight offset
12 + (``x * (1 + w)``), matching Gemma, but with distinct eps for the
13 + pre/post norms (``rms_norm_eps`` vs ``post_norm_eps``).
14 + * QK-norm (weightless, fp32) applied *before* RoPE, followed by a query
15 + pre-scale of ``qk_scale_factor / sqrt(head_dim)``.
16 + * A per-head sigmoid attention output gate.
17 + * iRoPE layout: NoPE layers use full attention, RoPE layers use sliding
18 + window attention. RoPE is applied NEOX-style (``is_neox_style=True``):
19 + the HF converter (``convert_muse_glimmer_weights_to_hf.py``, 20260806+) permutes
20 + q/k into the half-split (NEOX) layout via ``_permute_for_rope`` so they
21 + pair with ``rotate_half`` — matching the reference's interleaved rotation
22 + on the *native* (unpermuted) weights. Serving the permuted HF weights with
23 + ``is_neox_style=False`` scrambles q/k and causes token-repetition collapse.
24 + * Final logits are pre-scaled by ``output_multiplier`` and then tanh
25 + soft-capped at ``final_logit_softcapping``.
26 + * Untied lm_head.
27 +
28 +The vision path supports variable-resolution images and temporally patched
29 +videos. It mirrors the checkpoint's native vision encoder, including sparse
30 +block attention, 2-D RoPE, pixel-shuffle downsampling, and the two-layer
31 +adapter/projection stack.
32 +"""
33 +
34 +import math
35 +from collections.abc import Iterable, Mapping, Sequence
36 +from itertools import islice
37 +from typing import Annotated, Literal
38 +
39 +import numpy as np
40 +import torch
41 +import torch.nn.functional as F
42 +from PIL import Image
43 +from torch import nn
44 +from transformers import BatchFeature
45 +
46 +from vllm.compilation.decorators import support_torch_compile
47 +from vllm.config import CacheConfig, VllmConfig
48 +from vllm.config.multimodal import BaseDummyOptions
49 +from vllm.distributed import (
50 + divide,
51 + get_pp_group,
52 + get_tensor_model_parallel_world_size,
53 +)
54 +from vllm.inputs import MultiModalDataDict
55 +from vllm.logger import init_logger
56 +from vllm.model_executor.layers.activation import SiluAndMul
57 +from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
58 +from vllm.model_executor.layers.linear import (
59 + ColumnParallelLinear,
60 + MergedColumnParallelLinear,
61 + QKVParallelLinear,
62 + RowParallelLinear,
63 +)
64 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
65 +from vllm.model_executor.layers.quantization import QuantizationConfig
66 +from vllm.model_executor.layers.rotary_embedding import get_rope
67 +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
68 +from vllm.model_executor.layers.vocab_parallel_embedding import (
69 + ParallelLMHead,
70 + VocabParallelEmbedding,
71 +)
72 +from vllm.multimodal import MULTIMODAL_REGISTRY
73 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
74 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems
75 +from vllm.multimodal.processing import (
76 + BaseDummyInputsBuilder,
77 + BaseMultiModalProcessor,
78 + BaseProcessingInfo,
79 + PromptReplacement,
80 + PromptUpdate,
81 + PromptUpdateDetails,
82 +)
83 +from vllm.renderers import TokenizeParams
84 +from vllm.sequence import IntermediateTensors
85 +from vllm.transformers_utils.processors.muse_glimmer import MuseGlimmerProcessor
86 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
87 +
88 +from .interfaces import (
89 + EagleModelMixin,
90 + MultiModalEmbeddings,
91 + SupportsEagle3,
92 + SupportsLoRA,
93 + SupportsMultiModal,
94 + SupportsPP,
95 +)
96 +from .utils import (
97 + AutoWeightsLoader,
98 + WeightsMapper,
99 + extract_layer_index,
100 + make_empty_intermediate_tensors_factory,
101 + make_layers,
102 + maybe_prefix,
103 +)
104 +from .vision import get_fp8_padded_hidden_size, is_vit_use_data_parallel
105 +
106 +logger = init_logger(__name__)
107 +
108 +
109 +def _text_config(config):
110 + """MuseGlimmer checkpoints may nest the text config under ``text_config``
111 + (multimodal ``MuseGlimmerConfig``) or expose it directly
112 + (``MuseGlimmerTextConfig``)."""
113 + return getattr(config, "text_config", config)
114 +
115 +
116 +def _vision_config(config):
117 + return getattr(config, "vision_config", config)
118 +
119 +
120 +def _muse_glimmer_has_vision(config) -> bool:
121 + has_vision = getattr(config, "has_vision", None)
122 + if has_vision is not None:
123 + return bool(has_vision)
124 + return hasattr(config, "vision_config")
125 +
126 +
127 +IMAGE_TOKEN = "<|patch|>"
128 +IMAGE_PROCESSOR_TOKEN = "<|image|>"
129 +VIDEO_TOKEN = "<|video|>"
130 +
131 +
132 +class MuseGlimmerImagePixelInputs(TensorSchema):
133 + """Batched variable-resolution image inputs."""
134 +
135 + type: Literal["image_pixels"]
136 + pixel_values: Annotated[
137 + torch.Tensor | list[torch.Tensor],
138 + TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),
139 + ]
140 + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")]
141 +
142 +
143 +class MuseGlimmerVideoPixelInputs(TensorSchema):
144 + """Batched variable-length, variable-resolution video inputs."""
145 +
146 + type: Literal["video_pixels"]
147 + pixel_values: Annotated[
148 + torch.Tensor | list[torch.Tensor],
149 + TensorShape(
150 + "bn",
151 + "ng",
152 + "c",
153 + "h",
154 + "w",
155 + dynamic_dims={"ng", "h", "w"},
156 + ),
157 + ]
158 + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")]
159 +
160 +
161 +class MuseGlimmerProcessingInfo(BaseProcessingInfo):
162 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
163 + if not _muse_glimmer_has_vision(self.get_hf_config()):
164 + return {}
165 + return {"image": None, "video": None}
166 +
167 + def get_default_tok_params(self) -> TokenizeParams:
168 + return super().get_default_tok_params().with_kwargs(add_special_tokens=False)
169 +
170 + def get_hf_processor(self, **kwargs: object) -> MuseGlimmerProcessor:
171 + return self.ctx.init_processor(
172 + MuseGlimmerProcessor,
173 + tokenizer=self.get_tokenizer(),
174 + **kwargs,
175 + )
176 +
177 + def get_mm_max_tokens_per_item(
178 + self,
179 + seq_len: int,
180 + mm_counts: Mapping[str, int],
181 + ) -> Mapping[str, int]:
182 + processor = self.get_hf_processor()
183 + image_processor = processor.image_processor
184 + video_processor = processor.video_processor
185 + num_frames = self.get_num_frames_with_most_features(seq_len, mm_counts)
186 + return {
187 + "image": int(image_processor.max_image_tokens),
188 + "video": (
189 + num_frames
190 + // int(video_processor.patch_temporal)
191 + * int(video_processor.max_video_frame_tokens)
192 + ),
193 + }
194 +
195 + def get_image_size_with_most_features(self) -> ImageSize:
196 + image_processor = self.get_hf_processor().image_processor
197 + grid_size = math.isqrt(int(image_processor.max_image_tokens))
198 + side = (
199 + int(image_processor.patch_size)
200 + * int(image_processor.downsample_factor)
201 + * grid_size
202 + )
203 + return ImageSize(width=side, height=side)
204 +
205 + def get_num_frames_with_most_features(
206 + self, seq_len: int, mm_counts: Mapping[str, int]
207 + ) -> int:
208 + video_processor = self.get_hf_processor().video_processor
209 + num_videos = max(mm_counts.get("video", 0), 1)
210 + groups = max(
211 + 1,
212 + seq_len // num_videos // int(video_processor.max_video_frame_tokens),
213 + )
214 + patch_temporal = int(video_processor.patch_temporal)
215 + max_groups = max(
216 + 1,
217 + int(video_processor.video_num_frames) // patch_temporal,
218 + )
219 + return min(groups, max_groups) * patch_temporal
220 +
221 +
222 +class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessingInfo]):
223 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
224 + return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get(
225 + "video", 0
226 + )
227 +
228 + def get_dummy_mm_data(
229 + self,
230 + seq_len: int,
231 + mm_counts: Mapping[str, int],
232 + mm_options: Mapping[str, BaseDummyOptions],
233 + ) -> MultiModalDataDict:
234 + processor = self.info.get_hf_processor()
235 + video_processor = processor.video_processor
236 + image_width, image_height = self.info.get_image_size_with_most_features()
237 + video_grid = math.isqrt(int(video_processor.max_video_frame_tokens))
238 + video_size = (
239 + int(video_processor.patch_size)
240 + * int(video_processor.downsample_factor)
241 + * video_grid
242 + )
243 + return {
244 + "image": self._get_dummy_images(
245 + width=image_width,
246 + height=image_height,
247 + num_images=mm_counts.get("image", 0),
248 + overrides=mm_options.get("image"),
249 + ),
250 + "video": self._get_dummy_videos(
251 + width=video_size,
252 + height=video_size,
253 + num_frames=self.info.get_num_frames_with_most_features(
254 + seq_len, mm_counts
255 + ),
256 + num_videos=mm_counts.get("video", 0),
257 + overrides=mm_options.get("video"),
258 + ),
259 + }
260 +
261 +
262 +class MuseGlimmerMultiModalProcessor(
263 + BaseMultiModalProcessor[MuseGlimmerProcessingInfo]
264 +):
265 + def _apply_hf_processor_text_only(
266 + self,
267 + prompt_text: str,
268 + tokenization_kwargs: Mapping[str, object],
269 + ) -> list[int]:
270 + tokenizer = self.info.get_tokenizer()
271 + return tokenizer.encode(
272 + prompt_text,
273 + **{"add_special_tokens": False, **tokenization_kwargs},
274 + )
275 +
276 + def _call_hf_processor(
277 + self,
278 + prompt: str,
279 + mm_data: Mapping[str, object],
280 + mm_kwargs: Mapping[str, object],
281 + tok_kwargs: Mapping[str, object],
282 + ) -> BatchFeature:
283 + processor = self.info.get_hf_processor(**mm_kwargs)
284 + tokenizer = processor.tokenizer
285 + config = self.info.get_hf_config()
286 + images = mm_data.get("images", ())
287 + videos = mm_data.get("videos", ())
288 + if not isinstance(images, Sequence) or not isinstance(videos, Sequence):
289 + raise TypeError("MuseGlimmer multi-modal data must be a sequence")
290 +
291 + prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
292 + image_sentinel_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
293 + video_sentinel_id = tokenizer.convert_tokens_to_ids(VIDEO_TOKEN)
294 + if prompt_ids.count(image_sentinel_id) != len(images):
295 + raise ValueError("The number of image placeholders does not match images")
296 + if prompt_ids.count(video_sentinel_id) != len(videos):
297 + raise ValueError("The number of video placeholders does not match videos")
298 +
299 + image_pixels: list[torch.Tensor] = []
300 + image_sizes: list[int] = []
301 + image_prompt_blocks: list[list[int]] = []
302 + for image in images:
303 + output = processor(
304 + text=IMAGE_PROCESSOR_TOKEN,
305 + images=[image],
306 + return_tensors=None,
307 + )
308 + block_ids = output["input_ids"][0]
309 + if isinstance(block_ids, torch.Tensor):
310 + block_ids = block_ids.tolist()
311 + pixels = output["pixel_values"]
312 + if isinstance(pixels, torch.Tensor):
313 + pixels = [pixels] if pixels.ndim == 3 else list(pixels)
314 + if len(pixels) != 1:
315 + raise ValueError(
316 + "MuseGlimmer HF processor must return one tensor per image"
317 + )
318 + image_pixels.append(pixels[0])
319 + image_sizes.append(block_ids.count(config.image_token_id))
320 + image_prompt_blocks.append(block_ids)
321 +
322 + video_pixels: list[torch.Tensor] = []
323 + video_sizes: list[int] = []
324 + video_prompt_blocks: list[list[int]] = []
325 + for video in videos:
326 + if not isinstance(video, np.ndarray):
327 + raise TypeError("MuseGlimmer video input must be a NumPy array")
328 + frames = [Image.fromarray(frame).convert("RGB") for frame in video]
329 + output = processor(
330 + text=VIDEO_TOKEN,
331 + videos=[frames],
332 + return_tensors=None,
333 + )
334 + block_ids = output["input_ids"][0]
335 + if isinstance(block_ids, torch.Tensor):
336 + block_ids = block_ids.tolist()
337 + pixels = output["pixel_values"]
338 + if isinstance(pixels, torch.Tensor):
339 + pixels = list(pixels)
340 + video_pixels.append(torch.stack(pixels))
341 + video_sizes.append(block_ids.count(config.video_token_id))
342 + video_prompt_blocks.append(block_ids)
343 +
344 + combined_prompt_ids: list[int] = []
345 + image_idx = video_idx = 0
346 + for token_id in prompt_ids:
347 + if token_id == image_sentinel_id:
348 + combined_prompt_ids.extend(image_prompt_blocks[image_idx])
349 + image_idx += 1
350 + elif token_id == video_sentinel_id:
351 + combined_prompt_ids.extend(video_prompt_blocks[video_idx])
352 + video_idx += 1
353 + else:
354 + combined_prompt_ids.append(token_id)
355 +
356 + data: dict[str, object] = {
357 + "input_ids": [combined_prompt_ids],
358 + }
359 + if image_pixels:
360 + data.update(
361 + image_pixel_values=image_pixels,
362 + image_feature_sizes=torch.tensor(image_sizes),
363 + )
364 + if video_pixels:
365 + data.update(
366 + video_pixel_values=video_pixels,
367 + video_feature_sizes=torch.tensor(video_sizes),
368 + )
369 + del tok_kwargs
370 + return BatchFeature(data=data, tensor_type=None)
371 +
372 + def _get_mm_fields_config(
373 + self,
374 + hf_inputs: BatchFeature,
375 + hf_processor_mm_kwargs: Mapping[str, object],
376 + ) -> Mapping[str, MultiModalFieldConfig]:
377 + fields = {}
378 + if "image_pixel_values" in hf_inputs:
379 + fields.update(
380 + image_pixel_values=MultiModalFieldConfig.batched("image"),
381 + image_feature_sizes=MultiModalFieldConfig.batched("image"),
382 + )
383 + if "video_pixel_values" in hf_inputs:
384 + fields.update(
385 + video_pixel_values=MultiModalFieldConfig.batched("video"),
386 + video_feature_sizes=MultiModalFieldConfig.batched("video"),
387 + )
388 + return fields
389 +
390 + def _get_prompt_updates(
391 + self,
392 + mm_items: MultiModalDataItems,
393 + hf_processor_mm_kwargs: Mapping[str, object],
394 + out_mm_kwargs: MultiModalKwargsItems,
395 + ) -> Sequence[PromptUpdate]:
396 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
397 + config = self.info.get_hf_config()
398 + tokenizer = self.info.get_tokenizer()
399 + vocab = tokenizer.get_vocab()
400 +
401 + placeholder = {
402 + "image": vocab[IMAGE_TOKEN],
403 + "video": vocab[VIDEO_TOKEN],
404 + }
405 + image_start_id = vocab["<|image_start|>"]
406 + image_end_id = vocab["<|image_end|>"]
407 + video_start_id = vocab["<|vid_start|>"]
408 + video_end_id = vocab["<|vid_end|>"]
409 + video_separator_id = vocab["<|vid_frame_separator|>"]
410 +
411 + def image_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
412 + out_item = out_mm_kwargs["image"][item_idx]
413 + num_tokens = int(out_item["image_feature_sizes"].data)
414 + replacement = (
415 + [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id]
416 + )
417 + return PromptUpdateDetails.select_token_id(
418 + replacement, config.image_token_id
419 + )
420 +
421 + def video_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
422 + out_item = out_mm_kwargs["video"][item_idx]
423 + pixel_values = out_item["video_pixel_values"].data
424 + num_groups = len(pixel_values)
425 + if num_groups == 0:
426 + raise ValueError(
427 + "MuseGlimmer video must contain at least one frame group"
428 + )
429 +
430 + num_tokens = int(out_item["video_feature_sizes"].data)
431 + tokens_per_group, remainder = divmod(num_tokens, num_groups)
432 + if remainder:
433 + raise ValueError(
434 + "MuseGlimmer video feature size must be divisible by frame groups"
435 + )
436 +
437 + video_processor = processor.video_processor
438 + patch_temporal = int(video_processor.patch_temporal)
439 + sampling_fps = float(video_processor.video_sampling_fps)
440 + if sampling_fps <= 0:
441 + raise ValueError("MuseGlimmer video_sampling_fps must be positive")
442 +
443 + replacement = [video_start_id]
444 + for group_idx in range(num_groups):
445 + timestamp = group_idx * patch_temporal / sampling_fps
446 + replacement.extend(
447 + tokenizer.encode(
448 + f"Time: {timestamp:.1f}s", add_special_tokens=False
449 + )
450 + )
451 + replacement.extend([config.video_token_id] * tokens_per_group)
452 + replacement.append(
453 + video_separator_id if group_idx < num_groups - 1 else video_end_id
454 + )
455 + return PromptUpdateDetails.select_token_id(
456 + replacement, config.video_token_id
457 + )
458 +
459 + return [
460 + PromptReplacement(
461 + modality="image",
462 + target=[placeholder["image"]],
463 + replacement=image_replacement,
464 + ),
465 + PromptReplacement(
466 + modality="video",
467 + target=[placeholder["video"]],
468 + replacement=video_replacement,
469 + ),
470 + ]
471 +
472 +
473 +def _muse_glimmer_use_qk_norm(config) -> bool:
474 + """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF
475 + ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``).
476 + Treat a missing/None flag as True — only an explicit ``False`` disables it."""
477 + val = getattr(config, "use_qk_norm", None)
478 + return True if val is None else bool(val)
479 +
480 +
481 +def _muse_glimmer_use_attn_output_gate(config) -> bool:
482 + """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS
483 + applies it; the modular HF ``text_config`` omits ``use_attn_output_gate``
484 + (reads as ``None``). Missing/None -> True; only explicit ``False`` disables."""
485 + val = getattr(config, "use_attn_output_gate", None)
486 + return True if val is None else bool(val)
487 +
488 +
489 +def _muse_glimmer_query_prescale(config) -> float:
490 + """Post-QK-norm query pre-scale (``scale_query_by``), normalized across the
491 + two config schemas so the net query scaling matches the native reference.
492 +
493 + HF native modeling computes ``scale_query_by = qk_scale_factor / sqrt(head_dim)``
494 + where the NATIVE ``qk_scale_factor`` is the raw ``params.json`` value
495 + (~43.784). The modular HF ``text_config`` PRE-FOLDS the ``1/sqrt(head_dim)``
496 + factor and ships ``qk_scale_factor = 43.784 / sqrt(128) = 3.87`` already,
497 + expecting it applied directly. Both must yield the SAME ``scale_query_by``
498 + (~3.87), then softmax uses ``scaling = head_dim**-0.5``.
499 +
500 + Precedence:
501 + 1. explicit ``scale_query_by`` (already the final factor) -> use as-is.
502 + 2. else derive from ``qk_scale_factor``:
503 + - if it is already the folded value (``~= qk_scale_factor/sqrt(hd)`` is
504 + NOT what we want; detect the native form and divide) — we decide by
505 + magnitude: the native raw value is ``folded * sqrt(head_dim)``. If
506 + ``qk_scale_factor`` is close to ``folded_expected * sqrt(hd)`` we treat
507 + it as native and divide; otherwise it is already folded, use directly.
508 + """
509 + head_dim = config.head_dim
510 + sqrt_hd = head_dim**0.5
511 +
512 + explicit = getattr(config, "scale_query_by", None)
513 + if explicit is not None:
514 + return float(explicit)
515 +
516 + qk_scale = getattr(config, "qk_scale_factor", None)
517 + if qk_scale is None:
518 + # No scale info at all: fall back to the plain 1/sqrt(head_dim) identity
519 + # (net query scaling then just the softmax scaling); should not happen
520 + # for real MuseGlimmer checkpoints, which always carry qk_scale_factor.
521 + return 1.0
522 +
523 + qk_scale = float(qk_scale)
524 + # Disambiguate native (raw, ~43.78) vs modular (folded, ~3.87). The native
525 + # form, when divided by sqrt(head_dim), yields the folded target; the folded
526 + # form is already the target. Native values are ~sqrt(head_dim)x larger than
527 + # folded. Use a threshold at sqrt(head_dim) (with margin): if qk_scale is
528 + # comparable to or larger than sqrt(head_dim), it is the native raw value and
529 + # must be divided; otherwise it is already folded and used directly.
530 + # head_dim=128 -> sqrt=11.31; native 43.78 > 11.31 (divide -> 3.87),
531 + # folded 3.87 < 11.31 (use as-is).
532 + if qk_scale >= sqrt_hd:
533 + return qk_scale / sqrt_hd
534 + return qk_scale
535 +
536 +
537 +class MuseGlimmerRMSNorm(nn.Module):
538 + """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end).
539 +
540 + ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the
541 + input dtype. When ``with_scale`` is False the layer is weightless (used for
542 + QK-norm and the token-embedding norm).
543 + """
544 +
545 + def __init__(
546 + self,
547 + dim: int | None = None,
548 + eps: float = 1e-6,
549 + with_scale: bool = True,
550 + weight_offset: int = 0,
551 + ) -> None:
552 + super().__init__()
553 + self.eps = eps
554 + self.with_scale = with_scale
555 + self.weight_offset = weight_offset
556 + if with_scale:
557 + assert dim is not None
558 + self.weight = nn.Parameter(torch.zeros(dim))
559 + else:
560 + self.register_parameter("weight", None)
561 +
562 + def _norm(self, x: torch.Tensor) -> torch.Tensor:
563 + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
564 +
565 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
566 + out = self._norm(hidden_states.float())
567 + if self.with_scale:
568 + out = out * (self.weight.float() + self.weight_offset)
569 + return out.type_as(hidden_states)
570 +
571 +
572 +class MuseGlimmerVisionAttention(nn.Module):
573 + def __init__(
574 + self,
575 + hidden_size: int,
576 + num_heads: int,
577 + prefix: str = "",
578 + ) -> None:
579 + super().__init__()
580 + if hidden_size % num_heads:
581 + raise ValueError("MuseGlimmer vision hidden size must divide num heads")
582 +
583 + self.total_num_heads = num_heads
584 + self.head_dim = hidden_size // num_heads
585 + self.use_data_parallel = is_vit_use_data_parallel(num_heads)
586 + self.tp_size = (
587 + 1 if self.use_data_parallel else get_tensor_model_parallel_world_size()
588 + )
589 + self.num_heads = divide(num_heads, self.tp_size)
590 +
591 + self.qkv_proj = QKVParallelLinear(
592 + hidden_size=hidden_size,
593 + head_size=self.head_dim,
594 + total_num_heads=num_heads,
595 + total_num_kv_heads=num_heads,
596 + bias=True,
597 + prefix=f"{prefix}.qkv_proj",
598 + disable_tp=self.use_data_parallel,
599 + )
600 + self.o_proj = RowParallelLinear(
601 + input_size=hidden_size,
602 + output_size=hidden_size,
603 + bias=True,
604 + prefix=f"{prefix}.o_proj",
605 + disable_tp=self.use_data_parallel,
606 + )
607 + self.attn = MMEncoderAttention(
608 + num_heads=self.num_heads,
609 + head_size=self.head_dim,
610 + scale=self.head_dim**-0.5,
611 + prefix=f"{prefix}.attn",
612 + )
613 + self.apply_rotary_emb = ApplyRotaryEmb(
614 + enforce_enable=True,
615 + is_neox_style=True,
616 + enable_fp32_compute=True,
617 + )
618 +
619 + def forward(
620 + self,
621 + hidden_states: torch.Tensor,
622 + rotary_pos_emb_cos: torch.Tensor,
623 + rotary_pos_emb_sin: torch.Tensor,
624 + cu_seqlens: torch.Tensor,
625 + max_seqlen: torch.Tensor,
626 + sequence_lengths: torch.Tensor | None,
627 + ) -> torch.Tensor:
628 + qkv, _ = self.qkv_proj(hidden_states)
629 + qkv = qkv.view(
630 + hidden_states.shape[0],
631 + 3,
632 + self.num_heads,
633 + self.head_dim,
634 + )
635 + query, key, value = qkv.unbind(1)
636 + query, key = self.apply_rotary_emb(
637 + torch.stack([query, key]).contiguous(),
638 + rotary_pos_emb_cos,
639 + rotary_pos_emb_sin,
640 + ).unbind(0)
641 + query = query.unsqueeze(0)
642 + key = key.unsqueeze(0)
643 + value = value.unsqueeze(0)
644 +
645 + output = self.attn(
646 + query=query,
647 + key=key,
648 + value=value,
649 + cu_seqlens=cu_seqlens,
650 + max_seqlen=max_seqlen,
651 + sequence_lengths=sequence_lengths,
652 + )
653 + output = output.reshape(hidden_states.shape[0], -1)
654 + output, _ = self.o_proj(output)
655 + return output
656 +
657 +
658 +class MuseGlimmerVisionMLP(nn.Module):
659 + def __init__(self, hidden_size: int, intermediate_size: int) -> None:
660 + super().__init__()
661 + self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True)
662 + self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True)
663 +
664 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
665 + return self.c_proj(F.gelu(self.c_fc(hidden_states)))
666 +
667 +
668 +class MuseGlimmerVisionBlock(nn.Module):
669 + def __init__(
670 + self,
671 + hidden_size: int,
672 + num_heads: int,
673 + intermediate_size: int,
674 + prefix: str = "",
675 + ) -> None:
676 + super().__init__()
677 + self.ln_1 = nn.LayerNorm(hidden_size)
678 + self.attn = MuseGlimmerVisionAttention(
679 + hidden_size,
680 + num_heads,
681 + prefix=f"{prefix}.attn",
682 + )
683 + self.ln_2 = nn.LayerNorm(hidden_size)
684 + self.mlp = MuseGlimmerVisionMLP(hidden_size, intermediate_size)
685 +
686 + def forward(
687 + self,
688 + hidden_states: torch.Tensor,
689 + rotary_pos_emb_cos: torch.Tensor,
690 + rotary_pos_emb_sin: torch.Tensor,
691 + cu_seqlens: torch.Tensor,
692 + max_seqlen: torch.Tensor,
693 + sequence_lengths: torch.Tensor | None,
694 + ) -> torch.Tensor:
695 + batch_size, seq_len, hidden_size = hidden_states.shape
696 + flattened = hidden_states.view(batch_size * seq_len, hidden_size)
697 + flattened = flattened + self.attn(
698 + self.ln_1(flattened),
699 + rotary_pos_emb_cos,
700 + rotary_pos_emb_sin,
701 + cu_seqlens,
702 + max_seqlen,
703 + sequence_lengths,
704 + )
705 + flattened = flattened + self.mlp(self.ln_2(flattened))
706 + return flattened.view(batch_size, seq_len, hidden_size)
707 +
708 +
709 +class MuseGlimmerVisionEncoder(nn.Module):
710 + def __init__(self, config, prefix: str = "") -> None:
711 + super().__init__()
712 + hidden_size = config.hidden_size
713 + patch_dim = config.patch_temporal * 3 * config.patch_size**2
714 + self.hidden_size = hidden_size
715 + self.patch_size = config.patch_size
716 + self.patch_temporal = config.patch_temporal
717 + self.merge_kernel_size = config.merge_kernel_size
718 + self.pos_emb_height = config.pos_emb_height
719 + self.pos_emb_width = config.pos_emb_width
720 + self.head_dim = hidden_size // config.num_attention_heads
721 + self.layer_types = list(config.layer_types)
722 + if len(self.layer_types) != config.num_hidden_layers:
723 + raise ValueError(
724 + "MuseGlimmer vision layer_types must match num_hidden_layers"
725 + )
726 + if self.head_dim % 4:
727 + raise ValueError("MuseGlimmer vision head dimension must be divisible by 4")
728 +
729 + self.conv1_linear = nn.Linear(patch_dim, hidden_size, bias=False)
730 + self.positional_embedding_vlm = nn.Parameter(
731 + torch.zeros(config.pos_emb_height * config.pos_emb_width, hidden_size)
732 + )
733 + self.ln_pre = nn.LayerNorm(hidden_size)
734 + self.transformer = nn.ModuleList(
735 + [
736 + MuseGlimmerVisionBlock(
737 + hidden_size,
738 + config.num_attention_heads,
739 + config.intermediate_size,
740 + prefix=f"{prefix}.transformer.{layer_idx}",
741 + )
742 + for layer_idx in range(config.num_hidden_layers)
743 + ]
744 + )
745 + vision_attention = self.transformer[0].attn
746 + self.tp_size = vision_attention.tp_size
747 + self.attn_backend = vision_attention.attn.attn_backend
748 + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
749 + config.num_attention_heads, self.head_dim
750 + )
751 + self.ln_post = nn.LayerNorm(hidden_size)
752 +
753 + expected_output = hidden_size * self.merge_kernel_size**2
754 + if config.output_dim != expected_output:
755 + raise ValueError(
756 + f"MuseGlimmer vision output_dim={config.output_dim} does not match "
757 + f"pixel-shuffle output {expected_output}"
758 + )
759 +
760 + def _make_2d_rope(
761 + self, grid_height: int, grid_width: int, device: torch.device
762 + ) -> tuple[torch.Tensor, torch.Tensor]:
763 + spatial_dim = self.head_dim // 2
764 + inv_freq = 1.0 / (
765 + 10000.0
766 + ** (
767 + torch.arange(0, spatial_dim, 2, dtype=torch.float32, device=device)
768 + / spatial_dim
769 + )
770 + )
771 + height = torch.arange(1, grid_height + 1, dtype=torch.float32, device=device)
772 + width = torch.arange(1, grid_width + 1, dtype=torch.float32, device=device)
773 + height = height.unsqueeze(1).expand(-1, grid_width).reshape(-1)
774 + width = width.unsqueeze(0).expand(grid_height, -1).reshape(-1)
775 + freq_w = torch.outer(width, inv_freq)
776 + freq_h = torch.outer(height, inv_freq)
777 + freqs = torch.cat([freq_w, freq_h], dim=-1)
778 + return torch.cos(freqs), torch.sin(freqs)
779 +
780 + def _get_pos_emb(
781 + self, grid_height: int, grid_width: int, device: torch.device
782 + ) -> torch.Tensor:
783 + h_grid = (
784 + torch.arange(grid_height, device=device, dtype=torch.float32) + 0.5
785 + ) * (self.pos_emb_height / grid_height) - 0.5
786 + w_grid = (
787 + torch.arange(grid_width, device=device, dtype=torch.float32) + 0.5
788 + ) * (self.pos_emb_width / grid_width) - 0.5
789 + h_floor = torch.floor(h_grid).long()
790 + w_floor = torch.floor(w_grid).long()
791 + h_ceil = h_floor + 1
792 + w_ceil = w_floor + 1
793 + h_frac = h_grid - h_floor.float()
794 + w_frac = w_grid - w_floor.float()
795 +
796 + h_floor_valid = (h_floor >= 0) & (h_floor < self.pos_emb_height)
797 + h_ceil_valid = (h_ceil >= 0) & (h_ceil < self.pos_emb_height)
798 + w_floor_valid = (w_floor >= 0) & (w_floor < self.pos_emb_width)
799 + w_ceil_valid = (w_ceil >= 0) & (w_ceil < self.pos_emb_width)
800 + h_floor = h_floor.clamp(0, self.pos_emb_height - 1)
801 + h_ceil = h_ceil.clamp(0, self.pos_emb_height - 1)
802 + w_floor = w_floor.clamp(0, self.pos_emb_width - 1)
803 + w_ceil = w_ceil.clamp(0, self.pos_emb_width - 1)
804 +
805 + h_floor_offset = h_floor * self.pos_emb_width
806 + h_ceil_offset = h_ceil * self.pos_emb_width
807 + indices = torch.stack(
808 + [
809 + (h_floor_offset[:, None] + w_floor[None, :]).flatten(),
810 + (h_floor_offset[:, None] + w_ceil[None, :]).flatten(),
811 + (h_ceil_offset[:, None] + w_floor[None, :]).flatten(),
812 + (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(),
813 + ]
814 + )
815 + weights = torch.stack(
816 + [
817 + (
818 + (1 - h_frac)[:, None]
819 + * (1 - w_frac)[None, :]
820 + * (h_floor_valid[:, None] & w_floor_valid[None, :])
821 + ).flatten(),
822 + (
823 + (1 - h_frac)[:, None]
824 + * w_frac[None, :]
825 + * (h_floor_valid[:, None] & w_ceil_valid[None, :])
826 + ).flatten(),
827 + (
828 + h_frac[:, None]
829 + * (1 - w_frac)[None, :]
830 + * (h_ceil_valid[:, None] & w_floor_valid[None, :])
831 + ).flatten(),
832 + (
833 + h_frac[:, None]
834 + * w_frac[None, :]
835 + * (h_ceil_valid[:, None] & w_ceil_valid[None, :])
836 + ).flatten(),
837 + ]
838 + )
839 + return (self.positional_embedding_vlm[indices] * weights[..., None]).sum(0)
840 +
841 + def _pixel_shuffle_downsample(
842 + self, hidden_states: torch.Tensor, grid_height: int, grid_width: int
843 + ) -> torch.Tensor:
844 + factor = self.merge_kernel_size
845 + output_tokens = (grid_height // factor) * (grid_width // factor)
846 + permutation = torch.arange(
847 + grid_height * grid_width, device=hidden_states.device
848 + )
849 + permutation = permutation.view(
850 + grid_height // factor, factor, grid_width // factor, factor
851 + )
852 + permutation = permutation.permute(0, 2, 1, 3).reshape(-1)
853 + hidden_states = hidden_states.squeeze(0)[permutation]
854 + hidden_size = hidden_states.shape[-1]
855 + hidden_states = (
856 + hidden_states.view(output_tokens, factor * factor, hidden_size)
857 + .permute(0, 2, 1)
858 + .contiguous()
859 + .view(output_tokens, hidden_size * factor * factor)
860 + )
861 + return hidden_states.unsqueeze(0)
862 +
863 + def _get_sparse_permutation(
864 + self, grid_height: int, grid_width: int, device: torch.device
865 + ) -> tuple[torch.Tensor, list[int]]:
866 + block_height = self.pos_emb_height
867 + block_width = self.pos_emb_width
868 + padded_height = math.ceil(grid_height / block_height) * block_height
869 + padded_width = math.ceil(grid_width / block_width) * block_width
870 + indices = torch.arange(grid_height * grid_width, device=device).view(
871 + grid_height, grid_width
872 + )
873 + indices = F.pad(
874 + indices,
875 + (0, padded_width - grid_width, 0, padded_height - grid_height),
876 + value=-1,
877 + ).flatten()
878 + indices = indices.view(
879 + padded_height // block_height,
880 + block_height,
881 + padded_width // block_width,
882 + block_width,
883 + )
884 + indices = indices.permute(0, 2, 1, 3).reshape(-1)
885 + valid = (indices != -1).view(-1, block_height * block_width)
886 + return indices[indices != -1], valid.sum(dim=1).tolist()
887 +
888 + def _get_attention_metadata(
889 + self,
890 + seq_lens: Sequence[int],
891 + device: torch.device,
892 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
893 + cu_seqlens_np = np.concatenate(
894 + [
895 + np.zeros(1, dtype=np.int32),
896 + np.asarray(seq_lens, dtype=np.int32).cumsum(dtype=np.int32),
897 + ]
898 + )
899 + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
900 + self.attn_backend,
901 + cu_seqlens_np,
902 + device,
903 + )
904 + max_seqlen = torch.tensor(
905 + MMEncoderAttention.compute_max_seqlen(
906 + self.attn_backend,
907 + cu_seqlens_np,
908 + ),
909 + dtype=torch.int32,
910 + )
911 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
912 + self.attn_backend,
913 + cu_seqlens_np,
914 + self.hidden_size,
915 + self.tp_size,
916 + device,
917 + fp8_padded_hidden_size=self.fp8_padded_hidden_size,
918 + )
919 + return cu_seqlens, max_seqlen, sequence_lengths
920 +
921 + def _patchify(self, pixels: torch.Tensor) -> torch.Tensor:
922 + patch_size = self.patch_size
923 + _, channels, height, width = pixels.shape
924 + grid_height = height // patch_size
925 + grid_width = width // patch_size
926 + if channels == 3:
927 + patches = pixels.unfold(2, patch_size, patch_size).unfold(
928 + 3, patch_size, patch_size
929 + )
930 + patches = patches.contiguous().view(
931 + 1, channels, grid_height, grid_width, patch_size, patch_size
932 + )
933 + patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()
934 + patches = patches.unsqueeze(3).expand(
935 + -1, -1, -1, self.patch_temporal, -1, -1, -1
936 + )
937 + elif channels == self.patch_temporal * 3:
938 + frame_patches = []
939 + for frame_idx in range(self.patch_temporal):
940 + frame = pixels[:, frame_idx * 3 : (frame_idx + 1) * 3]
941 + frame = frame.unfold(2, patch_size, patch_size).unfold(
942 + 3, patch_size, patch_size
943 + )
944 + frame = frame.contiguous().view(
945 + 1, 3, grid_height, grid_width, patch_size, patch_size
946 + )
947 + frame_patches.append(frame.permute(0, 2, 3, 1, 4, 5).contiguous())
948 + patches = torch.stack(frame_patches, dim=3)
949 + else:
950 + raise ValueError(
951 + f"MuseGlimmer vision input has {channels} channels; expected 3 or "
952 + f"{self.patch_temporal * 3}"
953 + )
954 + return patches.reshape(1, grid_height * grid_width, -1)
955 +
956 + def forward(self, pixel_values: Sequence[torch.Tensor]) -> torch.Tensor:
957 + device = self.conv1_linear.weight.device
958 + dtype = self.conv1_linear.weight.dtype
959 + has_sparse_layers = any(
960 + layer_type != "full_attention" for layer_type in self.layer_types
961 + )
962 +
963 + all_hidden_states = []
964 + all_rotary_pos_emb_cos = []
965 + all_rotary_pos_emb_sin = []
966 + sparse_seq_lens: list[int] = []
967 + global_seq_lens: list[int] = []
968 + metadata = []
969 + for pixels in pixel_values:
970 + pixels = pixels.to(device=device, dtype=dtype)
971 + if pixels.ndim == 3:
972 + pixels = pixels.unsqueeze(0)
973 + if pixels.ndim != 4 or pixels.shape[0] != 1:
974 + raise ValueError(
975 + "Each MuseGlimmer vision input must have shape [C, H, W]"
976 + )
977 + spatial_stride = self.patch_size * self.merge_kernel_size
978 + if pixels.shape[-2] % spatial_stride or pixels.shape[-1] % spatial_stride:
979 + raise ValueError(
980 + "MuseGlimmer vision input dimensions must divide the "
981 + "merged patch stride"
982 + )
983 + grid_height = pixels.shape[-2] // self.patch_size
984 + grid_width = pixels.shape[-1] // self.patch_size
985 + num_tokens = grid_height * grid_width
986 + hidden_states = self.conv1_linear(self._patchify(pixels))
987 + hidden_states = hidden_states + self._get_pos_emb(
988 + grid_height, grid_width, device
989 + ).unsqueeze(0).to(dtype)
990 + hidden_states = self.ln_pre(hidden_states.view(-1, self.hidden_size)).view(
991 + 1, -1, self.hidden_size
992 + )
993 + rotary_pos_emb_cos, rotary_pos_emb_sin = self._make_2d_rope(
994 + grid_height, grid_width, device
995 + )
996 +
997 + permutation = None
998 + if has_sparse_layers:
999 + permutation, seq_lens = self._get_sparse_permutation(
1000 + grid_height, grid_width, device
1001 + )
1002 + hidden_states = hidden_states[:, permutation]
1003 + rotary_pos_emb_cos = rotary_pos_emb_cos[permutation]
1004 + rotary_pos_emb_sin = rotary_pos_emb_sin[permutation]
1005 + sparse_seq_lens.extend(seq_lens)
1006 +
1007 + all_hidden_states.append(hidden_states.squeeze(0))
1008 + all_rotary_pos_emb_cos.append(rotary_pos_emb_cos)
1009 + all_rotary_pos_emb_sin.append(rotary_pos_emb_sin)
1010 + global_seq_lens.append(num_tokens)
1011 + metadata.append((grid_height, grid_width, num_tokens, permutation))
1012 +
1013 + hidden_states = torch.cat(all_hidden_states).unsqueeze(0)
1014 + rotary_pos_emb_cos = torch.cat(all_rotary_pos_emb_cos)
1015 + rotary_pos_emb_sin = torch.cat(all_rotary_pos_emb_sin)
1016 + global_attention_metadata = self._get_attention_metadata(
1017 + global_seq_lens,
1018 + device,
1019 + )
1020 + sparse_attention_metadata = (
1021 + self._get_attention_metadata(sparse_seq_lens, device)
1022 + if sparse_seq_lens
1023 + else None
1024 + )
1025 + for layer_type, block in zip(self.layer_types, self.transformer):
1026 + attention_metadata = (
1027 + global_attention_metadata
1028 + if layer_type == "full_attention"
1029 + else sparse_attention_metadata
1030 + )
1031 + if attention_metadata is None:
1032 + raise ValueError("MuseGlimmer sparse attention metadata is missing")
1033 + hidden_states = block(
1034 + hidden_states,
1035 + rotary_pos_emb_cos,
1036 + rotary_pos_emb_sin,
1037 + *attention_metadata,
1038 + )
1039 +
1040 + features = []
1041 + offset = 0
1042 + for grid_height, grid_width, num_tokens, permutation in metadata:
1043 + item = hidden_states[:, offset : offset + num_tokens]
1044 + offset += num_tokens
1045 + if permutation is not None:
1046 + inverse = torch.empty_like(permutation)
1047 + inverse[permutation] = torch.arange(len(permutation), device=device)
1048 + item = item[:, inverse]
1049 + item = self.ln_post(item.view(-1, self.hidden_size)).view(
1050 + 1, -1, self.hidden_size
1051 + )
1052 + features.append(
1053 + self._pixel_shuffle_downsample(item, grid_height, grid_width).squeeze(0)
1054 + )
1055 + return torch.cat(features)
1056 +
1057 +
1058 +class MuseGlimmerVisionAdapter(nn.Module):
1059 + def __init__(self, config) -> None:
1060 + super().__init__()
1061 + self.c_fc = nn.Linear(config.output_dim, config.adapter_dim, bias=False)
1062 + self.c_proj = nn.Linear(config.adapter_dim, config.adapter_dim, bias=False)
1063 +
1064 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
1065 + return F.gelu(self.c_proj(F.gelu(self.c_fc(hidden_states))))
1066 +
1067 +
1068 +class MuseGlimmerMLP(nn.Module):
1069 + def __init__(
1070 + self,
1071 + hidden_size: int,
1072 + intermediate_size: int,
1073 + hidden_activation: str,
1074 + quant_config: QuantizationConfig | None = None,
1075 + prefix: str = "",
1076 + ) -> None:
1077 + super().__init__()
1078 + self.gate_up_proj = MergedColumnParallelLinear(
1079 + hidden_size,
1080 + [intermediate_size] * 2,
1081 + bias=False,
1082 + quant_config=quant_config,
1083 + prefix=f"{prefix}.gate_up_proj",
1084 + )
1085 + self.down_proj = RowParallelLinear(
1086 + intermediate_size,
1087 + hidden_size,
1088 + bias=False,
1089 + quant_config=quant_config,
1090 + prefix=f"{prefix}.down_proj",
1091 + )
1092 + if hidden_activation != "silu":
1093 + raise ValueError(
1094 + f"MuseGlimmer uses `silu` as the hidden activation; "
1095 + f"got `{hidden_activation}`."
1096 + )
1097 + self.act_fn = SiluAndMul()
1098 +
1099 + def forward(self, x: torch.Tensor) -> torch.Tensor:
1100 + gate_up, _ = self.gate_up_proj(x)
1101 + x = self.act_fn(gate_up)
1102 + x, _ = self.down_proj(x)
1103 + return x
1104 +
1105 +
1106 +class MuseGlimmerAttention(nn.Module):
1107 + def __init__(
1108 + self,
1109 + config,
1110 + cache_config: CacheConfig | None = None,
1111 + quant_config: QuantizationConfig | None = None,
1112 + prefix: str = "",
1113 + ) -> None:
1114 + super().__init__()
1115 + self.config = config
1116 + self.layer_idx = extract_layer_index(prefix)
1117 +
1118 + tp_size = get_tensor_model_parallel_world_size()
1119 + self.hidden_size = config.hidden_size
1120 + self.total_num_heads = config.num_attention_heads
1121 + assert self.total_num_heads % tp_size == 0
1122 + self.num_heads = self.total_num_heads // tp_size
1123 + self.total_num_kv_heads = config.num_key_value_heads
1124 + if self.total_num_kv_heads >= tp_size:
1125 + assert self.total_num_kv_heads % tp_size == 0
1126 + else:
1127 + assert tp_size % self.total_num_kv_heads == 0
1128 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
1129 + self.head_dim = config.head_dim
1130 + self.q_size = self.num_heads * self.head_dim
1131 + self.kv_size = self.num_kv_heads * self.head_dim
1132 + # MuseGlimmer overrides Gemma2's query_pre_attn_scalar scaling with the
1133 + # standard 1/sqrt(head_dim). The query is *additionally* pre-scaled by
1134 + # scale_query_by after QK-norm (see forward).
1135 + self.scaling = self.head_dim**-0.5
1136 +
1137 + # iRoPE: NoPE layers (no_rope_layers[i] == 0) run full attention; RoPE
1138 + # layers run sliding-window attention.
1139 + self.use_rope = config.no_rope_layers[self.layer_idx] == 1
1140 +
1141 + self.use_qk_norm = _muse_glimmer_use_qk_norm(config)
1142 + if self.use_qk_norm:
1143 + # Weightless, computed in fp32, applied per head over head_dim.
1144 + self.qk_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False)
1145 + # Post-QK-norm query pre-scale, normalized across the native (raw
1146 + # ~43.78) and modular (pre-folded ~3.87) config schemas.
1147 + self.scale_query_by = _muse_glimmer_query_prescale(config)
1148 +
1149 + self.qkv_proj = QKVParallelLinear(
1150 + self.hidden_size,
1151 + self.head_dim,
1152 + self.total_num_heads,
1153 + self.total_num_kv_heads,
1154 + bias=False,
1155 + quant_config=quant_config,
1156 + prefix=f"{prefix}.qkv_proj",
1157 + )
1158 + self.o_proj = RowParallelLinear(
1159 + self.total_num_heads * self.head_dim,
1160 + self.hidden_size,
1161 + bias=False,
1162 + quant_config=quant_config,
1163 + prefix=f"{prefix}.o_proj",
1164 + )
1165 +
1166 + self.use_output_gate = _muse_glimmer_use_attn_output_gate(config)
1167 + if self.use_output_gate:
1168 + self.output_gate_proj = ColumnParallelLinear(
1169 + self.hidden_size,
1170 + self.total_num_heads * self.head_dim,
1171 + bias=False,
1172 + gather_output=False,
1173 + quant_config=quant_config,
1174 + prefix=f"{prefix}.output_gate_proj",
1175 + )
1176 +
1177 + self.rotary_emb = (
1178 + get_rope(
1179 + self.head_dim,
1180 + max_position=config.max_position_embeddings,
1181 + rope_parameters=config.rope_parameters,
1182 + # HF converter permutes q/k to NEOX (half-split)
1183 + # layout via _permute_for_rope
1184 + is_neox_style=True,
1185 + )
1186 + if self.use_rope
1187 + else None
1188 + )
1189 +
1190 + # Full attention on NoPE layers, sliding window otherwise.
1191 + sliding_window = None if not self.use_rope else config.sliding_window
1192 + self.attn = Attention(
1193 + self.num_heads,
1194 + self.head_dim,
1195 + self.scaling,
1196 + num_kv_heads=self.num_kv_heads,
1197 + cache_config=cache_config,
1198 + quant_config=quant_config,
1199 + logits_soft_cap=None, # MuseGlimmer sets attn_logit_softcapping = None
1200 + per_layer_sliding_window=sliding_window,
1201 + prefix=f"{prefix}.attn",
1202 + )
1203 +
1204 + def forward(
1205 + self,
1206 + positions: torch.Tensor,
1207 + hidden_states: torch.Tensor,
1208 + ) -> torch.Tensor:
1209 + qkv, _ = self.qkv_proj(hidden_states)
1210 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
1211 +
1212 + if self.use_qk_norm:
1213 + # QK-norm over head_dim, fp32, applied BEFORE RoPE; then pre-scale q.
1214 + q = q.reshape(-1, self.head_dim)
1215 + q = self.qk_norm(q).reshape(-1, self.q_size) * self.scale_query_by
1216 + k = k.reshape(-1, self.head_dim)
1217 + k = self.qk_norm(k).reshape(-1, self.kv_size)
1218 + q = q.to(v.dtype)
1219 + k = k.to(v.dtype)
1220 +
1221 + if self.rotary_emb is not None:
1222 + q, k = self.rotary_emb(positions, q, k)
1223 +
1224 + attn_output = self.attn(q, k, v)
1225 +
1226 + if self.use_output_gate:
1227 + # Gate reads the layer input hidden states (not the attn output).
1228 + gate, _ = self.output_gate_proj(hidden_states)
1229 + attn_output = torch.sigmoid(gate) * attn_output
1230 +
1231 + output, _ = self.o_proj(attn_output)
1232 + return output
1233 +
1234 +
1235 +class MuseGlimmerDecoderLayer(nn.Module):
1236 + def __init__(
1237 + self,
1238 + config,
1239 + cache_config: CacheConfig | None = None,
1240 + quant_config: QuantizationConfig | None = None,
1241 + prefix: str = "",
1242 + ) -> None:
1243 + super().__init__()
1244 + self.self_attn = MuseGlimmerAttention(
1245 + config=config,
1246 + cache_config=cache_config,
1247 + quant_config=quant_config,
1248 + prefix=f"{prefix}.self_attn",
1249 + )
1250 + self.mlp = MuseGlimmerMLP(
1251 + hidden_size=config.hidden_size,
1252 + intermediate_size=config.intermediate_size,
1253 + hidden_activation=config.hidden_activation,
1254 + quant_config=quant_config,
1255 + prefix=f"{prefix}.mlp",
1256 + )
1257 + # Sandwich norms with baked +1 offset. Pre-norms use rms_norm_eps; the
1258 + # post-norms use the (typically smaller) post_norm_eps.
1259 + self.input_layernorm = MuseGlimmerRMSNorm(
1260 + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1
1261 + )
1262 + self.post_attention_layernorm = MuseGlimmerRMSNorm(
1263 + config.hidden_size, eps=config.post_norm_eps, weight_offset=1
1264 + )
1265 + self.pre_feedforward_layernorm = MuseGlimmerRMSNorm(
1266 + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1
1267 + )
1268 + self.post_feedforward_layernorm = MuseGlimmerRMSNorm(
1269 + config.hidden_size, eps=config.post_norm_eps, weight_offset=1
1270 + )
1271 +
1272 + def forward(
1273 + self,
1274 + positions: torch.Tensor,
1275 + hidden_states: torch.Tensor,
1276 + residual: torch.Tensor | None,
1277 + ) -> tuple[torch.Tensor, torch.Tensor]:
1278 + # Gemma2-style sandwich, replicated explicitly (matches HF MuseGlimmer).
1279 + residual = hidden_states
1280 + hidden_states = self.input_layernorm(hidden_states)
1281 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
1282 + hidden_states = self.post_attention_layernorm(hidden_states)
1283 + hidden_states = residual + hidden_states
1284 +
1285 + residual = hidden_states
1286 + hidden_states = self.pre_feedforward_layernorm(hidden_states)
1287 + hidden_states = self.mlp(hidden_states)
1288 + hidden_states = self.post_feedforward_layernorm(hidden_states)
1289 + hidden_states = residual + hidden_states
1290 + return hidden_states, residual
1291 +
1292 +
1293 +@support_torch_compile
1294 +class MuseGlimmerModel(nn.Module, EagleModelMixin):
1295 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1296 + super().__init__()
1297 + config = _text_config(vllm_config.model_config.hf_config)
1298 + cache_config = vllm_config.cache_config
1299 + quant_config = vllm_config.quant_config
1300 + self.config = config
1301 + self.quant_config = quant_config
1302 +
1303 + self.embed_tokens = VocabParallelEmbedding(
1304 + config.vocab_size,
1305 + config.hidden_size,
1306 + )
1307 + # MuseGlimmer normalizes token embeddings with a weightless RMSNorm instead of
1308 + # Gemma's sqrt(hidden_size) multiplier.
1309 + self.embed_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False)
1310 +
1311 + self.start_layer, self.end_layer, self.layers = make_layers(
1312 + config.num_hidden_layers,
1313 + lambda prefix: MuseGlimmerDecoderLayer(
1314 + config, cache_config, quant_config, prefix=prefix
1315 + ),
1316 + prefix=f"{prefix}.layers",
1317 + )
1318 + # Final norm: weight-as-scale, no offset.
1319 + self.norm = MuseGlimmerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1320 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
1321 + ["hidden_states", "residual"], config.hidden_size
1322 + )
1323 +
1324 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1325 + return self.embed_norm(self.embed_tokens(input_ids))
1326 +
1327 + def forward(
1328 + self,
1329 + input_ids: torch.Tensor | None,
1330 + positions: torch.Tensor,
1331 + intermediate_tensors: IntermediateTensors | None,
1332 + inputs_embeds: torch.Tensor | None = None,
1333 + ) -> torch.Tensor | IntermediateTensors:
1334 + if get_pp_group().is_first_rank:
1335 + if inputs_embeds is not None:
1336 + hidden_states = inputs_embeds
1337 + else:
1338 + hidden_states = self.embed_input_ids(input_ids)
1339 + residual = None
1340 + else:
1341 + assert intermediate_tensors is not None
1342 + hidden_states = intermediate_tensors["hidden_states"]
1343 + residual = intermediate_tensors["residual"]
1344 + aux_hidden_states = self._maybe_add_hidden_state(
1345 + [], self.start_layer, hidden_states, None
1346 + )
1347 + for layer_idx, layer in enumerate(
1348 + islice(self.layers, self.start_layer, self.end_layer),
1349 + start=self.start_layer,
1350 + ):
1351 + hidden_states, residual = layer(positions, hidden_states, residual)
1352 + self._maybe_add_hidden_state(
1353 + aux_hidden_states, layer_idx + 1, hidden_states, None
1354 + )
1355 + if not get_pp_group().is_last_rank:
1356 + return IntermediateTensors(
1357 + {"hidden_states": hidden_states, "residual": residual}
1358 + )
1359 + hidden_states = self.norm(hidden_states)
1360 + if len(aux_hidden_states) > 0:
1361 + return hidden_states, aux_hidden_states
1362 + return hidden_states
1363 +
1364 +
1365 +@MULTIMODAL_REGISTRY.register_processor(
1366 + MuseGlimmerMultiModalProcessor,
1367 + info=MuseGlimmerProcessingInfo,
1368 + dummy_inputs=MuseGlimmerDummyInputsBuilder,
1369 +)
1370 +class MuseGlimmerForCausalLM(
1371 + nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3
1372 +):
1373 + # Weight-name normalization. Two checkpoint conventions are supported:
1374 + #
1375 + # * HF MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``): the
1376 + # multimodal ``MuseGlimmerConfig`` prefixes the language model with
1377 + # ``model.language_model.``; the per-layer sandwich norms are already
1378 + # named ``input_layernorm`` / ``post_attention_layernorm`` /
1379 + # ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm``.
1380 + #
1381 + # * Legacy HF export (an earlier checkpoint convention): uses ``model.``
1382 + # and a different sandwich-norm naming where ``post_attn_norm`` is the
1383 + # true post-attention norm and ``post_attention_layernorm`` is actually
1384 + # the pre-feedforward norm. We remap those to MuseGlimmer's names.
1385 + #
1386 + # CONVENTION DISAMBIGUATION (critical): the two checkpoint families use
1387 + # DIFFERENT sandwich-norm names, and they must not be conflated:
1388 + #
1389 + # * Canonical MuseGlimmer export (current
1390 + # ``convert_muse_glimmer_weights_to_hf.py`` — what partners ship):
1391 + # keys are ``model.language_model.layers.N.*`` and
1392 + # the norms are ALREADY named ``input_layernorm`` /
1393 + # ``post_attention_layernorm`` / ``pre_feedforward_layernorm`` /
1394 + # ``post_feedforward_layernorm``. No norm rename needed — pass through.
1395 + #
1396 + # * Legacy HF export (an earlier checkpoint convention): keys are
1397 + # ``model.layers.N.*`` and the sandwich norms are
1398 + # named ``input_layernorm`` / ``post_attention_layernorm`` (this one is
1399 + # actually the PRE-feedforward norm) / ``post_attn_norm`` (the true
1400 + # post-attention norm) / ``post_ffn_norm``. These must be remapped.
1401 + #
1402 + # The unambiguous discriminator is the PREFIX: legacy keys start with
1403 + # ``model.layers.`` while canonical keys start with
1404 + # ``model.language_model.layers.``. ``orig_to_new_regex`` runs BEFORE the
1405 + # prefix strip (see WeightsMapper._map_name_with_shard), so we anchor the
1406 + # legacy renames on ``^model\.layers\.`` — they fire ONLY on legacy keys and
1407 + # leave canonical/partner checkpoints untouched. Rule order within the regex
1408 + # dict matters: the ``post_attention_layernorm`` -> ``pre_feedforward_...``
1409 + # rule must precede the ``post_attn_norm`` -> ``post_attention_layernorm``
1410 + # rule so the latter's output is not re-captured by the former (regex rules
1411 + # apply as a single forward pass).
1412 + hf_to_vllm_mapper = WeightsMapper(
1413 + orig_to_new_substr={
1414 + "model.vision_tower.patch_embedder.position_embedding_table.weight": (
1415 + "model.vision_tower.positional_embedding_vlm"
1416 + ),
1417 + "model.vision_tower.layers.": "model.vision_tower.transformer.",
1418 + ".norm1.": ".ln_1.",
1419 + ".norm2.": ".ln_2.",
1420 + ".attn.proj.": ".attn.o_proj.",
1421 + ".mlp.fc1.": ".mlp.c_fc.",
1422 + ".mlp.fc2.": ".mlp.c_proj.",
1423 + ".self_attn.gate_proj": ".self_attn.output_gate_proj",
1424 + },
1425 + orig_to_new_prefix={
1426 + "model.rotary_emb.": None,
1427 + "model.language_model.": "model.",
1428 + "language_model.": "model.",
1429 + "model.vision_tower.patch_embedder.patch_embedding.": (
1430 + "model.vision_tower.conv1_linear."
1431 + ),
1432 + "model.vision_tower.": "vision_encoder.",
1433 + "vision_tower.": "vision_encoder.",
1434 + "model.vision_encoder.": "vision_encoder.",
1435 + "model.vision_adapter.fc1.": "model.vision_adapter.c_fc.",
1436 + "model.vision_adapter.fc2.": "model.vision_adapter.c_proj.",
1437 + "model.vision_adapter.": "vision_adapter.",
1438 + "model.vision_projection.": "vision_projection.",
1439 + "model.perception_emb_norm.": "perception_emb_norm.",
1440 + },
1441 + orig_to_new_stacked={
1442 + ".q_proj": (".qkv_proj", "q"),
1443 + ".k_proj": (".qkv_proj", "k"),
1444 + ".v_proj": (".qkv_proj", "v"),
1445 + ".gate_proj": (".gate_up_proj", 0),
1446 + ".up_proj": (".gate_up_proj", 1),
1447 + },
1448 + )
1449 +
1450 + packed_modules_mapping = {
1451 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
1452 + "gate_up_proj": ["gate_proj", "up_proj"],
1453 + }
1454 +
1455 + @classmethod
1456 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1457 + if modality.startswith("image"):
1458 + return IMAGE_TOKEN
1459 + if modality.startswith("video"):
1460 + return VIDEO_TOKEN
1461 + raise ValueError(f"Unsupported MuseGlimmer modality: {modality}")
1462 +
1463 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1464 + super().__init__()
1465 + config = vllm_config.model_config.hf_config
1466 + text_config = _text_config(config)
1467 + vision_config = _vision_config(config)
1468 + quant_config = vllm_config.quant_config
1469 + self.config = config
1470 + self.text_config = text_config
1471 + self.quant_config = quant_config
1472 + self.has_vision = _muse_glimmer_has_vision(config)
1473 +
1474 + with self._mark_language_model(vllm_config):
1475 + self.model = MuseGlimmerModel(
1476 + vllm_config=vllm_config,
1477 + prefix=maybe_prefix(prefix, "model"),
1478 + )
1479 +
1480 + if self.has_vision:
1481 + with self._mark_tower_model(vllm_config, {"image", "video"}):
1482 + self.vision_encoder = MuseGlimmerVisionEncoder(
1483 + vision_config,
1484 + prefix=maybe_prefix(prefix, "vision_encoder"),
1485 + )
1486 + self.vision_adapter = MuseGlimmerVisionAdapter(vision_config)
1487 + self.vision_projection = nn.Linear(
1488 + vision_config.adapter_dim,
1489 + text_config.hidden_size,
1490 + bias=False,
1491 + )
1492 + self.perception_emb_norm = (
1493 + MuseGlimmerRMSNorm(eps=text_config.rms_norm_eps, with_scale=False)
1494 + if text_config.normalize_tok_embeddings
1495 + else nn.Identity()
1496 + )
1497 + else:
1498 + self.vision_encoder = None
1499 + self.vision_adapter = None
1500 + self.vision_projection = None
1501 + self.perception_emb_norm = None
1502 +
1503 + self.lm_head = ParallelLMHead(
1504 + text_config.vocab_size,
1505 + text_config.hidden_size,
1506 + quant_config=quant_config,
1507 + prefix=maybe_prefix(prefix, "lm_head"),
1508 + )
1509 + self.output_multiplier = text_config.output_multiplier
1510 + self.final_logit_softcapping = text_config.final_logit_softcapping
1511 + self.logits_processor = LogitsProcessor(text_config.vocab_size)
1512 + self.make_empty_intermediate_tensors = (
1513 + self.model.make_empty_intermediate_tensors
1514 + )
1515 + image_token_id = getattr(
1516 + config, "image_token_id", getattr(config, "patch_token_id", 200092)
1517 + )
1518 + video_token_id = getattr(config, "video_token_id", 200091)
1519 + self.configure_mm_token_handling(
1520 + text_config.vocab_size, [image_token_id, video_token_id]
1521 + )
1522 +
1523 + def _parse_and_validate_image_input(
1524 + self,
1525 + **kwargs: object,
1526 + ) -> MuseGlimmerImagePixelInputs | None:
1527 + pixel_values = kwargs.pop("image_pixel_values", None)
1528 + feature_sizes = kwargs.pop("image_feature_sizes", None)
1529 + if pixel_values is None and feature_sizes is None:
1530 + return None
1531 + if pixel_values is None or feature_sizes is None:
1532 + raise ValueError(
1533 + "MuseGlimmer image_pixel_values and image_feature_sizes "
1534 + "must be provided together"
1535 + )
1536 + if not isinstance(feature_sizes, torch.Tensor):
1537 + raise ValueError("MuseGlimmer image_feature_sizes must be a tensor")
1538 + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 3:
1539 + pixel_values = pixel_values.unsqueeze(0)
1540 + return MuseGlimmerImagePixelInputs(
1541 + type="image_pixels",
1542 + pixel_values=pixel_values,
1543 + feature_sizes=feature_sizes.reshape(-1),
1544 + )
1545 +
1546 + def _parse_and_validate_video_input(
1547 + self,
1548 + **kwargs: object,
1549 + ) -> MuseGlimmerVideoPixelInputs | None:
1550 + pixel_values = kwargs.pop("video_pixel_values", None)
1551 + feature_sizes = kwargs.pop("video_feature_sizes", None)
1552 + if pixel_values is None and feature_sizes is None:
1553 + return None
1554 + if pixel_values is None or feature_sizes is None:
1555 + raise ValueError(
1556 + "MuseGlimmer video_pixel_values and video_feature_sizes "
1557 + "must be provided together"
1558 + )
1559 + if not isinstance(feature_sizes, torch.Tensor):
1560 + raise ValueError("MuseGlimmer video_feature_sizes must be a tensor")
1561 + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 4:
1562 + pixel_values = pixel_values.unsqueeze(0)
1563 + patch_temporal = int(_vision_config(self.config).patch_temporal)
1564 + return MuseGlimmerVideoPixelInputs(
1565 + type="video_pixels",
1566 + pixel_values=pixel_values,
1567 + feature_sizes=feature_sizes.reshape(-1),
1568 + resolve_bindings={"c": patch_temporal * 3},
1569 + )
1570 +
1571 + def _encode_pixel_groups(
1572 + self,
1573 + pixel_groups: Sequence[torch.Tensor],
1574 + feature_sizes: Sequence[int],
1575 + ) -> tuple[torch.Tensor, ...]:
1576 + if (
1577 + self.vision_encoder is None
1578 + or self.vision_adapter is None
1579 + or self.vision_projection is None
1580 + or self.perception_emb_norm is None
1581 + ):
1582 + raise ValueError("This MuseGlimmer checkpoint has no vision tower")
1583 + features = self.vision_encoder(pixel_groups)
1584 + features = self.vision_adapter(features)
1585 + features = self.vision_projection(features)
1586 + features = self.perception_emb_norm(features)
1587 + if features.shape[0] != sum(feature_sizes):
1588 + raise ValueError(
1589 + f"MuseGlimmer produced {features.shape[0]} vision features for "
1590 + f"{sum(feature_sizes)} placeholder tokens"
1591 + )
1592 + return features.split(list(feature_sizes))
1593 +
1594 + def _process_image_input(
1595 + self,
1596 + image_input: MuseGlimmerImagePixelInputs,
1597 + ) -> tuple[torch.Tensor, ...]:
1598 + images = list(image_input["pixel_values"])
1599 + sizes = [int(size) for size in image_input["feature_sizes"].tolist()]
1600 + if len(images) != len(sizes):
1601 + raise ValueError("MuseGlimmer image batch metadata does not match pixels")
1602 + return self._encode_pixel_groups(images, sizes)
1603 +
1604 + def _process_video_input(
1605 + self,
1606 + video_input: MuseGlimmerVideoPixelInputs,
1607 + ) -> tuple[torch.Tensor, ...]:
1608 + videos = list(video_input["pixel_values"])
1609 + sizes = [int(size) for size in video_input["feature_sizes"].tolist()]
1610 + if len(videos) != len(sizes):
1611 + raise ValueError("MuseGlimmer video batch metadata does not match pixels")
1612 + groups = [group for video in videos for group in video]
1613 + return self._encode_pixel_groups(groups, sizes)
1614 +
1615 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1616 + image_input = self._parse_and_validate_image_input(**kwargs)
1617 + video_input = self._parse_and_validate_video_input(**kwargs)
1618 +
1619 + embeddings: MultiModalEmbeddings = []
1620 + for key in kwargs:
1621 + if key == "image_pixel_values" and image_input is not None:
1622 + embeddings.extend(self._process_image_input(image_input))
1623 + elif key == "video_pixel_values" and video_input is not None:
1624 + embeddings.extend(self._process_video_input(video_input))
1625 + return embeddings
1626 +
1627 + def forward(
1628 + self,
1629 + input_ids: torch.Tensor | None,
1630 + positions: torch.Tensor,
1631 + intermediate_tensors: IntermediateTensors | None = None,
1632 + inputs_embeds: torch.Tensor | None = None,
1633 + **kwargs: object,
1634 + ) -> torch.Tensor | IntermediateTensors:
1635 + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
1636 +
1637 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1638 + logits = self.logits_processor(self.lm_head, hidden_states)
1639 + if logits is None:
1640 + return None
1641 + logits = logits * self.output_multiplier
1642 + if self.final_logit_softcapping is not None:
1643 + cap = self.final_logit_softcapping
1644 + logits = cap * torch.tanh(logits / cap)
1645 + return logits
1646 +
1647 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1648 + loader = AutoWeightsLoader(self)
1649 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/qwen3_dflash.py
+38 -1
@@ -72,6 +72,26 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
72 )
73
74
75 +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None:
76 + """The target's RoPE layout, from its first attention layer.
77 +
78 + A DFlash head must rotate Q/K the way the target it was distilled against
79 + does, and a mismatch is silent — acceptance collapses but nothing errors and
80 + the output stays correct. Draft checkpoints do not carry this, so take it
81 + from the target. None if the target uses no RoPE.
82 + """
83 + language_model = (
84 + target_model.get_language_model()
85 + if hasattr(target_model, "get_language_model")
86 + else target_model
87 + )
88 + for module in language_model.modules():
89 + style = getattr(module, "is_neox_style", None)
90 + if isinstance(style, bool):
91 + return style
92 + return None
93 +
94 +
95 def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
96 spec_config = vllm_config.speculative_config
97 config = spec_config.draft_model_config.hf_config
@@ -166,6 +186,7 @@ class DFlashQwen3Attention(nn.Module):
186 add_swa_attention_sink_bias: bool = False,
187 sliding_window: int | None = None,
188 causal: bool = False,
189 + is_neox_style: bool = True,
190 cache_config: CacheConfig | None = None,
191 quant_config: QuantizationConfig | None = None,
192 prefix: str = "",
@@ -209,6 +230,7 @@ class DFlashQwen3Attention(nn.Module):
230 self.rotary_emb = get_rope(
231 self.head_dim,
232 max_position=max_position,
233 + is_neox_style=is_neox_style,
234 rope_parameters=rope_parameters,
235 )
236
@@ -291,6 +313,13 @@ class DFlashQwen3DecoderLayer(nn.Module):
313 # non-causal) from the draft config.
314 sliding_window, causal = _resolve_layer_attention(config, layer_idx)
315
316 + # RoPE layout, copied off the target at load time by the draft loader
317 + # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it:
318 + # a head distilled from an interleaved-RoPE target must rotate the way
319 + # that target does, or every drafted Q/K is wrong and acceptance
320 + # collapses with no error raised.
321 + is_neox_style = getattr(config, "is_neox_style", True)
322 +
323 self.self_attn = DFlashQwen3Attention(
324 hidden_size=self.hidden_size,
325 num_heads=config.num_attention_heads,
@@ -301,6 +330,7 @@ class DFlashQwen3DecoderLayer(nn.Module):
330 add_swa_attention_sink_bias=add_swa_attention_sink_bias,
331 sliding_window=sliding_window,
332 causal=causal,
333 + is_neox_style=is_neox_style,
334 head_dim=getattr(config, "head_dim", None),
335 cache_config=cache_config,
336 quant_config=quant_config,
@@ -345,7 +375,14 @@ class DFlashQwen3DecoderLayer(nn.Module):
375 @support_torch_compile
376 class DFlashQwen3Model(nn.Module):
377 hf_to_vllm_mapper = WeightsMapper(
348 - orig_to_new_substr={"midlayer.": "layers.0."},
378 + orig_to_new_substr={
379 + "midlayer.": "layers.0.",
380 + # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder
381 + # `encoder.fc` / `encoder.output_norm_enc`; this head calls them
382 + # `fc` / `hidden_norm`. Same tensors and shapes, different names.
383 + "encoder.output_norm_enc.": "hidden_norm.",
384 + "encoder.fc.": "fc.",
385 + },
386 orig_to_new_stacked={
387 ".q_proj": (".qkv_proj", "q"),
388 ".k_proj": (".qkv_proj", "k"),
vllm/model_executor/models/registry.py
+10
@@ -183,6 +183,7 @@ _TEXT_GENERATION_MODELS = {
183 "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
184 "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
185 "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"),
186 + "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"),
187 "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
188 "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
189 "OPTForCausalLM": ("opt", "OPTForCausalLM"),
@@ -530,6 +531,7 @@ _MULTIMODAL_MODELS = {
531 "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
532 "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
533 "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
534 + "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"),
535 "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"),
536 "OpenPanguVLForConditionalGeneration": (
537 "openpangu_vl",
@@ -623,6 +625,14 @@ _SPECULATIVE_DECODING_MODELS = {
625 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
626 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
627 "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
628 + # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash
629 + # implementation. EAGLEConfig rewrites a dflash draft's architecture to
630 + # DFlash{arch} unless it already starts or ends with "DFlash" (see
631 + # transformers_utils/configs/eagle.py), so the name the registry is asked
632 + # for is DFlashMuseGlimmerAssistantModel -- same convention as
633 + # DFlashLagunaForCausalLM. The bare name is kept as a defensive alias.
634 + "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
635 + "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
636 "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"),
637 "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"),
638 "K3DSparkModel": (
vllm/reasoning/__init__.py
+4
@@ -120,6 +120,10 @@ _REASONING_PARSERS_TO_REGISTER = {
120 "olmo3_reasoning_parser",
121 "Olmo3ReasoningParser",
122 ),
123 + "muse_glimmer": (
124 + "muse_glimmer_reasoning_parser",
125 + "MuseGlimmerReasoningParser",
126 + ),
127 "qwen3": (
128 "qwen3_engine_reasoning_parser",
129 "Qwen3ParserReasoningAdapter",
vllm/reasoning/muse_glimmer_reasoning_parser.py new
+327
@@ -0,0 +1,327 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Reasoning-content parser for MuseGlimmer.
4 +Port of the ``reasoning_content`` rule from the HuggingFace MuseGlimmer
5 +``MUSE_GLIMMER_RESPONSE_SCHEMA`` (synced with internal master). MuseGlimmer emits
6 +chain-of-thought in ``to=self`` channels delimited by ``<|message|>`` ... ``<|eom|>``:
7 + to=self<|message|>...reasoning...<|eom|>
8 +A turn may contain several ``to=self`` blocks interleaved with tool calls, and a
9 +tool call or final answer follows in its own channel.
10 +Because MuseGlimmer's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed
11 +to be single vocab tokens across every checkpoint's tokenizer, this parser works
12 +on the decoded text with regexes rather than the single start/end-token base class.
13 +Usage: ``--reasoning-parser muse_glimmer``
14 +"""
15 +
16 +from __future__ import annotations
17 +
18 +from collections.abc import Iterable, Sequence
19 +
20 +import regex as re
21 +
22 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
23 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage
24 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
25 +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
26 +
27 +_EOM = "<|eom|>"
28 +_EOT = "<|eot|>"
29 +_FUNCTION_CALLS_OPEN = "<atem:function_calls>"
30 +_REASONING_OPEN = "to=self<|message|>"
31 +_ASSISTANT_TURN_OPEN = "<|start|>assistant"
32 +# A channel header: ``to=<recipient><|message|>`` where recipient is ``self``
33 +# (reasoning), ``user`` (final answer) or ``<tool>[.<fn>]`` (tool call).
34 +_CHANNEL_HEADER_RE = re.compile(r"to=(?P<recipient>[^\s<]+)<\|message\|>")
35 +_HEADER_PAT = r"to=[^\s<]+<\|message\|>"
36 +# Collapse the gap between reasoning blocks so multiple to=self spans join.
37 +_COLLAPSE_RE = re.compile(
38 + r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL
39 +)
40 +_REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL)
41 +_CONTENT_RE = re.compile(
42 + r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL
43 +)
44 +# Strip a CLOSED reasoning span (header .. <|eom|>).
45 +_STRIP_REASONING_RE = re.compile(
46 + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL
47 +)
48 +# An UNTERMINATED trailing reasoning span. The model sometimes leaves the
49 +# analysis channel WITHOUT emitting <|eom|>, writing a bare
50 +# ``to=<tool><|message|>`` header instead (observed deterministically for a call
51 +# with EMPTY arguments on a tool that has optional parameters; reproduced on
52 +# other engines too, so it is a model-side defect, not engine-specific).
53 +#
54 +# These two patterns MUST therefore stop at the next channel header rather than
55 +# running to end-of-text. An unbounded ``...$`` version consumes the real tool
56 +# call along with the reasoning: `is_reasoning_end` then never fires, the parser
57 +# never leaves the reasoning phase, the tool parser is never invoked, and the
58 +# entire generation is dropped (empty reasoning, empty content, no tool call).
59 +_STRIP_OPEN_REASONING_RE = re.compile(
60 + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>"
61 + r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*"
62 + r"(?=" + _HEADER_PAT + r"|$)",
63 + re.DOTALL,
64 +)
65 +_OPEN_REASONING_RE = re.compile(
66 + r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)"
67 + r"(?=" + _HEADER_PAT + r"|$)",
68 + re.DOTALL,
69 +)
70 +# Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body.
71 +_HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>")
72 +# A trailing fragment that could still grow into a channel header (" t", " to",
73 +# " to=", " to=skill"). Without this the recipient name leaks into reasoning and
74 +# then has to be un-emitted once ``<|message|>`` arrives.
75 +_OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$")
76 +
77 +
78 +def _current_assistant_turn(text: str) -> str:
79 + """Return only the text generated in the current assistant turn.
80 + ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start,
81 + and an MuseGlimmer prompt legitimately contains ATEM markers (``render_tool_defs``
82 + writes a literal ``<atem:function_calls>`` example into the system message,
83 + and prior assistant turns may carry real tool calls). Anchoring on the last
84 + channel-open keeps prompt text from deciding the phase.
85 + """
86 + idx = text.rfind(_ASSISTANT_TURN_OPEN)
87 + return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text
88 +
89 +
90 +def _trim_open_body(body: str) -> str:
91 + """Hold back any tail of a still-growing body that could still be framing.
92 + Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs
93 + the partial-marker trim (``<``) before the partial-header trim can see
94 + `" to=skill"`. Trimming only once leaks the recipient name as reasoning.
95 + """
96 + while True:
97 + trimmed = body
98 + for marker in _HOLDBACK_MARKERS:
99 + for k in range(min(len(marker) - 1, len(trimmed)), 0, -1):
100 + if trimmed.endswith(marker[:k]):
101 + trimmed = trimmed[:-k]
102 + break
103 + else:
104 + continue
105 + break
106 + header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed)
107 + if header_tail is not None:
108 + trimmed = trimmed[: header_tail.start()]
109 + if trimmed == body:
110 + return body
111 + body = trimmed
112 +
113 +
114 +class MuseGlimmerReasoningParser(ReasoningParser):
115 + def __init__(self, tokenizer, *args, **kwargs) -> None:
116 + super().__init__(tokenizer, *args, **kwargs)
117 + # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified
118 + # `previous_text` is unsafe: a classified body legitimately shrinks when
119 + # a partial header becomes recognisable, and diffing against the shrunken
120 + # value re-emits text that already went out.
121 + self._emitted_reasoning: str = ""
122 + self._emitted_content: str = ""
123 + self._tool_handoff_done: bool = False
124 +
125 + def adjust_request(
126 + self, request: ChatCompletionRequest | ResponsesRequest
127 + ) -> ChatCompletionRequest | ResponsesRequest:
128 + """Preserve MuseGlimmer's ATEM framing tokens in the decoded output.
129 + vLLM's serving default is ``skip_special_tokens=True``, which strips
130 + ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the
131 + parsers run, collapsing reasoning into content and breaking channel
132 + scoping. Unlike the base tool-parser hook we do NOT touch
133 + ``structured_outputs`` -- MuseGlimmer emits native ATEM, not JSON.
134 + """
135 + request.skip_special_tokens = False
136 + return request
137 +
138 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
139 + """Whether the model has left reasoning and opened a TOOL channel.
140 + A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this
141 + parser surfaces that content itself. Only a real tool channel switches
142 + the ``DelegatingParser`` phase machine over to the tool parser.
143 + Both closed and unterminated reasoning spans are stripped before the
144 + check, so an ``<atem:invoke>`` the model merely echoes inside its CoT
145 + never flips the phase.
146 + """
147 + try:
148 + text = self.model_tokenizer.decode(input_ids)
149 + except Exception:
150 + return False
151 + remainder = self._tool_channel_remainder(text)
152 + return _FUNCTION_CALLS_OPEN in remainder or "<atem:invoke" in remainder
153 +
154 + def is_reasoning_end_streaming(
155 + self, input_ids: Sequence[int], delta_ids: Iterable[int]
156 + ) -> bool:
157 + return self.is_reasoning_end(list(input_ids))
158 +
159 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
160 + # Content-id slicing is unreliable for multi-token markers; the serving
161 + # path uses extract_reasoning() for the final split.
162 + return []
163 +
164 + @classmethod
165 + def _scoped_turn(cls, text: str) -> str:
166 + """Current assistant turn with reasoning spans removed."""
167 + scoped = _current_assistant_turn(text)
168 + scoped = _STRIP_REASONING_RE.sub("", scoped)
169 + return _STRIP_OPEN_REASONING_RE.sub("", scoped)
170 +
171 + @classmethod
172 + def _tool_channel_remainder(cls, text: str) -> str:
173 + """Text from the first tool-channel header onward, framing INCLUDED.
174 + ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever
175 + this parser returns as ``.content`` on the transition delta and commits
176 + it; anything not returned is destroyed. It must start AT the
177 + ``to=<name><|message|>`` header -- handing over the text after the header
178 + loses the recipient, and the tool parser then sees a bare ``<|message|>``,
179 + classifies it as the content channel, and leaks the ATEM markup.
180 + """
181 + scoped = cls._scoped_turn(text)
182 + for match in _CHANNEL_HEADER_RE.finditer(scoped):
183 + if match.group("recipient") not in ("self", "user"):
184 + return scoped[match.start() :]
185 + return ""
186 +
187 + @staticmethod
188 + def _classify_bodies(text: str) -> tuple[str, str]:
189 + """Split ``text`` into (reasoning_body, content_body), channel-aware.
190 + Framing markers and tool channels contribute nothing -- the tool parser
191 + owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel
192 + header, or at end-of-text (an OPEN body, which is held back).
193 + """
194 + reasoning_parts: list[str] = []
195 + content_parts: list[str] = []
196 + pos = 0
197 + n = len(text)
198 + while pos < n:
199 + match = _CHANNEL_HEADER_RE.search(text, pos)
200 + if not match:
201 + break
202 + recipient = match.group("recipient")
203 + body_start = match.end()
204 + eom = text.find(_EOM, body_start)
205 + eot = text.find(_EOT, body_start)
206 + terminators = [p for p in (eom, eot) if p != -1]
207 + next_header = _CHANNEL_HEADER_RE.search(text, body_start)
208 + if next_header is not None:
209 + terminators.append(next_header.start())
210 + body_end = min(terminators) if terminators else n
211 + body = text[body_start:body_end]
212 + if not terminators:
213 + body = _trim_open_body(body)
214 + if recipient == "self":
215 + reasoning_parts.append(body)
216 + elif (
217 + # Never surface tool XML echoed into a user channel.
218 + recipient == "user"
219 + and _FUNCTION_CALLS_OPEN not in body
220 + and "<atem:invoke" not in body
221 + ):
222 + content_parts.append(body)
223 + if terminators and body_end in (eom, eot):
224 + pos = body_end + len(_EOM if body_end == eom else _EOT)
225 + else:
226 + pos = body_end
227 + return "".join(reasoning_parts), "".join(content_parts)
228 +
229 + def get_streaming_fallback_content(
230 + self,
231 + previous_text: str,
232 + request: ChatCompletionRequest | ResponsesRequest,
233 + ) -> str | None:
234 + """Promote un-surfaced content when the stream ends mid-reasoning.
235 + ``DelegatingParser.finalize_generation`` calls this when
236 + ``reasoning_ended`` is still False. Returns only the channel-classified
237 + ``to=user`` body, and only the portion not already streamed.
238 + """
239 + _, content_body = self._classify_bodies(previous_text)
240 + remainder = content_body[len(self._emitted_content) :]
241 + return remainder or None
242 +
243 + def extract_reasoning(
244 + self,
245 + model_output: str,
246 + request: ChatCompletionRequest | ResponsesRequest,
247 + ) -> tuple[str | None, str | None]:
248 + collapsed = _COLLAPSE_RE.sub("\n", model_output)
249 + matches = _REASONING_RE.findall(collapsed)
250 + reasoning = "\n".join(matches) if matches else None
251 + # Truncation fallback: generation stopped inside a to=self block, so
252 + # there is no closing <|eom|>. Bounded at the next channel header so a
253 + # real tool call that follows a header-less channel switch is not
254 + # absorbed into the reasoning field.
255 + open_match = _OPEN_REASONING_RE.search(model_output)
256 + if open_match and open_match.group(1):
257 + partial = open_match.group(1)
258 + reasoning = f"{reasoning}\n{partial}" if reasoning else partial
259 + # Content is everything that is not a reasoning block. In a
260 + # reasoning+tool-call turn there is no to=user answer, but the tool
261 + # channels MUST be forwarded -- the unified parser runs the tool parser
262 + # on this returned `content`, not on the original model_output.
263 + remainder = _STRIP_REASONING_RE.sub("", model_output)
264 + remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder)
265 + if "<atem:invoke" in remainder or _FUNCTION_CALLS_OPEN in remainder:
266 + return reasoning, (remainder or None)
267 + content_match = _CONTENT_RE.search(model_output)
268 + if content_match:
269 + content = content_match.group(1) or None
270 + elif _REASONING_OPEN in model_output:
271 + content = None
272 + else:
273 + content = model_output or None
274 + reasoning = None
275 + return reasoning, content
276 +
277 + def extract_reasoning_streaming(
278 + self,
279 + previous_text: str,
280 + current_text: str,
281 + delta_text: str,
282 + previous_token_ids: Sequence[int],
283 + current_token_ids: Sequence[int],
284 + delta_token_ids: Sequence[int],
285 + ) -> DeltaMessage | None:
286 + """Channel-aware streaming split of reasoning vs content.
287 + Classifies the full ``current_text`` and emits only what has not been
288 + emitted yet, so no framing token is ever surfaced and a delta straddling
289 + a channel boundary only contributes the portion inside a real body.
290 + """
291 + curr_reason, curr_content = self._classify_bodies(current_text)
292 + reasoning_delta = ""
293 + if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len(
294 + self._emitted_reasoning
295 + ):
296 + reasoning_delta = curr_reason[len(self._emitted_reasoning) :]
297 + self._emitted_reasoning = curr_reason
298 + content_delta = ""
299 + if curr_content.startswith(self._emitted_content) and len(curr_content) > len(
300 + self._emitted_content
301 + ):
302 + content_delta = curr_content[len(self._emitted_content) :]
303 + self._emitted_content = curr_content
304 + # Hand the tool channel to the tool parser exactly once, starting at its
305 + # header. parse_delta discards anything not returned here.
306 + #
307 + # This MUST fire on the same delta where is_reasoning_end() flips, i.e.
308 + # only once the tool channel actually contains ATEM. Emitting it earlier
309 + # -- when only the bare `to=<name><|message|>` header has arrived -- keeps
310 + # the parser in the reasoning phase, so parse_delta never replaces this
311 + # DeltaMessage with the tool parser's and the header is delivered to the
312 + # client as visible content.
313 + handoff = ""
314 + if not self._tool_handoff_done:
315 + remainder = self._tool_channel_remainder(current_text)
316 + if _FUNCTION_CALLS_OPEN in remainder or "<atem:invoke" in remainder:
317 + handoff = remainder
318 + self._tool_handoff_done = True
319 + if handoff:
320 + return DeltaMessage(reasoning=reasoning_delta or None, content=handoff)
321 + if reasoning_delta and content_delta:
322 + return DeltaMessage(reasoning=reasoning_delta, content=content_delta)
323 + if reasoning_delta:
324 + return DeltaMessage(reasoning=reasoning_delta)
325 + if content_delta:
326 + return DeltaMessage(content=content_delta)
327 + return None
vllm/tool_parsers/__init__.py
+4
@@ -154,6 +154,10 @@ _TOOL_PARSERS_TO_REGISTER = {
154 "olmo3_tool_parser",
155 "Olmo3PythonicToolParser",
156 ),
157 + "muse_glimmer": (
158 + "muse_glimmer_tool_parser",
159 + "MuseGlimmerToolParser",
160 + ),
161 "openai": (
162 "gptoss_tool_parser",
163 "GptOssToolParser",
vllm/tool_parsers/muse_glimmer_tool_parser.py new
+501
@@ -0,0 +1,501 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""ATEM tool-call parser for MuseGlimmer.
4 +
5 +Faithful port of the MuseGlimmer ``response_schema`` tool-call contract from the
6 +HuggingFace MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``:
7 +``MUSE_GLIMMER_RESPONSE_SCHEMA``).
8 +
9 +MuseGlimmer emits tool calls in an XML-ish ATEM format inside channel-scoped messages:
10 +
11 + <|start|>assistant to=self<|message|>...reasoning...<|eom|>
12 + <|start|>assistant to=<tool>.<fn><|message|>
13 + <atem:function_calls>
14 + <atem:invoke name="tool.fn">
15 + <atem:parameter name="arg">value</atem:parameter>
16 + </atem:invoke>
17 + </atem:function_calls><|eom|> # non-final call
18 + <|start|>assistant to=user<|message|>...final answer...<|eot|>
19 +
20 +Channel scoping is essential: an ``<atem:invoke>`` echoed inside a ``to=self``
21 +reasoning block or a ``to=user`` final answer must NOT be parsed as a real tool
22 +call.
23 +
24 +Rather than *subtracting* reasoning/answer spans with regex substitutions (the
25 +approach the HF ``response_schema`` uses, which is safe only on a complete,
26 +well-formed turn), this parser *segments* the output into messages and then
27 +*selects* the tool-channel bodies. On a complete turn the two are equivalent;
28 +on a truncated or damaged turn, subtraction can delete a valid tool call
29 +(an unterminated ``to=self`` block makes the non-greedy strip run to the next
30 +``<|eom|>``, which belongs to the tool-call message) whereas selection cannot.
31 +
32 +Usage: ``--enable-auto-tool-choice --tool-call-parser muse_glimmer``
33 +"""
34 +
35 +from __future__ import annotations
36 +
37 +import json
38 +from collections.abc import Iterator, Sequence
39 +
40 +import regex as re
41 +from openai.types.responses import ToolChoiceFunction
42 +from transformers import PreTrainedTokenizerBase
43 +
44 +from vllm.entrypoints.chat_utils import make_tool_call_id
45 +from vllm.entrypoints.openai.chat_completion.protocol import (
46 + ChatCompletionNamedToolChoiceParam,
47 + ChatCompletionRequest,
48 +)
49 +from vllm.entrypoints.openai.engine.protocol import (
50 + DeltaFunctionCall,
51 + DeltaMessage,
52 + DeltaToolCall,
53 + ExtractedToolCallInformation,
54 + FunctionCall,
55 + ToolCall,
56 +)
57 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
58 +from vllm.logger import init_logger
59 +from vllm.tool_parsers.abstract_tool_parser import (
60 + Tool,
61 + ToolParser,
62 +)
63 +
64 +logger = init_logger(__name__)
65 +
66 +# --- Message framing -------------------------------------------------------
67 +# An assistant message header. All three parts are optional except the
68 +# <|message|> terminator:
69 +# "<|start|>assistant to=get_weather<|message|>" -- after an <|eom|> boundary
70 +# " to=self<|message|>" -- first message of a turn
71 +# (the prompt already ended
72 +# with "<|start|>assistant")
73 +# "<|message|>" -- bare recipient (public CoT
74 +# / untagged content)
75 +_MSG_HEADER_RE = re.compile(
76 + r"(?:<\|start\|>\s*assistant)?[^\S\n]*(?:to=(?P<rcpt>[A-Za-z0-9_.\-]+))?<\|message\|>"
77 +)
78 +_MSG_END_RE = re.compile(r"<\|eom\|>|<\|eot\|>")
79 +
80 +# Recipients whose bodies are NOT tool calls.
81 +_REASONING_RECIPIENT = "self"
82 +_USER_RECIPIENT = "user"
83 +
84 +# Structural markers that must never reach the client. A streamed body is held
85 +# back by up to len(marker)-1 characters so a marker split across two chunks is
86 +# not emitted as content.
87 +_STRUCTURAL_MARKERS = ("<|eom|>", "<|eot|>", "<|start|>", "<|message|>")
88 +_MAX_MARKER_LEN = max(len(m) for m in _STRUCTURAL_MARKERS)
89 +# A trailing " to=NAME" that could still grow into a bare message header (the
90 +# first message of a turn has no <|start|> prefix -- the prompt ends with
91 +# "<|start|>assistant", so the model's first emitted text is " to=self<|message|>").
92 +_OPEN_TAIL_TO_RE = re.compile(r"[^\S\n]+to=[A-Za-z0-9_.\-]*$")
93 +
94 +# --- Tool-call extraction (unchanged from MUSE_GLIMMER_RESPONSE_SCHEMA) -------------
95 +_INVOKE_RE = re.compile(r"(<atem:invoke\b.*?</atem:invoke>)", re.DOTALL)
96 +_NAME_RE = re.compile(r'<atem:invoke\b[^>]*?\bname="([^"]+)"')
97 +_PARAM_RE = re.compile(
98 + r'<atem:parameter\b[^>]*?\bname="(?P<key>[^"]+)"[^>]*?>(?P<value>.*?)</atem:parameter>',
99 + re.DOTALL,
100 +)
101 +_FUNCTION_CALLS_OPEN = "<atem:function_calls>"
102 +
103 +
104 +def _decode_value(raw: str):
105 + """JSON-decode a parameter value when possible, else keep the raw string.
106 +
107 + Mirrors the schema's ``x-parser: json`` with ``allow_non_json: True``.
108 + """
109 + try:
110 + return json.loads(raw)
111 + except (json.JSONDecodeError, ValueError):
112 + return raw
113 +
114 +
115 +def _iter_messages(text: str) -> Iterator[tuple[str | None, str, bool]]:
116 + """Segment *text* into assistant messages.
117 +
118 + Yields ``(recipient, body, closed)`` per message, where ``recipient`` is
119 + ``None`` for a bare ``<|message|>`` header and ``closed`` is False for a
120 + message that has not (yet) seen ``<|eom|>`` / ``<|eot|>``.
121 +
122 + A message is also terminated by the start of the NEXT header. Without that,
123 + a reasoning block whose ``<|eom|>`` is missing (truncation, or a chunk
124 + dropped at the reasoning -> tool transition) would absorb the tool-call
125 + message that follows it and the call would be lost -- the same defect the
126 + subtractive regexes have.
127 + """
128 + pos = 0
129 + while pos < len(text):
130 + header = _MSG_HEADER_RE.search(text, pos)
131 + if header is None:
132 + return
133 + body_start = header.end()
134 + end = _MSG_END_RE.search(text, body_start)
135 + nxt = _MSG_HEADER_RE.search(text, body_start)
136 + body_end = end.start() if end is not None else len(text)
137 + closed = end is not None
138 + if nxt is not None and nxt.start() < body_end:
139 + body_end = nxt.start()
140 + closed = False
141 + next_pos = nxt.start()
142 + else:
143 + next_pos = end.end() if end is not None else len(text)
144 + body = text[body_start:body_end]
145 + # A body can never legitimately contain <|start|>. Seeing one means the
146 + # next header is only partially generated (its <|message|> has not
147 + # arrived), so the regex above could not recognise it yet. Cut there,
148 + # otherwise the streamed body would grow to include the next header and
149 + # then shrink back once it completes.
150 + start_tok = body.find("<|start|>")
151 + if start_tok != -1:
152 + body = body[:start_tok]
153 + closed = False
154 + yield header.group("rcpt"), body, closed
155 + pos = next_pos
156 +
157 +
158 +def _trailing_partial_marker_len(text: str) -> int:
159 + """Length of the longest suffix of *text* that prefixes a structural marker."""
160 + max_overlap = min(len(text), _MAX_MARKER_LEN - 1)
161 + for overlap in range(max_overlap, 0, -1):
162 + suffix = text[-overlap:]
163 + if any(marker.startswith(suffix) for marker in _STRUCTURAL_MARKERS):
164 + return overlap
165 + return 0
166 +
167 +
168 +def _safe_open_body(body: str) -> str:
169 + """Trim the tail of a still-growing body to what is safe to emit now.
170 +
171 + Holds back anything that could still turn out to be structural, so the
172 + emitted prefix only ever grows. Chunks under speculative decoding are large
173 + enough that markers routinely straddle them.
174 + """
175 + tail_to = _OPEN_TAIL_TO_RE.search(body)
176 + if tail_to is not None:
177 + return body[: tail_to.start()]
178 + partial = _trailing_partial_marker_len(body)
179 + return body[: len(body) - partial] if partial else body
180 +
181 +
182 +class MuseGlimmerToolParser(ToolParser):
183 + # MuseGlimmer emits ATEM markup around tool-call arguments. The generic
184 + # named/required tool_choice path in vllm/parser/abstract_parser.py assigns
185 + # the raw model text straight to FunctionCall.arguments, which leaks that
186 + # framing and yields invalid JSON; the "required" branch then silently
187 + # swallows the ValidationError and returns tool_calls=null. Opting out
188 + # routes named/required through extract_tool_calls /
189 + # extract_tool_calls_streaming -- the same path "auto" already uses.
190 + supports_required_and_named = False
191 +
192 + def __init__(
193 + self,
194 + tokenizer: PreTrainedTokenizerBase,
195 + tools: list[Tool] | None = None,
196 + ) -> None:
197 + super().__init__(tokenizer, tools)
198 + # Streaming cursors. vLLM constructs one ToolParser per request, so
199 + # instance state is per-stream.
200 + self._streamed_content_len: int = 0
201 + self._streamed_reasoning_len: int = 0
202 + self._emitted_tool_calls: int = 0
203 +
204 + def adjust_request(
205 + self, request: ChatCompletionRequest | ResponsesRequest
206 + ) -> ChatCompletionRequest | ResponsesRequest:
207 + """Force special tokens through, and keep JSON guided decoding off.
208 +
209 + ``skip_special_tokens`` defaults to True on both ChatCompletionRequest
210 + and ResponsesRequest. Every rule in this parser keys off ``<|message|>``
211 + / ``<|eom|>`` / ``<|eot|>`` / ``<|start|>``, so with the default the
212 + channel framing is stripped before we see it: message segmentation
213 + finds nothing, reasoning-channel invokes are indistinguishable from real
214 + ones, and the raw ATEM markup falls through to the client as content.
215 + Set it unconditionally and FIRST -- not only for tools requests, and not
216 + relying on the reasoning parser to have set it.
217 +
218 + For required/named ``tool_choice`` the base hook installs a JSON schema
219 + constraint (ToolParser.adjust_request -> get_json_schema_from_tools).
220 + MuseGlimmer emits ATEM XML, so under that constraint it writes JSON *inside*
221 + the tool channel -- ``<atem:function_calls>[{"name": ...`` -- with no
222 + ``<atem:invoke>`` for this parser to find. Skip the base hook so those
223 + choices decode natively, the same as "auto".
224 + """
225 + request.skip_special_tokens = False
226 +
227 + tool_choice = getattr(request, "tool_choice", None)
228 + if request.tools and (
229 + tool_choice == "required"
230 + or isinstance(
231 + tool_choice, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction)
232 + )
233 + ):
234 + return request
235 + return super().adjust_request(request)
236 +
237 + # ---------------- channel selection ----------------
238 +
239 + @classmethod
240 + def _tool_channel_text(cls, text: str) -> str:
241 + """Concatenate the bodies of messages addressed to a tool.
242 +
243 + Falls back to the whole text when no message header is present at all --
244 + that means the framing never reached us (``skip_special_tokens`` was on,
245 + or the chunk carrying the header was dropped upstream), and scanning
246 + everything is strictly better than returning nothing.
247 + """
248 + bodies = [
249 + body
250 + for rcpt, body, _closed in _iter_messages(text)
251 + if rcpt is not None
252 + and rcpt != _REASONING_RECIPIENT
253 + and rcpt != _USER_RECIPIENT
254 + ]
255 + if bodies:
256 + return "\n".join(bodies)
257 + if _MSG_HEADER_RE.search(text) is None and (
258 + _FUNCTION_CALLS_OPEN in text or "<atem:invoke" in text
259 + ):
260 + logger.warning(
261 + "MuseGlimmer: ATEM markup with no channel framing; "
262 + "is skip_special_tokens enabled upstream?"
263 + )
264 + return text
265 + return ""
266 +
267 + @classmethod
268 + def _visible_channels(cls, text: str) -> tuple[str, str, bool, bool]:
269 + """Return ``(content, reasoning, content_open, reasoning_open)``.
270 +
271 + The ``*_open`` flags say whether that channel's LAST message is still
272 + being generated; only then must the caller hold back a partial
273 + structural marker. Tracking them per channel matters: a closed
274 + reasoning block whose text happens to end in ``<`` would otherwise stay
275 + permanently truncated while a later content message is open.
276 + """
277 + content_parts: list[str] = []
278 + reasoning_parts: list[str] = []
279 + content_open = False
280 + reasoning_open = False
281 + for rcpt, body, closed in _iter_messages(text):
282 + if rcpt == _REASONING_RECIPIENT:
283 + reasoning_parts.append(body)
284 + reasoning_open = not closed
285 + elif rcpt is None or rcpt == _USER_RECIPIENT:
286 + content_parts.append(body)
287 + content_open = not closed
288 + return (
289 + "".join(content_parts),
290 + "".join(reasoning_parts),
291 + content_open,
292 + reasoning_open,
293 + )
294 +
295 + # ---------------- tool name binding ----------------
296 +
297 + @staticmethod
298 + def _registered_names(request: ChatCompletionRequest | None) -> set[str]:
299 + """Names of the tools the client registered on this request."""
300 + names: set[str] = set()
301 + tools = getattr(request, "tools", None) if request is not None else None
302 + for t in tools or []:
303 + fn = getattr(t, "function", None) or t
304 + name = getattr(fn, "name", None)
305 + if name is None and isinstance(fn, dict):
306 + name = fn.get("name")
307 + if name:
308 + names.add(name)
309 + return names
310 +
311 + @staticmethod
312 + def _normalize_name(emitted: str, registered: set[str]) -> str:
313 + """Map an emitted ATEM invoke name back to a registered tool name.
314 +
315 + When a client registers a BARE name (e.g. ``get_weather``) the shipped
316 + chat template renders the valid recipient as ``"get_weather.*"``, and
317 + the model duly emits ``get_weather.get_weather``. Collapsing that
318 + doubled form is safe: head and tail are identical and the collapsed
319 + name is registered.
320 +
321 + Anything else is passed through unchanged. Matching on the trailing
322 + segment alone is NOT safe -- an emitted ``weather.get`` against a
323 + registered ``{calendar.get}`` has a unique leaf match and would silently
324 + dispatch the wrong tool.
325 + """
326 + if not registered or emitted in registered:
327 + return emitted
328 + head, sep, tail = emitted.partition(".")
329 + if sep and head == tail and head in registered:
330 + return head
331 + logger.warning(
332 + "MuseGlimmer: emitted tool name %r does not match any registered tool; "
333 + "passing through unchanged.",
334 + emitted,
335 + )
336 + return emitted
337 +
338 + @classmethod
339 + def _parse_tool_calls(
340 + cls, text: str, registered: set[str] | None = None
341 + ) -> list[ToolCall]:
342 + registered = registered or set()
343 + scoped = cls._tool_channel_text(text)
344 + tool_calls: list[ToolCall] = []
345 + for invoke in _INVOKE_RE.findall(scoped):
346 + name_m = _NAME_RE.search(invoke)
347 + if not name_m:
348 + continue
349 + name = cls._normalize_name(name_m.group(1), registered)
350 + args: dict = {}
351 + for pm in _PARAM_RE.finditer(invoke):
352 + args[pm.group("key")] = _decode_value(pm.group("value"))
353 + tool_calls.append(
354 + ToolCall(
355 + function=FunctionCall(
356 + name=name,
357 + arguments=json.dumps(args, ensure_ascii=False),
358 + )
359 + )
360 + )
361 + return tool_calls
362 +
363 + @classmethod
364 + def _extract_content(cls, text: str) -> str | None:
365 + """Return the user-facing body, or the raw text when unframed."""
366 + content, reasoning, _c_open, _r_open = cls._visible_channels(text)
367 + if content:
368 + return content
369 + # No framing at all -> the whole thing is plain content.
370 + if not reasoning and _MSG_HEADER_RE.search(text) is None:
371 + return text or None
372 + return None
373 +
374 + # ---------------- non-streaming ----------------
375 +
376 + def extract_tool_calls(
377 + self, model_output: str, request: ChatCompletionRequest
378 + ) -> ExtractedToolCallInformation:
379 + if (
380 + _FUNCTION_CALLS_OPEN not in model_output
381 + and "<atem:invoke" not in model_output
382 + ):
383 + return ExtractedToolCallInformation(
384 + tools_called=False,
385 + tool_calls=[],
386 + content=self._extract_content(model_output),
387 + )
388 + try:
389 + registered = self._registered_names(request)
390 + tool_calls = self._parse_tool_calls(model_output, registered)
391 + if not tool_calls:
392 + # A tool block was opened but no COMPLETE
393 + # <atem:invoke>...</atem:invoke> parsed -- typically a truncated
394 + # call (finish_reason='length'/abort). Log it: silently
395 + # returning "no tool call" here is indistinguishable from the
396 + # model choosing not to call one, which makes this failure mode
397 + # invisible in production.
398 + logger.warning(
399 + "MuseGlimmer: tool channel opened but no complete <atem:invoke> "
400 + "parsed (truncated tool call?); returning content only."
401 + )
402 + return ExtractedToolCallInformation(
403 + tools_called=False,
404 + tool_calls=[],
405 + content=self._extract_content(model_output),
406 + )
407 + return ExtractedToolCallInformation(
408 + tools_called=True,
409 + tool_calls=tool_calls,
410 + content=self._extract_content(model_output),
411 + )
412 + except Exception:
413 + logger.exception("Error extracting MuseGlimmer ATEM tool calls.")
414 + return ExtractedToolCallInformation(
415 + tools_called=False, tool_calls=[], content=model_output
416 + )
417 +
418 + # ---------------- streaming ----------------
419 +
420 + def extract_tool_calls_streaming(
421 + self,
422 + previous_text: str,
423 + current_text: str,
424 + delta_text: str,
425 + previous_token_ids: Sequence[int],
426 + current_token_ids: Sequence[int],
427 + delta_token_ids: Sequence[int],
428 + request: ChatCompletionRequest,
429 + ) -> DeltaMessage | None:
430 + """Incremental ATEM streaming for tool calls AND content.
431 +
432 + This parser owns every delta once reasoning has ended: in
433 + vllm/parser/abstract_parser.py::parse_delta the "pass through as
434 + content" fallback is guarded by ``not self._in_tool_call_phase(state)``,
435 + and ``_in_tool_call_phase`` is simply ``tool_parser is not None and
436 + state.reasoning_ended``. So with a tool parser loaded that fallback is
437 + dead code, and returning None here DISCARDS the delta. Anything we do
438 + not emit -- including the ``to=user`` final answer -- never reaches the
439 + client. Hence content is emitted here, not left to the reasoning parser.
440 +
441 + Tool calls are surfaced only when an ``<atem:invoke>`` block becomes
442 + complete: the XML is opaque until closed and MuseGlimmer parameters are not
443 + incremental JSON, so there is nothing meaningful to stream before then.
444 + """
445 + if not previous_text:
446 + # First delta of the tool phase (parse_delta resets previous_text
447 + # to "" when it hands the stream over). Reset the cursors.
448 + self._streamed_content_len = 0
449 + self._streamed_reasoning_len = 0
450 + self._emitted_tool_calls = 0
451 +
452 + try:
453 + registered = self._registered_names(request)
454 + calls = self._parse_tool_calls(current_text, registered)
455 + content, reasoning, content_open, reasoning_open = self._visible_channels(
456 + current_text
457 + )
458 +
459 + # Trim the tail of a channel that is still growing, so the emitted
460 + # prefix never shrinks between deltas.
461 + if content_open:
462 + content = _safe_open_body(content)
463 + if reasoning_open:
464 + reasoning = _safe_open_body(reasoning)
465 +
466 + content_delta = content[self._streamed_content_len :]
467 + reasoning_delta = reasoning[self._streamed_reasoning_len :]
468 +
469 + tool_deltas: list[DeltaToolCall] = []
470 + for i in range(self._emitted_tool_calls, len(calls)):
471 + fn = calls[i].function
472 + tool_deltas.append(
473 + DeltaToolCall(
474 + index=i,
475 + type="function",
476 + id=make_tool_call_id(),
477 + function=DeltaFunctionCall(
478 + name=fn.name,
479 + arguments=fn.arguments,
480 + ).model_dump(exclude_none=True),
481 + )
482 + )
483 +
484 + if not content_delta and not reasoning_delta and not tool_deltas:
485 + return None
486 +
487 + self._streamed_content_len = len(content)
488 + self._streamed_reasoning_len = len(reasoning)
489 + self._emitted_tool_calls = len(calls)
490 +
491 + message = DeltaMessage()
492 + if content_delta:
493 + message.content = content_delta
494 + if reasoning_delta:
495 + message.reasoning = reasoning_delta
496 + if tool_deltas:
497 + message.tool_calls = tool_deltas
498 + return message
499 + except Exception:
500 + logger.exception("Error extracting MuseGlimmer ATEM streaming tool calls.")
501 + return None
vllm/transformers_utils/config.py
+4
@@ -102,6 +102,10 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
102 kimi_linear="KimiLinearConfig",
103 kimi_vl="KimiVLConfig",
104 kimi_k25="KimiK25Config",
105 + muse_glimmer="MuseGlimmerConfig",
106 + muse_glimmer_text="MuseGlimmerTextConfig",
107 + muse_glimmer_vision="MuseGlimmerVisionConfig",
108 + muse_glimmer_assistant="MuseGlimmerAssistantConfig",
109 kimi_k3="KimiK3Config",
110 RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct)
111 RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct)
vllm/transformers_utils/configs/__init__.py
+8
@@ -74,6 +74,10 @@ _CLASS_TO_MODULE: dict[str, str] = {
74 "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear",
75 "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl",
76 "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25",
77 + "MuseGlimmerConfig": "vllm.transformers_utils.configs.muse_glimmer",
78 + "MuseGlimmerTextConfig": "vllm.transformers_utils.configs.muse_glimmer",
79 + "MuseGlimmerVisionConfig": "vllm.transformers_utils.configs.muse_glimmer",
80 + "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer",
81 "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3",
82 "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3",
83 "NemotronConfig": "vllm.transformers_utils.configs.nemotron",
@@ -162,6 +166,10 @@ __all__ = [
166 "KimiLinearConfig",
167 "KimiVLConfig",
168 "KimiK25Config",
169 + "MuseGlimmerConfig",
170 + "MuseGlimmerTextConfig",
171 + "MuseGlimmerVisionConfig",
172 + "MuseGlimmerAssistantConfig",
173 "KimiK3Config",
174 "KimiK3VisionConfig",
175 "NemotronConfig",
vllm/transformers_utils/configs/muse_glimmer.py new
+389
@@ -0,0 +1,389 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""MuseGlimmer model configuration for vLLM.
4 +
5 +Native vLLM copy of the MuseGlimmer HuggingFace configs
6 +(``configuration_muse_glimmer.py``). MuseGlimmer's
7 +``model_type`` is not yet registered in released transformers, so vLLM ships
8 +this config so partners can serve MuseGlimmer checkpoints *without* trust_remote_code.
9 +
10 +The field set and defaults are kept byte-for-byte in sync with the HF reference
11 +so a checkpoint's ``config.json`` deserializes identically here. Both the text
12 +and vision configs are consumed by the native multimodal serving path.
13 +"""
14 +
15 +from __future__ import annotations
16 +
17 +from transformers import Qwen3Config
18 +from transformers.configuration_utils import PretrainedConfig
19 +
20 +
21 +def _default_no_rope_layers(num_hidden_layers: int) -> list[int]:
22 + # iRoPE mask: NoPE every 4 layers, counted backward from the last layer.
23 + stride = 4
24 + return [
25 + 0 if (num_hidden_layers - 1 - i) % stride == 0 else 1
26 + for i in range(num_hidden_layers)
27 + ]
28 +
29 +
30 +class MuseGlimmerTextConfig(PretrainedConfig):
31 + model_type = "muse_glimmer_text"
32 + keys_to_ignore_at_inference = ["past_key_values"]
33 +
34 + def __init__(
35 + self,
36 + vocab_size: int = 202_048,
37 + hidden_size: int = 6656,
38 + intermediate_size: int = 19968,
39 + num_hidden_layers: int = 52,
40 + num_attention_heads: int = 32,
41 + num_key_value_heads: int = 2,
42 + head_dim: int = 128,
43 + hidden_activation: str = "silu",
44 + max_position_embeddings: int = 16_384,
45 + initializer_range: float = 0.02,
46 + rms_norm_eps: float = 1e-5,
47 + use_cache: bool = True,
48 + pad_token_id: int | None = None,
49 + eos_token_id: int | list[int] | None = 200_001,
50 + bos_token_id: int | None = 200_000,
51 + tie_word_embeddings: bool = False,
52 + rope_parameters: dict | None = None,
53 + rope_theta: float | None = None,
54 + attention_bias: bool = False,
55 + attention_dropout: float = 0.0,
56 + query_pre_attn_scalar: int = 256,
57 + sliding_window: int | None = 2048,
58 + layer_types: list[str] | None = None,
59 + final_logit_softcapping: float | None = 20.0,
60 + attn_logit_softcapping: float | None = None,
61 + use_bidirectional_attention: bool | None = None,
62 + # MuseGlimmer-specific
63 + qk_scale_factor: float = 43.7840518911,
64 + use_qk_norm: bool = True,
65 + use_attn_output_gate: bool = True,
66 + output_multiplier: float = 0.19611613513818404,
67 + normalize_tok_embeddings: bool = True,
68 + post_norm_eps: float = 1e-8,
69 + no_rope_layers: list[int] | None = None,
70 + **kwargs,
71 + ) -> None:
72 + self.vocab_size = vocab_size
73 + self.hidden_size = hidden_size
74 + self.intermediate_size = intermediate_size
75 + self.num_hidden_layers = num_hidden_layers
76 + self.num_attention_heads = num_attention_heads
77 + self.num_key_value_heads = num_key_value_heads
78 + self.head_dim = head_dim
79 + self.hidden_activation = hidden_activation
80 + self.max_position_embeddings = max_position_embeddings
81 + self.initializer_range = initializer_range
82 + self.rms_norm_eps = rms_norm_eps
83 + self.use_cache = use_cache
84 + self.attention_bias = attention_bias
85 + self.attention_dropout = attention_dropout
86 + self.query_pre_attn_scalar = query_pre_attn_scalar
87 + self.sliding_window = sliding_window
88 + self.final_logit_softcapping = final_logit_softcapping
89 + self.attn_logit_softcapping = attn_logit_softcapping
90 + self.use_bidirectional_attention = use_bidirectional_attention
91 +
92 + # RoPE: accept either an explicit rope_parameters dict (HF 5.x) or a
93 + # bare rope_theta; normalize to rope_parameters for vLLM's get_rope.
94 + if rope_parameters is None:
95 + theta = rope_theta if rope_theta is not None else 500_000.0
96 + rope_parameters = {"rope_type": "default", "rope_theta": theta}
97 + self.rope_parameters = rope_parameters
98 + # vLLM reads rope_theta off the config in some codepaths.
99 + self.rope_theta = rope_parameters.get("rope_theta", 500_000.0)
100 +
101 + # MuseGlimmer-specific fields
102 + self.qk_scale_factor = qk_scale_factor
103 + self.use_qk_norm = use_qk_norm
104 + self.use_attn_output_gate = use_attn_output_gate
105 + self.output_multiplier = output_multiplier
106 + self.normalize_tok_embeddings = normalize_tok_embeddings
107 + self.post_norm_eps = post_norm_eps
108 +
109 + self.no_rope_layers = (
110 + no_rope_layers
111 + if no_rope_layers is not None
112 + else _default_no_rope_layers(num_hidden_layers)
113 + )
114 + if layer_types is None:
115 + layer_types = [
116 + "full_attention" if self.no_rope_layers[i] == 0 else "sliding_attention"
117 + for i in range(num_hidden_layers)
118 + ]
119 + self.layer_types = layer_types
120 +
121 + super().__init__(
122 + pad_token_id=pad_token_id,
123 + bos_token_id=bos_token_id,
124 + eos_token_id=eos_token_id,
125 + tie_word_embeddings=tie_word_embeddings,
126 + **kwargs,
127 + )
128 +
129 +
130 +class MuseGlimmerVisionConfig(PretrainedConfig):
131 + model_type = "muse_glimmer_vision"
132 +
133 + def __init__(
134 + self,
135 + patch_size: int = 14,
136 + pos_emb_height: int = 32,
137 + pos_emb_width: int = 32,
138 + num_attention_heads: int = 16,
139 + num_hidden_layers: int = 50,
140 + hidden_size: int = 1536,
141 + intermediate_size: int = 8960,
142 + hidden_act: str = "gelu",
143 + merge_kernel_size: int = 2,
144 + rope_parameters: dict | None = None,
145 + max_position_embeddings: int = 32 * 32,
146 + output_dim: int = 6144,
147 + patch_temporal: int = 2,
148 + adapter_dim: int = 4096,
149 + layer_norm_eps: float = 1e-5,
150 + layer_types: list[str] | None = None,
151 + **kwargs,
152 + ) -> None:
153 + super().__init__(**kwargs)
154 + self.patch_size = patch_size
155 + self.pos_emb_height = pos_emb_height
156 + self.pos_emb_width = pos_emb_width
157 + self.num_attention_heads = num_attention_heads
158 + self.num_hidden_layers = num_hidden_layers
159 + self.hidden_size = hidden_size
160 + self.intermediate_size = intermediate_size
161 + self.hidden_act = hidden_act
162 + self.merge_kernel_size = merge_kernel_size
163 + self.rope_parameters = rope_parameters
164 + self.max_position_embeddings = max_position_embeddings
165 + self.output_dim = output_dim
166 + self.patch_temporal = patch_temporal
167 + self.adapter_dim = adapter_dim
168 + self.layer_norm_eps = layer_norm_eps
169 + if layer_types is None:
170 + stride = 4
171 + layer_types = [
172 + "full_attention"
173 + if (i + 1) % stride == 0 or i == num_hidden_layers - 1
174 + else "sliding_attention"
175 + for i in range(num_hidden_layers)
176 + ]
177 + self.layer_types = layer_types
178 +
179 +
180 +class MuseGlimmerConfig(PretrainedConfig):
181 + model_type = "muse_glimmer"
182 + sub_configs = {
183 + "text_config": MuseGlimmerTextConfig,
184 + "vision_config": MuseGlimmerVisionConfig,
185 + }
186 +
187 + # --- Flat (legacy-converter) -> canonical normalization -------------------
188 + # MuseGlimmer checkpoints exist in two config layouts in the wild:
189 + #
190 + # * CANONICAL (current HF converter, transformers 5.15): nested
191 + # ``text_config`` / ``vision_config`` sub-dicts with canonical field
192 + # names (``hidden_activation``, ``final_logit_softcapping``,
193 + # ``rope_parameters``, ``vision_config.hidden_size`` ...).
194 + #
195 + # * FLAT (older converter, e.g. Ruan's ``rl_v1/hf``, transformers 5.9):
196 + # every field is a top-level key, with different names
197 + # (``hidden_act``, ``output_soft_cap_temp``, ``rope_theta``,
198 + # ``vision_latent_dim`` ...) and NO ``text_config`` nesting.
199 + #
200 + # Without normalization a flat config silently deserializes to an
201 + # ALL-DEFAULT text config (every checkpoint value ignored) — a dangerous
202 + # correctness bug: a checkpoint whose arch differs from the defaults would
203 + # load into a wrong-shaped model with no error. So when we detect a flat
204 + # config we hoist the flat fields into ``text_config`` / ``vision_config``
205 + # and rename them to canonical names.
206 +
207 + # flat text field name -> canonical MuseGlimmerTextConfig field name
208 + _FLAT_TEXT_RENAMES = {
209 + "hidden_act": "hidden_activation",
210 + "output_soft_cap_temp": "final_logit_softcapping",
211 + }
212 + # canonical MuseGlimmerTextConfig constructor params that may appear flat
213 + _FLAT_TEXT_KEYS = frozenset(
214 + {
215 + "vocab_size",
216 + "hidden_size",
217 + "intermediate_size",
218 + "num_hidden_layers",
219 + "num_attention_heads",
220 + "num_key_value_heads",
221 + "head_dim",
222 + "hidden_activation",
223 + "hidden_act",
224 + "max_position_embeddings",
225 + "initializer_range",
226 + "rms_norm_eps",
227 + "use_cache",
228 + "tie_word_embeddings",
229 + "rope_parameters",
230 + "rope_theta",
231 + "attention_bias",
232 + "attention_dropout",
233 + "query_pre_attn_scalar",
234 + "sliding_window",
235 + "layer_types",
236 + "final_logit_softcapping",
237 + "output_soft_cap_temp",
238 + "attn_logit_softcapping",
239 + "use_bidirectional_attention",
240 + "qk_scale_factor",
241 + "use_qk_norm",
242 + "use_attn_output_gate",
243 + "output_multiplier",
244 + "normalize_tok_embeddings",
245 + "post_norm_eps",
246 + "no_rope_layers",
247 + }
248 + )
249 + # flat vision field name -> canonical MuseGlimmerVisionConfig field name
250 + _FLAT_VISION_RENAMES = {
251 + "vision_latent_dim": "hidden_size",
252 + "vision_heads": "num_attention_heads",
253 + "vision_layers": "num_hidden_layers",
254 + "vision_output_dim": "output_dim",
255 + "vision_patch_size": "patch_size",
256 + "vision_patch_temporal": "patch_temporal",
257 + "vision_adapter_dim": "adapter_dim",
258 + "vision_pos_emb_grid_h": "pos_emb_height",
259 + "vision_pos_emb_grid_w": "pos_emb_width",
260 + "vision_downsample_factor": "merge_kernel_size",
261 + }
262 +
263 + @classmethod
264 + def _looks_flat(cls, kwargs: dict) -> bool:
265 + # Flat if there is no explicit text_config but there ARE text-level
266 + # fields at the top level (e.g. hidden_size / num_hidden_layers).
267 + if kwargs.get("text_config") is not None:
268 + return False
269 + return any(k in kwargs for k in ("hidden_size", "num_hidden_layers"))
270 +
271 + @classmethod
272 + def _normalize_flat(cls, kwargs: dict) -> dict:
273 + kwargs = dict(kwargs)
274 + text: dict = {}
275 + for key in list(kwargs.keys()):
276 + if key in cls._FLAT_TEXT_KEYS:
277 + canon = cls._FLAT_TEXT_RENAMES.get(key, key)
278 + text.setdefault(canon, kwargs.pop(key))
279 +
280 + vision_mlp_ratio = kwargs.pop("vision_mlp_ratio", None)
281 + sparse_factor = kwargs.pop("vision_sparse_attention_factor", None)
282 + vision: dict = {}
283 + for key in list(kwargs.keys()):
284 + if key in cls._FLAT_VISION_RENAMES:
285 + vision.setdefault(cls._FLAT_VISION_RENAMES[key], kwargs.pop(key))
286 +
287 + if vision_mlp_ratio is not None:
288 + hidden_size = vision.get("hidden_size", 1536)
289 + vision["intermediate_size"] = int(vision_mlp_ratio * hidden_size)
290 + if sparse_factor is not None:
291 + sparse_factor = int(sparse_factor)
292 + if sparse_factor <= 0:
293 + raise ValueError("vision_sparse_attention_factor must be positive")
294 + num_layers = vision.get("num_hidden_layers", 50)
295 + vision["layer_types"] = [
296 + "full_attention"
297 + if (layer_idx + 1) % sparse_factor == 0 or layer_idx == num_layers - 1
298 + else "sliding_attention"
299 + for layer_idx in range(num_layers)
300 + ]
301 +
302 + if text:
303 + kwargs["text_config"] = text
304 + if vision:
305 + kwargs["vision_config"] = vision
306 + return kwargs
307 +
308 + def __init__(
309 + self,
310 + text_config: dict | MuseGlimmerTextConfig | None = None,
311 + vision_config: dict | MuseGlimmerVisionConfig | None = None,
312 + image_token_id: int = 200092,
313 + video_token_id: int = 200091,
314 + **kwargs,
315 + ) -> None:
316 + # Detect + fold a flat config into nested text/vision before building
317 + # the sub-configs. (image/video token ids are handled below; a flat
318 + # ``patch_token_id`` alias is mapped too.)
319 + if text_config is None and self._looks_flat(kwargs):
320 + kwargs = self._normalize_flat(kwargs)
321 + text_config = kwargs.pop("text_config", None)
322 + vision_config = kwargs.pop("vision_config", vision_config)
323 + # flat image-token alias
324 + if "patch_token_id" in kwargs and "image_token_id" not in kwargs:
325 + image_token_id = kwargs.pop("patch_token_id")
326 +
327 + if text_config is None:
328 + self.text_config = MuseGlimmerTextConfig()
329 + elif isinstance(text_config, dict):
330 + self.text_config = MuseGlimmerTextConfig(**text_config)
331 + else:
332 + self.text_config = text_config
333 +
334 + if vision_config is None:
335 + self.vision_config = MuseGlimmerVisionConfig()
336 + elif isinstance(vision_config, dict):
337 + self.vision_config = MuseGlimmerVisionConfig(**vision_config)
338 + else:
339 + self.vision_config = vision_config
340 +
341 + self.image_token_id = image_token_id
342 + self.video_token_id = video_token_id
343 + self.patch_token_id = image_token_id
344 + super().__init__(**kwargs)
345 +
346 +
347 +class MuseGlimmerAssistantConfig(Qwen3Config):
348 + """Config for the Muse Glimmer DFlash draft head.
349 +
350 + The head is Qwen3-shaped and runs on vLLM's generic ``qwen3_dflash``
351 + implementation, so this derives from ``Qwen3Config``. It cannot BE
352 + ``Qwen3Config``, because two of the checkpoint's values do not survive that
353 + class:
354 +
355 + * ``sliding_window`` is gated behind ``use_sliding_window``, which defaults
356 + to False -- ``Qwen3Config(sliding_window=2048).sliding_window`` is None.
357 + The checkpoint declares ``sliding_window: 2048`` and five
358 + ``sliding_attention`` layers, so the window silently disappears and the
359 + DFlash path then raises "sliding attention requires a window size".
360 + * ``vocab_size`` is absent from the checkpoint, so Qwen3's default of
361 + 151936 applies instead of Muse Glimmer's 202048. That one is *silent*: it
362 + builds an all-zero ``draft_id_to_target_id`` remap, and also puts
363 + pad/bos/eos/mask token ids out of range.
364 +
365 + Both defaults are set here so an unmodified checkpoint loads correctly.
366 + """
367 +
368 + model_type = "muse_glimmer_assistant"
369 +
370 + def __init__(
371 + self,
372 + vocab_size: int = 202048,
373 + use_sliding_window: bool = True,
374 + dflash_config: dict | None = None,
375 + **kwargs,
376 + ) -> None:
377 + super().__init__(
378 + vocab_size=vocab_size, use_sliding_window=use_sliding_window, **kwargs
379 + )
380 + # muse_glimmer_assistant uses non-causal attention.
381 + self.dflash_config = {"causal": False} | (dflash_config or {})
382 +
383 +
384 +__all__ = [
385 + "MuseGlimmerTextConfig",
386 + "MuseGlimmerVisionConfig",
387 + "MuseGlimmerConfig",
388 + "MuseGlimmerAssistantConfig",
389 +]
vllm/transformers_utils/processors/muse_glimmer.py new
+582
@@ -0,0 +1,582 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HuggingFace processor for the MuseGlimmer multimodal model (text + image + video).
4 +
5 +This packages MuseGlimmer's vision preprocessing and token-span expansion the standard
6 +HF way, so you can do:
7 +
8 + processor = AutoProcessor.from_pretrained(hf_dir, trust_remote_code=True)
9 + messages = [{"role": "user", "content": [
10 + {"type": "image"},
11 + {"type": "text", "text": "Describe this image."},
12 + ]}]
13 + text = processor.apply_chat_template(messages, add_generation_prompt=True)
14 + inputs = processor(text=text, images=[pil_image], return_tensors="pt")
15 + out = model.generate(**inputs, max_new_tokens=256)
16 +
17 +It reproduces the exact token layout the model expects:
18 +
19 + image -> <|image_start|> + <|patch|> * N + <|image_end|>
20 + video -> <|vid_start|>
21 + ( "Time: X.Xs" + <|video|> * P [+ <|vid_frame_separator|>] )*
22 + + <|vid_end|>
23 +
24 +where N = floor(H/(patch*ds)) * floor(W/(patch*ds)) (<= 4096) for images and
25 +P = the same grid count per temporal frame-group (<= 144) for video, with the
26 +target (H, W) chosen by MuseGlimmerImageProcessor.compute_image_size /
27 +MuseGlimmerVideoProcessor.compute_video_frame_size (the same grid logic as
28 +modeling_muse_glimmer.MuseGlimmerVisionEncoder).
29 +
30 +The chat template emits one sentinel per media item -- ``<|image|>`` for images,
31 +``<|video|>`` for videos -- which ``MuseGlimmerProcessor.__call__`` expands into the
32 +spans above using each media item's computed token count. Vision features are
33 +returned in ``pixel_values`` (a list of per-image / per-frame-group tensors) in
34 +the order the sentinels appear, matching how ``MuseGlimmerModel`` merges them at the
35 +``<|patch|>`` / ``<|video|>`` positions.
36 +"""
37 +
38 +from __future__ import annotations
39 +
40 +import itertools
41 +import math
42 +from pathlib import Path
43 +
44 +import torch
45 +from PIL import Image
46 +from torchvision import transforms as T
47 +from transformers.feature_extraction_utils import BatchFeature
48 +from transformers.image_processing_utils import BaseImageProcessor
49 +from transformers.processing_utils import ProcessorMixin
50 +from transformers.video_processing_utils import BaseVideoProcessor
51 +
52 +# Single-token sentinels emitted by the chat template, one per media item.
53 +IMAGE_SENTINEL = "<|image|>" # id 200090; absent from the image/video spans
54 +VIDEO_SENTINEL = "<|video|>" # id 200091; also the per-frame token (see __call__)
55 +
56 +# Jinja chat template (multimodal superset of convert_muse_glimmer_to_hf.CHAT_TEMPLATE).
57 +# Message content may be a plain string OR a list of parts
58 +# ({"type": "text"|"image"|"video", ...}); image/video parts render as a single
59 +# sentinel that __call__ expands.
60 +MUSE_GLIMMER_MM_CHAT_TEMPLATE = (
61 + "{{- bos_token -}}"
62 + "{%- macro render_parts(content) -%}"
63 + "{%- if content is string -%}{{- content -}}"
64 + "{%- else -%}"
65 + "{%- for part in content -%}"
66 + "{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}"
67 + "{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}"
68 + "{%- elif part['type'] == 'text' -%}{{- part['text'] -}}"
69 + "{%- endif -%}"
70 + "{%- endfor -%}"
71 + "{%- endif -%}"
72 + "{%- endmacro -%}"
73 + # At inference (add_generation_prompt), inject a default system message when
74 + # the caller supplied none. Full-transcript rendering
75 + # (add_generation_prompt=False) is left unchanged.
76 + "{%- set ns = namespace(has_system=false) -%}"
77 + "{%- for m in messages -%}"
78 + "{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}"
79 + "{%- endfor -%}"
80 + "{%- if add_generation_prompt and not ns.has_system -%}"
81 + "{{- '<|start|>system<|message|>You are a helpful assistant.<|eot|>' -}}"
82 + "{%- endif -%}"
83 + "{%- for message in messages -%}"
84 + "{%- set role = message['role'] -%}"
85 + "{%- if role == 'assistant' -%}"
86 + "{%- set recipient = message.get('recipient') -%}"
87 + "{%- set end_turn = message.get('end_turn') -%}"
88 + "{%- if end_turn is none -%}"
89 + "{%- set end_turn = not (recipient and recipient != 'user') -%}"
90 + "{%- endif -%}"
91 + "{{- '<|start|>assistant' -}}"
92 + "{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}"
93 + "{{- '<|message|>' -}}{{- render_parts(message['content']) -}}"
94 + "{{- ('<|eot|>' if end_turn else '<|eom|>') -}}"
95 + "{%- elif role == 'tool' -%}"
96 + "{%- set name = message.get('name', '') -%}"
97 + # Tool content is emitted as-is (string body or interleaved image/text parts).
98 + # Any <tool_output ...> wrapper is baked into the SFT data content; the
99 + # tokenizer does not add it, so the template must not either.
100 + "{{- '<|start|>tool ' + name + '<|message|>' -}}"
101 + "{{- render_parts(message['content']) -}}"
102 + "{{- '<|eot|>' -}}"
103 + "{%- else -%}"
104 + "{%- set header = role -%}"
105 + "{%- if message.get('name') -%}"
106 + "{%- set header = role + ' ' + message['name'] -%}"
107 + "{%- endif -%}"
108 + "{{- '<|start|>' + header + '<|message|>' -}}"
109 + "{{- render_parts(message['content']) -}}"
110 + "{{- '<|eot|>' -}}"
111 + "{%- endif -%}"
112 + "{%- endfor -%}"
113 + "{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}"
114 +)
115 +
116 +
117 +def _grid_size(
118 + img_w: int, img_h: int, patch_hw: int, max_tokens: int
119 +) -> tuple[int, int, int]:
120 + """Pick the integer (H, W) grid closest to the aspect ratio under the token cap.
121 +
122 + Replicates MuseGlimmerVisionEncoder._compute_grid_size
123 + (modeling_muse_glimmer.py) so the processor needs no torch model import.
124 + Returns (target_h, target_w, n_tokens).
125 + """
126 + i_nph = img_h / patch_hw
127 + i_npw = img_w / patch_hw
128 + ratio = i_npw / i_nph if i_nph > 0 else 1.0
129 + if i_nph * i_npw > max_tokens:
130 + i_nph = (max_tokens / ratio) ** 0.5
131 + i_npw = i_nph * ratio
132 + candidates = list(
133 + set(
134 + itertools.product(
135 + [math.floor(i_nph), math.ceil(i_nph)],
136 + [math.floor(i_npw), math.ceil(i_npw)],
137 + )
138 + )
139 + )
140 + candidates = [
141 + (nph, npw)
142 + for nph, npw in candidates
143 + if nph >= 1 and npw >= 1 and nph * npw <= max_tokens
144 + ]
145 + if not candidates:
146 + candidates = [(max(1, round(i_nph)), max(1, round(i_npw)))]
147 + nph, npw = min(candidates, key=lambda c: abs(c[0] / c[1] - img_h / img_w))
148 + return nph * patch_hw, npw * patch_hw, nph * npw
149 +
150 +
151 +class MuseGlimmerImageProcessor(BaseImageProcessor):
152 + """Resize + normalize MuseGlimmer images and compute patch-token counts.
153 +
154 + Variable-resolution: each image is resized to the grid that best matches its
155 + aspect ratio under the per-image token cap, then normalized with mean/std 0.5.
156 + Returns per-image tensors (not stacked) because MuseGlimmer consumes a list of
157 + variable-size images. Video frames are handled by MuseGlimmerVideoProcessor.
158 + """
159 +
160 + model_input_names = ["pixel_values"]
161 +
162 + def __init__(
163 + self,
164 + patch_size: int = 14,
165 + downsample_factor: int = 2,
166 + max_image_tokens: int = 4096,
167 + image_mean: float = 0.5,
168 + image_std: float = 0.5,
169 + **kwargs,
170 + ):
171 + super().__init__(**kwargs)
172 + self.patch_size = patch_size
173 + self.downsample_factor = downsample_factor
174 + self.max_image_tokens = max_image_tokens
175 + self.image_mean = image_mean
176 + self.image_std = image_std
177 +
178 + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor:
179 + # Functional normalize -- no stored Normalize object, so the processor
180 + # stays JSON-serializable for to_dict() / save_pretrained.
181 + return T.functional.normalize(
182 + T.functional.to_tensor(image),
183 + [self.image_mean] * 3,
184 + [self.image_std] * 3,
185 + )
186 +
187 + # -- size computation (mirrors modeling_muse_glimmer.MuseGlimmerVisionEncoder) --
188 + def compute_image_size(self, img_w: int, img_h: int) -> tuple[int, int, int]:
189 + ph = self.patch_size * self.downsample_factor
190 + return _grid_size(img_w, img_h, ph, self.max_image_tokens)
191 +
192 + # -- preprocessing ---------------------------------------------------------
193 + def preprocess_image(self, image: Image.Image) -> tuple[torch.Tensor, int]:
194 + """Return (pixel tensor [3, H, W], n_patch_tokens) for one image."""
195 + image = image.convert("RGB")
196 + target_h, target_w, n_tokens = self.compute_image_size(
197 + image.width, image.height
198 + )
199 + image = image.resize((target_w, target_h), Image.LANCZOS)
200 + return self._to_norm_tensor(image), n_tokens
201 +
202 +
203 +# torchcodec is the training-faithful video decoder; other decoders (torchvision /
204 +# PyAV) seek + color-convert differently and diverge from training. Optional at
205 +# import time so text/image paths load without it; decode_video raises if used.
206 +try:
207 + import torchcodec
208 +except Exception:
209 + torchcodec = None
210 +
211 +
212 +class MuseGlimmerVideoProcessor(BaseVideoProcessor):
213 + """MuseGlimmer video preprocessing behind the standard HF
214 + ``AutoVideoProcessor`` API.
215 +
216 + Exposes MuseGlimmer's training-faithful video handling as a first-class video
217 + processor so you can call ``processor(videos="clip.mp4")`` and downstream
218 + users can discover it via the standard HF ``AutoVideoProcessor`` API plus a
219 + ``video_preprocessor_config.json``. It wraps:
220 +
221 + * torchcodec decode (matches training; other decoders diverge),
222 + * uniform frame sampling to a whole multiple of ``patch_temporal``,
223 + * real per-group PTS timestamps (rendered as ``Time: X.Xs`` by
224 + MuseGlimmerProcessor),
225 + * ``patch_temporal`` frame-grouping (frames cat on the channel axis ->
226 + ``[patch_temporal * 3, H, W]``; the encoder detects video by channel count).
227 +
228 + It deliberately overrides ``preprocess`` instead of the BaseVideoProcessor fast
229 + pipeline (group_videos_by_shape / smart_resize / single stacked tensor): MuseGlimmer
230 + consumes a LIST of variable-size group tensors and needs the real per-group PTS,
231 + neither of which the stacked fast path models.
232 + """
233 +
234 + model_input_names = ["pixel_values"]
235 +
236 + def __init__(
237 + self,
238 + patch_size: int = 14,
239 + downsample_factor: int = 2,
240 + patch_temporal: int = 2,
241 + max_video_frame_tokens: int = 144,
242 + image_mean: float = 0.5,
243 + image_std: float = 0.5,
244 + video_num_frames: int = 96,
245 + video_sampling_fps: float = 2.0,
246 + **kwargs,
247 + ):
248 + super().__init__(**kwargs)
249 + self.patch_size = patch_size
250 + self.downsample_factor = downsample_factor
251 + self.patch_temporal = patch_temporal
252 + self.max_video_frame_tokens = max_video_frame_tokens
253 + self.image_mean = image_mean
254 + self.image_std = image_std
255 + self.video_num_frames = video_num_frames
256 + self.video_sampling_fps = video_sampling_fps
257 +
258 + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor:
259 + return T.functional.normalize(
260 + T.functional.to_tensor(image),
261 + [self.image_mean] * 3,
262 + [self.image_std] * 3,
263 + )
264 +
265 + def decode_video(self, video_path: str) -> tuple[list[Image.Image], list[float]]:
266 + """Sample frames + per-group PTS with torchcodec (the training decode path).
267 +
268 + ``timestamps[g]`` is the ACTUAL decoded PTS of the first frame in temporal
269 + group ``g``; ``len(frames)`` is a whole multiple of ``patch_temporal``.
270 + """
271 + if torchcodec is None:
272 + raise RuntimeError(
273 + "torchcodec is required for video decoding (it matches the training "
274 + "decode path). See the Environment Setup steps in hf/README.md."
275 + )
276 + pt = self.patch_temporal
277 + reader = torchcodec.decoders.VideoDecoder(video_path)
278 + total = len(reader)
279 + assert reader.metadata.average_fps is not None, (
280 + f"Video has no FPS metadata: {video_path}"
281 + )
282 + fps = reader.metadata.average_fps
283 + assert self.video_sampling_fps and self.video_sampling_fps > 0, (
284 + f"video_sampling_fps must be positive, got {self.video_sampling_fps}"
285 + )
286 + n = min(
287 + int(total * self.video_sampling_fps / fps), self.video_num_frames, total
288 + )
289 + n = max(pt, (n // pt) * pt)
290 + n = min(n, total)
291 + if n < pt:
292 + raise ValueError(
293 + f"Video has only {total} decodable frame(s) but needs at least "
294 + f"{pt} (one temporal patch): {video_path}"
295 + )
296 + indices = torch.linspace(0, total - 1, n).long().tolist()
297 + frames: list[Image.Image] = []
298 + timestamps: list[float] = []
299 + for j, i in enumerate(indices):
300 + fr = reader[i]
301 + frames.append(
302 + Image.fromarray(fr.data.permute(1, 2, 0).numpy()).convert("RGB")
303 + )
304 + if j % pt == 0:
305 + pts = getattr(fr, "pts_seconds", None)
306 + timestamps.append(float(pts) if pts is not None else i / fps)
307 + return frames, timestamps
308 +
309 + def compute_video_frame_size(self, img_w: int, img_h: int) -> tuple[int, int, int]:
310 + """Pick (H, W, tokens_per_frame) for a video frame under the per-frame cap.
311 +
312 + Mirrors MuseGlimmerImageProcessor.compute_image_size but uses the (smaller)
313 + per-frame video token budget; shares the same ``_grid_size`` grid logic as
314 + modeling_muse_glimmer.MuseGlimmerVisionEncoder.
315 + """
316 + ph = self.patch_size * self.downsample_factor
317 + return _grid_size(img_w, img_h, ph, self.max_video_frame_tokens)
318 +
319 + def _group_frames(
320 + self, frames: list[Image.Image]
321 + ) -> tuple[list[torch.Tensor], int, int]:
322 + """Resize/normalize frames and cat ``patch_temporal`` frames per group.
323 +
324 + Returns (group_tensors, n_groups, tokens_per_group).
325 + """
326 + pt = self.patch_temporal
327 + if pt <= 0:
328 + raise ValueError(f"patch_temporal must be positive, got {pt}")
329 + if not frames:
330 + raise ValueError("video must contain at least one frame")
331 + if padding := (-len(frames)) % pt:
332 + frames = [*frames, *([frames[-1]] * padding)]
333 + first = frames[0].convert("RGB")
334 + target_h, target_w, n_tokens = self.compute_video_frame_size(
335 + first.width, first.height
336 + )
337 + groups: list[torch.Tensor] = []
338 + for i in range(0, len(frames), pt):
339 + grp = [
340 + self._to_norm_tensor(
341 + frames[i + j]
342 + .convert("RGB")
343 + .resize((target_w, target_h), Image.LANCZOS)
344 + )
345 + for j in range(pt)
346 + ]
347 + groups.append(torch.cat(grp, dim=0))
348 + return groups, len(groups), n_tokens
349 +
350 + @staticmethod
351 + def _normalize_videos(videos) -> list:
352 + """Normalize the ``videos`` arg to a list of per-video items.
353 +
354 + A video item is a path (str/Path) or a list of PIL frames. Accepts a
355 + single video (one path, or one list of frames) or a list of those.
356 + """
357 + if isinstance(videos, (str, Path)):
358 + return [videos]
359 + if videos and not isinstance(videos[0], (list, tuple, str, Path)):
360 + return [videos] # a single list of PIL frames
361 + return list(videos or [])
362 +
363 + def preprocess_one(
364 + self,
365 + video: str | Path | list[Image.Image],
366 + timestamps: list[float] | None = None,
367 + ) -> tuple[list[torch.Tensor], int, int, list[float]]:
368 + """Decode (if a path) + group ONE video.
369 +
370 + Returns (group_tensors, n_groups, tokens_per_group, group_timestamps). This
371 + is the single per-video implementation shared by ``preprocess`` (the HF
372 + AutoVideoProcessor batch API) and ``MuseGlimmerProcessor.__call__`` (prompt
373 + building), so the two paths cannot drift. For a path, frames + real
374 + per-group PTS come from ``decode_video``; for pre-decoded frames the given
375 + ``timestamps`` are used verbatim (empty if not supplied).
376 + """
377 + if isinstance(video, (str, Path)):
378 + frames, ts = self.decode_video(str(video))
379 + else:
380 + frames = [f.convert("RGB") for f in video]
381 + ts = list(timestamps) if timestamps is not None else []
382 + groups, n_groups, tokens_per_group = self._group_frames(frames)
383 + return groups, n_groups, tokens_per_group, ts
384 +
385 + def preprocess(
386 + self,
387 + videos,
388 + video_timestamps: list[list[float]] | None = None,
389 + return_tensors: str | None = "pt",
390 + **kwargs,
391 + ) -> BatchFeature:
392 + """Preprocess video file path(s) or pre-decoded frame list(s).
393 +
394 + Accepts a single video (path str, or list of PIL frames) or a list of
395 + those. Returns a BatchFeature carrying, in video order:
396 + * ``pixel_values`` -- flat list of [pt*3, H, W] group tensors,
397 + * ``video_num_groups`` -- groups per video,
398 + * ``video_tokens_per_group`` -- patch-token count per group per video,
399 + * ``video_timestamps`` -- per-group PTS per video,
400 + i.e. exactly what MuseGlimmerProcessor needs to build the
401 + ``<|vid_start|> ( Time: X.Xs <|video|>*P [sep] )* <|vid_end|>`` block.
402 + """
403 + videos = self._normalize_videos(videos)
404 + pixel_values: list[torch.Tensor] = []
405 + num_groups: list[int] = []
406 + tokens_per_group: list[int] = []
407 + out_ts: list[list[float]] = []
408 + for idx, v in enumerate(videos):
409 + ts_in = video_timestamps[idx] if video_timestamps else None
410 + groups, ng, tpg, ts = self.preprocess_one(v, ts_in)
411 + pixel_values += groups
412 + num_groups.append(ng)
413 + tokens_per_group.append(tpg)
414 + out_ts.append(ts)
415 + batch = BatchFeature(
416 + data={
417 + "video_num_groups": num_groups,
418 + "video_tokens_per_group": tokens_per_group,
419 + "video_timestamps": out_ts,
420 + },
421 + tensor_type=None,
422 + )
423 + # Variable-size list -- BatchFeature would fail to stack into one tensor.
424 + batch["pixel_values"] = pixel_values
425 + return batch
426 +
427 +
428 +class MuseGlimmerProcessor(ProcessorMixin):
429 + """Bundle MuseGlimmerImageProcessor + MuseGlimmerVideoProcessor + tokenizer;
430 + expand media sentinels into spans.
431 +
432 + Images go through ``image_processor`` (MuseGlimmerImageProcessor) and videos through
433 + ``video_processor`` (MuseGlimmerVideoProcessor) -- one preprocessing implementation
434 + each, no overlap.
435 + """
436 +
437 + attributes = ["image_processor", "video_processor", "tokenizer"]
438 + image_processor_class = "AutoImageProcessor"
439 + video_processor_class = "AutoVideoProcessor"
440 + tokenizer_class = "PreTrainedTokenizerFast"
441 +
442 + def __init__(
443 + self,
444 + image_processor=None,
445 + video_processor=None,
446 + tokenizer=None,
447 + chat_template=None,
448 + **kwargs,
449 + ):
450 + if image_processor is None:
451 + image_processor = MuseGlimmerImageProcessor()
452 + if video_processor is None:
453 + video_processor = MuseGlimmerVideoProcessor()
454 + super().__init__(
455 + image_processor,
456 + video_processor,
457 + tokenizer,
458 + chat_template=chat_template or MUSE_GLIMMER_MM_CHAT_TEMPLATE,
459 + **kwargs,
460 + )
461 +
462 + def _sid(self, token: str) -> int:
463 + return self.tokenizer.convert_tokens_to_ids(token)
464 +
465 + def _image_block(self, n_tokens: int) -> list[int]:
466 + return (
467 + [self._sid("<|image_start|>")]
468 + + [self._sid("<|patch|>")] * n_tokens
469 + + [self._sid("<|image_end|>")]
470 + )
471 +
472 + def _video_block(
473 + self,
474 + n_groups: int,
475 + tokens_per_group: int,
476 + timestamps: list[float] | None = None,
477 + ) -> list[int]:
478 + """Per-group ``Time: X.Xs`` + <|video|>*P, separated/terminated.
479 +
480 + ``timestamps`` (one per temporal group, in seconds) should be the ACTUAL
481 + decoded frame times -- training renders the real factored frame PTS, not a
482 + uniform grid. Falls back to ``g*patch_temporal/fps`` only when timestamps
483 + are not supplied (e.g. callers that don't track frame times); that
484 + approximation can differ from training by the per-frame jitter.
485 + """
486 + vid = self._sid("<|video|>")
487 + sep = self._sid("<|vid_frame_separator|>")
488 + pt = self.video_processor.patch_temporal
489 + fps = self.video_processor.video_sampling_fps
490 + block = [self._sid("<|vid_start|>")]
491 + for g in range(n_groups):
492 + ts = timestamps[g] if timestamps is not None else g * pt / fps
493 + block += self.tokenizer.encode(f"Time: {ts:.1f}s", add_special_tokens=False)
494 + block += [vid] * tokens_per_group
495 + block.append(sep if g < n_groups - 1 else self._sid("<|vid_end|>"))
496 + return block
497 +
498 + def __call__(
499 + self,
500 + text: str | list[str] | None = None,
501 + images: list[Image.Image] | None = None,
502 + videos: list[list[Image.Image] | str] | None = None,
503 + video_timestamps: list[list[float]] | None = None,
504 + return_tensors: str | None = "pt",
505 + **kwargs,
506 + ) -> BatchFeature:
507 + if text is None:
508 + raise ValueError(
509 + "`text` is required (use apply_chat_template to build it)."
510 + )
511 + if isinstance(text, (list, tuple)):
512 + if len(text) != 1:
513 + raise ValueError(
514 + "MuseGlimmerProcessor supports a single text sample per call."
515 + )
516 + text = text[0]
517 +
518 + images = list(images or [])
519 + videos = list(videos or [])
520 + image_sentinel = self._sid(IMAGE_SENTINEL)
521 + video_sentinel = self._sid(VIDEO_SENTINEL)
522 +
523 + # Preprocess media up front so we can expand sentinels in document order.
524 + # Images via image_processor; videos via video_processor (each video item
525 + # is a path -- decoded here -- or a pre-decoded list of PIL frames).
526 + prepped_images = [self.image_processor.preprocess_image(im) for im in images]
527 + prepped_videos = [
528 + self.video_processor.preprocess_one(
529 + v, video_timestamps[i] if video_timestamps else None
530 + )
531 + for i, v in enumerate(videos)
532 + ]
533 +
534 + ids = self.tokenizer.encode(text, add_special_tokens=False)
535 + n_img = sum(1 for t in ids if t == image_sentinel)
536 + n_vid = sum(1 for t in ids if t == video_sentinel)
537 + if n_img != len(prepped_images):
538 + raise ValueError(
539 + f"{n_img} image sentinel(s) in text but "
540 + f"{len(prepped_images)} image(s) given."
541 + )
542 + if n_vid != len(prepped_videos):
543 + raise ValueError(
544 + f"{n_vid} video sentinel(s) in text but "
545 + f"{len(prepped_videos)} video(s) given."
546 + )
547 +
548 + out_ids: list[int] = []
549 + pixel_values: list[torch.Tensor] = []
550 + img_i = vid_i = 0
551 + for tid in ids:
552 + if tid == image_sentinel:
553 + tensor, n_tokens = prepped_images[img_i]
554 + img_i += 1
555 + out_ids += self._image_block(n_tokens)
556 + pixel_values.append(tensor)
557 + elif tid == video_sentinel:
558 + groups, n_groups, tokens_per_group, ts = prepped_videos[vid_i]
559 + vid_i += 1
560 + out_ids += self._video_block(n_groups, tokens_per_group, ts or None)
561 + pixel_values += groups
562 + else:
563 + out_ids.append(tid)
564 +
565 + data: dict = {
566 + "input_ids": [out_ids],
567 + "attention_mask": [[1] * len(out_ids)],
568 + }
569 + batch = BatchFeature(data=data, tensor_type=return_tensors)
570 + # Keep pixel_values as a list of variable-size tensors (MuseGlimmerModel
571 + # consumes a list); BatchFeature would fail to stack them into one tensor.
572 + if pixel_values:
573 + batch["pixel_values"] = pixel_values
574 + return batch
575 +
576 +
577 +__all__ = [
578 + "MuseGlimmerImageProcessor",
579 + "MuseGlimmerVideoProcessor",
580 + "MuseGlimmerProcessor",
581 + "MUSE_GLIMMER_MM_CHAT_TEMPLATE",
582 +]
vllm/v1/spec_decode/dflash.py
+15
@@ -80,6 +80,21 @@ class DFlashProposer(SpecDecodeBaseProposer):
80 self.draft_model_config.hf_config
81 )
82
83 + @override
84 + def load_model(self, target_model: torch.nn.Module) -> None:
85 + from vllm.model_executor.models.qwen3_dflash import (
86 + dflash_target_rope_is_neox_style,
87 + )
88 +
89 + draft_model_config = self.speculative_config.draft_model_config
90 + assert draft_model_config is not None
91 + # The drafter must rotate Q/K the way its target does. Take that from the
92 + # built target before super() constructs the draft.
93 + is_neox_style = dflash_target_rope_is_neox_style(target_model)
94 + if is_neox_style is not None:
95 + draft_model_config.hf_config.is_neox_style = is_neox_style
96 + super().load_model(target_model)
97 +
98 @override
99 def _create_draft_vllm_config(self) -> VllmConfig:
100 base = super()._create_draft_vllm_config()
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+13 -2
@@ -13,11 +13,19 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
13
14 def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module:
15 from vllm.compilation.backends import set_model_tag
16 - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
16 + from vllm.model_executor.models.qwen3_dflash import (
17 + dflash_has_any_non_causal,
18 + dflash_target_rope_is_neox_style,
19 + )
20
21 speculative_config = vllm_config.speculative_config
22 assert speculative_config is not None
23 draft_model_config = speculative_config.draft_model_config
24 + # The drafter must rotate Q/K the way its target does. Take that from the
25 + # built target before super() constructs the draft.
26 + is_neox_style = dflash_target_rope_is_neox_style(target_model)
27 + if is_neox_style is not None:
28 + draft_model_config.hf_config.is_neox_style = is_neox_style
29 # Select an attention backend that supports the drafter's attention: mixing
30 # a non-causal layer onto a causal-only backend would fail.
31 draft_vllm_config = replace(
@@ -46,7 +54,10 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
54 if hasattr(target_model, "get_language_model")
55 else target_model
56 )
49 - target_inner = target_language_model.model
57 + # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language
58 + # model, so get_language_model() already returns the inner module and has
59 + # no .model of its own.
60 + target_inner = getattr(target_language_model, "model", target_language_model)
61 draft_inner = dflash_model.model
62
63 # Skip embedding sharing under PP — each rank owns its own embedding.