diff --git a/examples/tool_chat_template_muse_glimmer.jinja b/examples/tool_chat_template_muse_glimmer.jinja new file mode 100644 index 0000000000..0ea2208934 --- /dev/null +++ b/examples/tool_chat_template_muse_glimmer.jinja @@ -0,0 +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 -%}{{- '\n\n' -}}{%- for k, v in args.items() -%}{{- '' -}}{%- 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 -%}{{- '\n' -}}{%- endfor -%}{{- '\n' -}}{%- 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 "" block like the following:\n' -}}{{- '\n\n$PARAMETER_VALUE\n...\n\n\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' -}}{{- '\n\n' -}}{{- 'value_1\n' -}}{{- 'This is the value for the second parameter\nthat can span\n"multiple" lines\n\n' -}}{{- '\n' -}}{%- 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|>\n' -}}{{- render_content(message['content']) -}}{{- '\n<|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 -%} diff --git a/tests/models/registry.py b/tests/models/registry.py index ed1786fdab..9db20e617e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1192,6 +1192,12 @@ _MULTIMODAL_EXAMPLE_MODELS = { # required by current PrefixLM implementation max_num_batched_tokens=31872, ), + "MuseGlimmerForConditionalGeneration": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + ), + "MuseGlimmerForCausalLM": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + ), "NVLM_D": _HfExamplesInfo("nvidia/NVLM-D-72B", trust_remote_code=True), "Llama_Nemotron_Nano_VL": _HfExamplesInfo( "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1", @@ -1489,6 +1495,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { max_num_seqs=32, min_transformers_version="4.56.3", # Required for Qwen3Next ), + "MuseGlimmerAssistantModel": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + speculative_model="meta-models/Muse-Glimmer-30B-assistant", + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env + max_num_seqs=32, + ), + "DFlashMuseGlimmerAssistantModel": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + speculative_model="meta-models/Muse-Glimmer-30B-assistant", + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env + max_num_seqs=32, + ), # [DSpark] "DSparkDraftModel": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Pro-DSpark", diff --git a/tests/models/utils.py b/tests/models/utils.py index f938a70223..6bd5431829 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -564,15 +564,20 @@ def dummy_hf_overrides( ) if hasattr(hf_config, "vision_config"): - hf_config.vision_config.update( + vision_config = hf_config.vision_config + vision_config.update( { "num_layers": 1, "num_hidden_layers": 1, } ) + # Keep per-layer metadata consistent with the reduced layer count. + if layer_types := getattr(vision_config, "layer_types", None): + vision_config.update({"layer_types": layer_types[:1]}) + if model_arch in ("Moondream3ForCausalLM", "HfMoondream"): - hf_config.vision_config.update({"enc_n_layers": 1}) + vision_config.update({"enc_n_layers": 1}) # e.g.: ibm-granite/granite-speech-3.3-2b if hasattr(hf_config, "encoder_config"): diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py new file mode 100644 index 0000000000..77f99a47ec --- /dev/null +++ b/tests/tool_use/test_muse_glimmer.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the MuseGlimmer ATEM tool parser and reasoning parser. + +MuseGlimmer writes every turn as a sequence of channel-scoped messages rather +than JSON, so the two parsers are tested together: the reasoning parser strips +the reasoning span and forwards the remaining channels as content, and the tool +parser reads ATEM markup out of those channels. + +Four areas, in order: + + 1. non-streaming tool-call extraction, including channel scoping (an + ```` echoed inside reasoning must never become a call); + 2. the reasoning -> tool-parser handoff, which regressed once by returning + ``content=None`` and starving the tool parser; + 3. streaming, where markers routinely straddle chunk boundaries, plus + truncation isolation for an unterminated ``to=self`` block; + 4. tool-name normalization against the tools registered on the request. + +These drive the parsers directly and need no checkpoint. The tests that require +a real tokenizer live in ``test_muse_glimmer_parse_delta.py``. +""" + +import json +from types import SimpleNamespace + +import pytest + +from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser +from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser + +R: MuseGlimmerReasoningParser +T: MuseGlimmerToolParser + + +@pytest.fixture(autouse=True) +def _fresh_parsers(): + """Give each test request-scoped parser state through the real constructors.""" + global R, T + R = MuseGlimmerReasoningParser(object()) + T = MuseGlimmerToolParser(object()) + + +# Any framing token that must NEVER appear in surfaced reasoning/content. +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "to=read.read", + "assistant to={name}<|message|>" + f'\n\n' + f'Paris\n' + f"\n" + ) + + +# ---------------------------------------------------------------- tool calls + + +def test_single_tool_call_after_reasoning(): + raw = ( + "to=self<|message|>Let me check the weather.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + 'celsius\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert out.tools_called and len(out.tool_calls) == 1 + assert out.tool_calls[0].function.name == "weather.get" + assert json.loads(out.tool_calls[0].function.arguments) == { + "city": "Paris", + "units": "celsius", + } + + +def test_parallel_calls_across_eom_boundaries(): + raw = ( + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n' + '2\n' + "\n<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n' + '4\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert out.tools_called and len(out.tool_calls) == 2, len(out.tool_calls) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"] + # JSON-typed values decode to ints + assert json.loads(out.tool_calls[0].function.arguments) == {"a": 1, "b": 2} + + +def test_echoed_invoke_in_reasoning_is_not_parsed(): + """Channel scoping: an invoke quoted inside reasoning is not a call.""" + raw = ( + 'to=self<|message|>I could call ' + '1 ' + "but I will not.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!" + assert out.content == "The answer is 42.", repr(out.content) + + +def test_plain_answer_yields_no_tool_calls(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, "to=user<|message|>Just a plain answer.<|eot|>", None + ) + assert not out.tools_called + + +def test_json_object_array_and_bool_params_decode(): + raw = ( + "<|start|>assistant to=api.call<|message|>" + '\n\n' + '{"nested": [1, 2, 3]}\n' + 'true\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert json.loads(out.tool_calls[0].function.arguments) == { + "payload": {"nested": [1, 2, 3]}, + "flag": True, + } + + +# ------------------------------------------------- reasoning -> tool handoff + + +def test_reasoning_to_toolcall_handoff(): + """The regression: content=None here starved the tool parser.""" + raw = ( + " to=self<|message|>Let me call the tool.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + "\n" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "Let me call the tool.", repr(reasoning) + assert content is not None and "thinking<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "thinking", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert not MuseGlimmerToolParser.extract_tool_calls(T, content, None).tools_called + + +def test_plain_content_without_framing_passes_through(): + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning( + R, "Just a direct answer.", None + ) + assert reasoning is None and content == "Just a direct answer.", ( + reasoning, + content, + ) + + +def test_reasoning_then_parallel_calls(): + raw = ( + " to=self<|message|>need two calls<|eom|>" + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n\n' + "<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n\n' + "<|eot|>" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "need two calls", repr(reasoning) + out = MuseGlimmerToolParser.extract_tool_calls(T, content, None) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"], ( + out.tool_calls + ) + + +# ----------------------------------------------------------------- streaming + + +def _stream(raw: str, chunk: int): + """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming + parsers; return (reasoning, content, tool_calls).""" + reasoning, content, toolcalls = [], [], [] + prev = "" + i = 0 + while i < len(raw): + cur = raw[: i + chunk] + delta = cur[len(prev) :] + dm = MuseGlimmerReasoningParser.extract_reasoning_streaming( + R, prev, cur, delta, [], [], [] + ) + if dm is not None: + if getattr(dm, "reasoning", None): + reasoning.append(dm.reasoning) + content_delta = getattr(dm, "content", None) + # Tool-channel content is an internal handoff to the tool parser, + # not client-visible content from the unified parser. + if content_delta and "" not in content_delta: + content.append(content_delta) + dt = MuseGlimmerToolParser.extract_tool_calls_streaming( + T, prev, cur, delta, [], [], [], _FakeReq() + ) + if dt is not None and dt.tool_calls: + toolcalls.extend(dt.tool_calls) + prev = cur + i += chunk + return "".join(reasoning), "".join(content), toolcalls + + +def _fn_of(tc): + fn = tc.function + if isinstance(fn, dict): + return fn.get("name"), fn.get("arguments") + return fn.name, fn.arguments + + +RAW_TOOLCALL = ( + " to=self<|message|>I should read the hostname file to answer.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n" +) + +RAW_ANSWER = ( + " to=self<|message|>Think about it.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" +) + +# NO closing <|eom|> -> truncated CoT +RAW_TRUNCATED = ( + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait" +) + + +def _check_toolcall_stream(chunk): + reasoning, content, tcs = _stream(RAW_TOOLCALL, chunk) + # (a) no framing token leaks into reasoning or content + for f in _FRAMING: + assert f not in reasoning, ( + f"framing {f!r} leaked into reasoning (chunk={chunk})" + ) + assert f not in content, f"framing {f!r} leaked into content (chunk={chunk})" + # (b) exactly one tool_call with correct name + args + assert len(tcs) == 1, f"expected 1 tool_call, got {len(tcs)} (chunk={chunk})" + name, args = _fn_of(tcs[0]) + assert name == "read.read", name + assert json.loads(args) == {"path": "/etc/hostname"}, args + assert tcs[0].index == 0 and tcs[0].type == "function" and tcs[0].id + # (c) reasoning captured separately and clean + assert reasoning == "I should read the hostname file to answer.", repr(reasoning) + assert content == "", repr(content) + + +def test_streaming_toolcall_chunk3(): + _check_toolcall_stream(3) + + +def test_streaming_toolcall_charwise(): + # worst case: markers arrive one char at a time (mid-marker deltas) + _check_toolcall_stream(1) + + +def test_streaming_toolcall_bigchunks(): + _check_toolcall_stream(17) + + +def test_streaming_reasoning_then_content(): + reasoning, content, tcs = _stream(RAW_ANSWER, 3) + for f in _FRAMING: + assert f not in reasoning and f not in content, f + assert reasoning == "Think about it.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tcs == [] + + +def test_truncated_cot_no_toolcall_nonstreaming(): + out = MuseGlimmerToolParser.extract_tool_calls(T, RAW_TRUNCATED, _FakeReq()) + assert not out.tools_called and out.tool_calls == [] + # partial reasoning must still be recovered by the reasoning parser + reasoning, _ = MuseGlimmerReasoningParser.extract_reasoning( + R, RAW_TRUNCATED, _FakeReq() + ) + assert reasoning and "Maybe I should call" in reasoning, repr(reasoning) + + +def test_truncated_cot_no_toolcall_streaming(): + _, _, tcs = _stream(RAW_TRUNCATED, 3) + assert tcs == [], f"truncated CoT invoke leaked as streaming tool call: {tcs}" + + +# ------------------------------------------------------- name normalization +# +# MuseGlimmer emits `get_weather.get_weather` for a bare-registered +# `get_weather`, and `weather.get` verbatim for a namespaced one. The parser +# normalizes against the tools actually registered on the request. + + +def test_doubled_bare_name_collapses(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), _req("get_weather") + ) + assert out.tools_called and out.tool_calls[0].function.name == "get_weather", ( + out.tool_calls[0].function.name + ) + + +def test_namespaced_name_preserved(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("weather.get"), _req("weather.get") + ) + assert out.tool_calls[0].function.name == "weather.get" + + +def test_unregistered_namespace_is_preserved(): + # Suffix-only matching can silently dispatch a tool from the wrong namespace. + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("foo.get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "foo.get_weather" + + +def test_trailing_segment_ambiguous_left_alone(): + # two registered tools share leaf 'get' -> ambiguous -> do NOT rewrite + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("x.get"), _req("weather.get", "time.get") + ) + assert out.tool_calls[0].function.name == "x.get" + + +def test_no_registered_tools_passthrough(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), None + ) + assert out.tool_calls[0].function.name == "get_weather.get_weather" + + +def test_exact_match_kept(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "get_weather" diff --git a/tests/tool_use/test_muse_glimmer_parse_delta.py b/tests/tool_use/test_muse_glimmer_parse_delta.py new file mode 100644 index 0000000000..805a46e13b --- /dev/null +++ b/tests/tool_use/test_muse_glimmer_parse_delta.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""parse_delta-level regression tests for the MuseGlimmer parsers. + +These drive the UNIFIED parser-engine streaming API the serving layer actually +uses (``DelegatingParser.parse_delta``), NOT the ``extract_*_streaming`` methods +directly. This is the coverage that was missing when a live no-tools streaming +request leaked raw harmony framing into ``content``: the phase machine marked +``reasoning_ended=True`` at the prompt boundary (because ``is_reasoning_end`` on +a prompt with no ``to=self`` wrongly returned True), skipping the reasoning +phase for the whole generation. + +Requires a real MuseGlimmer tokenizer; skipped if the checkpoint is unavailable. +""" + +import json +import os + +import pytest + +CKPT = os.environ.get("MUSE_GLIMMER_CKPT", "") + +pytestmark = pytest.mark.skipif( + not os.path.isdir(CKPT), reason=f"MuseGlimmer checkpoint not found at {CKPT}" +) + +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "Let me think step by step about the sum.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>", + ) + _assert_no_framing(content) + assert reasoning == "Let me think step by step about the sum.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tools == [] + + +def test_parse_delta_content_only(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=user<|message|>Just a direct answer.<|eot|>", + ) + _assert_no_framing(content) + assert content == "Just a direct answer.", repr(content) + assert tools == [] + + +def test_parse_delta_tool_call(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>I should read the hostname.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n", + ) + _assert_no_framing(content) + assert reasoning == "I should read the hostname.", repr(reasoning) + assert len(tools) == 1, tools + idx, name, args = tools[0] + assert idx == 0 and name == "read.read" + assert json.loads(args) == {"path": "/etc/hostname"} + + +def test_parse_delta_truncated_cot_no_toolcall(parser_cls, tok): + # Contemplated invoke inside an unterminated to=self block: no tool call, + # no framing leak into content, partial reasoning recovered. + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait", + ) + _assert_no_framing(content) + assert tools == [], f"contemplated invoke leaked as tool call: {tools}" + assert "Maybe I should call" in reasoning, repr(reasoning) + + +def test_parse_delta_reasoning_suppressed_when_not_requested(parser_cls, tok): + class _NoReasonReq: + tools = None + tool_choice = None + include_reasoning = False + + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>secret thoughts<|eom|>" + "<|start|>assistant to=user<|message|>Public answer.<|eot|>", + req=_NoReasonReq(), + ) + assert reasoning == "", repr(reasoning) # suppressed + assert content == "Public answer.", repr(content) + _assert_no_framing(content) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/transformers_utils/test_muse_glimmer_config.py b/tests/transformers_utils/test_muse_glimmer_config.py new file mode 100644 index 0000000000..ce9d4a819e --- /dev/null +++ b/tests/transformers_utils/test_muse_glimmer_config.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""MuseGlimmer config normalization. + +Two schema pairs have to converge on the same model, and each has produced a +silent-wrong-output bug: + + 1. flat (legacy converter) vs nested (canonical) ``config.json``. A FLAT + config once deserialized to an all-default text config, silently dropping + every checkpoint value. + 2. native vs modular attention config. The modular HF text_config OMITS + use_qk_norm / use_attn_output_gate (read as None) and ships a PRE-FOLDED + qk_scale_factor (43.784/sqrt(128)=3.87). Missing flags must read as True + (MuseGlimmer always applies QK-norm + output gate) and the query pre-scale + must normalize so native and modular land on the same number. +""" + +import math +from typing import Any + +from vllm.model_executor.models.muse_glimmer import ( + _muse_glimmer_query_prescale, + _muse_glimmer_use_attn_output_gate, + _muse_glimmer_use_qk_norm, +) +from vllm.transformers_utils.configs.muse_glimmer import MuseGlimmerConfig + +# A representative FLAT config (Ruan rl_v1/hf shape), trimmed. +FLAT: dict[str, Any] = { + "architectures": ["MuseGlimmerForCausalLM"], + "model_type": "muse_glimmer", + "has_vision": True, + "bos_token_id": 200000, + "eos_token_id": 200001, + "vocab_size": 202048, + "hidden_size": 6656, + "intermediate_size": 19968, + "num_hidden_layers": 52, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "hidden_act": "silu", + "max_position_embeddings": 16384, + "rms_norm_eps": 1e-5, + "post_norm_eps": 1e-8, + "qk_scale_factor": 43.7840518911, + "use_qk_norm": True, + "use_attn_output_gate": True, + "output_multiplier": 0.19611613513818404, + "output_soft_cap_temp": 20.0, + "normalize_tok_embeddings": True, + "rope_theta": 500000.0, + "sliding_window": 2048, + "patch_token_id": 200092, + "vision_latent_dim": 1536, + "vision_heads": 16, + "vision_layers": 50, + "vision_output_dim": 6144, + "vision_patch_size": 14, + "vision_patch_temporal": 2, + "vision_adapter_dim": 4096, + "vision_pos_emb_grid_h": 32, + "vision_pos_emb_grid_w": 32, +} + +NESTED: dict[str, Any] = { + "architectures": ["MuseGlimmerForCausalLM"], + "model_type": "muse_glimmer", + "image_token_id": 200092, + "text_config": { + "model_type": "muse_glimmer_text", + "vocab_size": 202048, + "hidden_size": 6656, + "num_hidden_layers": 52, + "hidden_activation": "silu", + "final_logit_softcapping": 20.0, + "qk_scale_factor": 43.7840518911, + "rope_parameters": {"rope_type": "default", "rope_theta": 500000.0}, + }, + "vision_config": { + "model_type": "muse_glimmer_vision", + "hidden_size": 1536, + "num_hidden_layers": 50, + }, +} + +HEAD_DIM = 128 +SQRT_HD = math.sqrt(HEAD_DIM) +NATIVE = 43.7840518911 +FOLDED = NATIVE / SQRT_HD # 3.8700... + + +class Cfg: + """Attention-config stand-in carrying only the fields under test.""" + + def __init__(self, **kw): + self.head_dim = HEAD_DIM + for k, v in kw.items(): + setattr(self, k, v) + + +# ------------------------------------------------------------ flat vs nested + + +def test_flat_config_values_respected(): + c = MuseGlimmerConfig(**FLAT) + t = c.text_config + assert t.hidden_size == 6656 + assert t.num_hidden_layers == 52 + assert t.vocab_size == 202048 + assert t.head_dim == 128 + assert t.hidden_activation == "silu" # renamed from hidden_act + assert t.final_logit_softcapping == 20.0 # renamed from output_soft_cap_temp + assert abs(t.output_multiplier - 0.19611613513818404) < 1e-12 + assert abs(t.qk_scale_factor - 43.7840518911) < 1e-9 + assert t.rope_parameters["rope_theta"] == 500000.0 # from flat rope_theta + # vision hoisted + renamed + assert c.vision_config.hidden_size == 1536 + assert c.vision_config.num_hidden_layers == 50 + assert c.vision_config.output_dim == 6144 + # flat patch_token_id -> image_token_id + assert c.image_token_id == 200092 + + +def test_flat_config_no_silent_default(): + # The regression: a non-default value MUST be honored, not silently dropped. + flat = dict(FLAT) + flat["hidden_size"] = 4096 + flat["num_hidden_layers"] = 40 + c = MuseGlimmerConfig(**flat) + assert c.text_config.hidden_size == 4096, "flat hidden_size silently ignored!" + assert c.text_config.num_hidden_layers == 40, "flat num_hidden_layers ignored!" + + +def test_nested_config_unchanged(): + c = MuseGlimmerConfig(**NESTED) + assert c.text_config.hidden_size == 6656 + assert c.text_config.hidden_activation == "silu" + assert c.text_config.final_logit_softcapping == 20.0 + assert c.vision_config.hidden_size == 1536 + assert c.vision_config.num_hidden_layers == 50 + assert c.image_token_id == 200092 + + +# -------------------------------------------------------- native vs modular + + +def test_qk_norm_missing_defaults_true(): + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=None)) is True # modular + assert _muse_glimmer_use_qk_norm(Cfg()) is True # absent + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=True)) is True # native + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=False)) is False # explicit off + + +def test_output_gate_missing_defaults_true(): + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=None)) is True + assert _muse_glimmer_use_attn_output_gate(Cfg()) is True + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=False)) is False + + +def test_query_prescale_native_and_modular_converge(): + # Both schemas must yield the SAME final scale_query_by (~3.87). + assert ( + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=NATIVE)) - FOLDED) < 1e-9 + ) + assert ( + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=FOLDED)) - FOLDED) < 1e-9 + ) + + +def test_query_prescale_explicit_wins(): + c = Cfg(scale_query_by=FOLDED, qk_scale_factor=NATIVE) + assert abs(_muse_glimmer_query_prescale(c) - FOLDED) < 1e-9 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 65d8e88d0a..2b93113b7e 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -943,7 +943,11 @@ class SpeculativeConfig: self.method = "eagle" elif "eagle3" in self.draft_model_config.model.lower(): self.method = "eagle3" - elif "dflash" in self.draft_model_config.model.lower(): + elif ( + "dflash" in self.draft_model_config.model.lower() + or "MuseGlimmerAssistantModel" + in self.draft_model_config.architectures + ): self.method = "dflash" elif ( "dspark" in self.draft_model_config.model.lower() diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 38fca287d5..88f0a2dc6f 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1567,13 +1567,15 @@ class SupportsEagle3(SupportsEagleBase, Protocol): parent_ref = self.get_language_model() elif hasattr(self, "language_model"): parent_ref = self.language_model - assert hasattr(parent_ref, "model"), ( - "Model instance must have 'model' attribute to set number of layers" - ) - assert isinstance(parent_ref.model, EagleModelMixin), ( + # A multimodal model that builds its decoder inside + # `_mark_language_model` has get_language_model() return the inner + # decoder itself, which IS the EagleModelMixin and has no further + # `.model`. Unwrap only when there is something to unwrap. + holder = getattr(parent_ref, "model", parent_ref) + assert isinstance(holder, EagleModelMixin), ( "Model instance must inherit from EagleModelMixin to set auxiliary layers" ) - parent_ref.model._set_aux_hidden_state_layers(layers) + holder._set_aux_hidden_state_layers(layers) def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: """ @@ -1590,13 +1592,12 @@ class SupportsEagle3(SupportsEagleBase, Protocol): parent_ref = self.get_language_model() elif hasattr(self, "language_model"): parent_ref = self.language_model - assert hasattr(parent_ref, "model"), ( - "Model instance must have 'model' attribute to get number of layers" - ) - assert hasattr(parent_ref.model, "layers"), ( + # Same unwrap-only-if-needed rule as set_aux_hidden_state_layers. + holder = getattr(parent_ref, "model", parent_ref) + assert hasattr(holder, "layers"), ( "Model instance must have 'layers' attribute to get number of layers" ) - num_layers = len(parent_ref.model.layers) + num_layers = len(holder.layers) return (2, num_layers // 2, num_layers - 3) diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py new file mode 100644 index 0000000000..74e0001318 --- /dev/null +++ b/vllm/model_executor/models/muse_glimmer.py @@ -0,0 +1,1649 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MuseGlimmer multimodal model for vLLM. + +Native port of the MuseGlimmer text decoder (``MuseGlimmerForCausalLM``). The text +stack is a Gemma2 derivative with the +following MuseGlimmer-specific deltas, each of which is handled explicitly here: + + * SiLU-gated MLP (``hidden_activation="silu"``), not Gemma's gelu-tanh. + * Scaleless RMSNorm on the token embeddings (no sqrt(hidden) scaling). + * Per-layer sandwich RMSNorms with a baked ``+1`` weight offset + (``x * (1 + w)``), matching Gemma, but with distinct eps for the + pre/post norms (``rms_norm_eps`` vs ``post_norm_eps``). + * QK-norm (weightless, fp32) applied *before* RoPE, followed by a query + pre-scale of ``qk_scale_factor / sqrt(head_dim)``. + * A per-head sigmoid attention output gate. + * iRoPE layout: NoPE layers use full attention, RoPE layers use sliding + window attention. RoPE is applied NEOX-style (``is_neox_style=True``): + the HF converter (``convert_muse_glimmer_weights_to_hf.py``, 20260806+) permutes + q/k into the half-split (NEOX) layout via ``_permute_for_rope`` so they + pair with ``rotate_half`` — matching the reference's interleaved rotation + on the *native* (unpermuted) weights. Serving the permuted HF weights with + ``is_neox_style=False`` scrambles q/k and causes token-repetition collapse. + * Final logits are pre-scaled by ``output_multiplier`` and then tanh + soft-capped at ``final_logit_softcapping``. + * Untied lm_head. + +The vision path supports variable-resolution images and temporally patched +videos. It mirrors the checkpoint's native vision encoder, including sparse +block attention, 2-D RoPE, pixel-shuffle downsampling, and the two-layer +adapter/projection stack. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from itertools import islice +from typing import Annotated, Literal + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torch import nn +from transformers import BatchFeature + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.distributed import ( + divide, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention, MMEncoderAttention +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.renderers import TokenizeParams +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processors.muse_glimmer import MuseGlimmerProcessor +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .interfaces import ( + EagleModelMixin, + MultiModalEmbeddings, + SupportsEagle3, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from .vision import get_fp8_padded_hidden_size, is_vit_use_data_parallel + +logger = init_logger(__name__) + + +def _text_config(config): + """MuseGlimmer checkpoints may nest the text config under ``text_config`` + (multimodal ``MuseGlimmerConfig``) or expose it directly + (``MuseGlimmerTextConfig``).""" + return getattr(config, "text_config", config) + + +def _vision_config(config): + return getattr(config, "vision_config", config) + + +def _muse_glimmer_has_vision(config) -> bool: + has_vision = getattr(config, "has_vision", None) + if has_vision is not None: + return bool(has_vision) + return hasattr(config, "vision_config") + + +IMAGE_TOKEN = "<|patch|>" +IMAGE_PROCESSOR_TOKEN = "<|image|>" +VIDEO_TOKEN = "<|video|>" + + +class MuseGlimmerImagePixelInputs(TensorSchema): + """Batched variable-resolution image inputs.""" + + type: Literal["image_pixels"] + pixel_values: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}), + ] + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] + + +class MuseGlimmerVideoPixelInputs(TensorSchema): + """Batched variable-length, variable-resolution video inputs.""" + + type: Literal["video_pixels"] + pixel_values: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape( + "bn", + "ng", + "c", + "h", + "w", + dynamic_dims={"ng", "h", "w"}, + ), + ] + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] + + +class MuseGlimmerProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + if not _muse_glimmer_has_vision(self.get_hf_config()): + return {} + return {"image": None, "video": None} + + def get_default_tok_params(self) -> TokenizeParams: + return super().get_default_tok_params().with_kwargs(add_special_tokens=False) + + def get_hf_processor(self, **kwargs: object) -> MuseGlimmerProcessor: + return self.ctx.init_processor( + MuseGlimmerProcessor, + tokenizer=self.get_tokenizer(), + **kwargs, + ) + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + processor = self.get_hf_processor() + image_processor = processor.image_processor + video_processor = processor.video_processor + num_frames = self.get_num_frames_with_most_features(seq_len, mm_counts) + return { + "image": int(image_processor.max_image_tokens), + "video": ( + num_frames + // int(video_processor.patch_temporal) + * int(video_processor.max_video_frame_tokens) + ), + } + + def get_image_size_with_most_features(self) -> ImageSize: + image_processor = self.get_hf_processor().image_processor + grid_size = math.isqrt(int(image_processor.max_image_tokens)) + side = ( + int(image_processor.patch_size) + * int(image_processor.downsample_factor) + * grid_size + ) + return ImageSize(width=side, height=side) + + def get_num_frames_with_most_features( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> int: + video_processor = self.get_hf_processor().video_processor + num_videos = max(mm_counts.get("video", 0), 1) + groups = max( + 1, + seq_len // num_videos // int(video_processor.max_video_frame_tokens), + ) + patch_temporal = int(video_processor.patch_temporal) + max_groups = max( + 1, + int(video_processor.video_num_frames) // patch_temporal, + ) + return min(groups, max_groups) * patch_temporal + + +class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get( + "video", 0 + ) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + processor = self.info.get_hf_processor() + video_processor = processor.video_processor + image_width, image_height = self.info.get_image_size_with_most_features() + video_grid = math.isqrt(int(video_processor.max_video_frame_tokens)) + video_size = ( + int(video_processor.patch_size) + * int(video_processor.downsample_factor) + * video_grid + ) + return { + "image": self._get_dummy_images( + width=image_width, + height=image_height, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), + ), + "video": self._get_dummy_videos( + width=video_size, + height=video_size, + num_frames=self.info.get_num_frames_with_most_features( + seq_len, mm_counts + ), + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), + ), + } + + +class MuseGlimmerMultiModalProcessor( + BaseMultiModalProcessor[MuseGlimmerProcessingInfo] +): + def _apply_hf_processor_text_only( + self, + prompt_text: str, + tokenization_kwargs: Mapping[str, object], + ) -> list[int]: + tokenizer = self.info.get_tokenizer() + return tokenizer.encode( + prompt_text, + **{"add_special_tokens": False, **tokenization_kwargs}, + ) + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + processor = self.info.get_hf_processor(**mm_kwargs) + tokenizer = processor.tokenizer + config = self.info.get_hf_config() + images = mm_data.get("images", ()) + videos = mm_data.get("videos", ()) + if not isinstance(images, Sequence) or not isinstance(videos, Sequence): + raise TypeError("MuseGlimmer multi-modal data must be a sequence") + + prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) + image_sentinel_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN) + video_sentinel_id = tokenizer.convert_tokens_to_ids(VIDEO_TOKEN) + if prompt_ids.count(image_sentinel_id) != len(images): + raise ValueError("The number of image placeholders does not match images") + if prompt_ids.count(video_sentinel_id) != len(videos): + raise ValueError("The number of video placeholders does not match videos") + + image_pixels: list[torch.Tensor] = [] + image_sizes: list[int] = [] + image_prompt_blocks: list[list[int]] = [] + for image in images: + output = processor( + text=IMAGE_PROCESSOR_TOKEN, + images=[image], + return_tensors=None, + ) + block_ids = output["input_ids"][0] + if isinstance(block_ids, torch.Tensor): + block_ids = block_ids.tolist() + pixels = output["pixel_values"] + if isinstance(pixels, torch.Tensor): + pixels = [pixels] if pixels.ndim == 3 else list(pixels) + if len(pixels) != 1: + raise ValueError( + "MuseGlimmer HF processor must return one tensor per image" + ) + image_pixels.append(pixels[0]) + image_sizes.append(block_ids.count(config.image_token_id)) + image_prompt_blocks.append(block_ids) + + video_pixels: list[torch.Tensor] = [] + video_sizes: list[int] = [] + video_prompt_blocks: list[list[int]] = [] + for video in videos: + if not isinstance(video, np.ndarray): + raise TypeError("MuseGlimmer video input must be a NumPy array") + frames = [Image.fromarray(frame).convert("RGB") for frame in video] + output = processor( + text=VIDEO_TOKEN, + videos=[frames], + return_tensors=None, + ) + block_ids = output["input_ids"][0] + if isinstance(block_ids, torch.Tensor): + block_ids = block_ids.tolist() + pixels = output["pixel_values"] + if isinstance(pixels, torch.Tensor): + pixels = list(pixels) + video_pixels.append(torch.stack(pixels)) + video_sizes.append(block_ids.count(config.video_token_id)) + video_prompt_blocks.append(block_ids) + + combined_prompt_ids: list[int] = [] + image_idx = video_idx = 0 + for token_id in prompt_ids: + if token_id == image_sentinel_id: + combined_prompt_ids.extend(image_prompt_blocks[image_idx]) + image_idx += 1 + elif token_id == video_sentinel_id: + combined_prompt_ids.extend(video_prompt_blocks[video_idx]) + video_idx += 1 + else: + combined_prompt_ids.append(token_id) + + data: dict[str, object] = { + "input_ids": [combined_prompt_ids], + } + if image_pixels: + data.update( + image_pixel_values=image_pixels, + image_feature_sizes=torch.tensor(image_sizes), + ) + if video_pixels: + data.update( + video_pixel_values=video_pixels, + video_feature_sizes=torch.tensor(video_sizes), + ) + del tok_kwargs + return BatchFeature(data=data, tensor_type=None) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields = {} + if "image_pixel_values" in hf_inputs: + fields.update( + image_pixel_values=MultiModalFieldConfig.batched("image"), + image_feature_sizes=MultiModalFieldConfig.batched("image"), + ) + if "video_pixel_values" in hf_inputs: + fields.update( + video_pixel_values=MultiModalFieldConfig.batched("video"), + video_feature_sizes=MultiModalFieldConfig.batched("video"), + ) + return fields + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + config = self.info.get_hf_config() + tokenizer = self.info.get_tokenizer() + vocab = tokenizer.get_vocab() + + placeholder = { + "image": vocab[IMAGE_TOKEN], + "video": vocab[VIDEO_TOKEN], + } + image_start_id = vocab["<|image_start|>"] + image_end_id = vocab["<|image_end|>"] + video_start_id = vocab["<|vid_start|>"] + video_end_id = vocab["<|vid_end|>"] + video_separator_id = vocab["<|vid_frame_separator|>"] + + def image_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]: + out_item = out_mm_kwargs["image"][item_idx] + num_tokens = int(out_item["image_feature_sizes"].data) + replacement = ( + [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id] + ) + return PromptUpdateDetails.select_token_id( + replacement, config.image_token_id + ) + + def video_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]: + out_item = out_mm_kwargs["video"][item_idx] + pixel_values = out_item["video_pixel_values"].data + num_groups = len(pixel_values) + if num_groups == 0: + raise ValueError( + "MuseGlimmer video must contain at least one frame group" + ) + + num_tokens = int(out_item["video_feature_sizes"].data) + tokens_per_group, remainder = divmod(num_tokens, num_groups) + if remainder: + raise ValueError( + "MuseGlimmer video feature size must be divisible by frame groups" + ) + + video_processor = processor.video_processor + patch_temporal = int(video_processor.patch_temporal) + sampling_fps = float(video_processor.video_sampling_fps) + if sampling_fps <= 0: + raise ValueError("MuseGlimmer video_sampling_fps must be positive") + + replacement = [video_start_id] + for group_idx in range(num_groups): + timestamp = group_idx * patch_temporal / sampling_fps + replacement.extend( + tokenizer.encode( + f"Time: {timestamp:.1f}s", add_special_tokens=False + ) + ) + replacement.extend([config.video_token_id] * tokens_per_group) + replacement.append( + video_separator_id if group_idx < num_groups - 1 else video_end_id + ) + return PromptUpdateDetails.select_token_id( + replacement, config.video_token_id + ) + + return [ + PromptReplacement( + modality="image", + target=[placeholder["image"]], + replacement=image_replacement, + ), + PromptReplacement( + modality="video", + target=[placeholder["video"]], + replacement=video_replacement, + ), + ] + + +def _muse_glimmer_use_qk_norm(config) -> bool: + """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF + ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``). + Treat a missing/None flag as True — only an explicit ``False`` disables it.""" + val = getattr(config, "use_qk_norm", None) + return True if val is None else bool(val) + + +def _muse_glimmer_use_attn_output_gate(config) -> bool: + """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS + applies it; the modular HF ``text_config`` omits ``use_attn_output_gate`` + (reads as ``None``). Missing/None -> True; only explicit ``False`` disables.""" + val = getattr(config, "use_attn_output_gate", None) + return True if val is None else bool(val) + + +def _muse_glimmer_query_prescale(config) -> float: + """Post-QK-norm query pre-scale (``scale_query_by``), normalized across the + two config schemas so the net query scaling matches the native reference. + + HF native modeling computes ``scale_query_by = qk_scale_factor / sqrt(head_dim)`` + where the NATIVE ``qk_scale_factor`` is the raw ``params.json`` value + (~43.784). The modular HF ``text_config`` PRE-FOLDS the ``1/sqrt(head_dim)`` + factor and ships ``qk_scale_factor = 43.784 / sqrt(128) = 3.87`` already, + expecting it applied directly. Both must yield the SAME ``scale_query_by`` + (~3.87), then softmax uses ``scaling = head_dim**-0.5``. + + Precedence: + 1. explicit ``scale_query_by`` (already the final factor) -> use as-is. + 2. else derive from ``qk_scale_factor``: + - if it is already the folded value (``~= qk_scale_factor/sqrt(hd)`` is + NOT what we want; detect the native form and divide) — we decide by + magnitude: the native raw value is ``folded * sqrt(head_dim)``. If + ``qk_scale_factor`` is close to ``folded_expected * sqrt(hd)`` we treat + it as native and divide; otherwise it is already folded, use directly. + """ + head_dim = config.head_dim + sqrt_hd = head_dim**0.5 + + explicit = getattr(config, "scale_query_by", None) + if explicit is not None: + return float(explicit) + + qk_scale = getattr(config, "qk_scale_factor", None) + if qk_scale is None: + # No scale info at all: fall back to the plain 1/sqrt(head_dim) identity + # (net query scaling then just the softmax scaling); should not happen + # for real MuseGlimmer checkpoints, which always carry qk_scale_factor. + return 1.0 + + qk_scale = float(qk_scale) + # Disambiguate native (raw, ~43.78) vs modular (folded, ~3.87). The native + # form, when divided by sqrt(head_dim), yields the folded target; the folded + # form is already the target. Native values are ~sqrt(head_dim)x larger than + # folded. Use a threshold at sqrt(head_dim) (with margin): if qk_scale is + # comparable to or larger than sqrt(head_dim), it is the native raw value and + # must be divided; otherwise it is already folded and used directly. + # head_dim=128 -> sqrt=11.31; native 43.78 > 11.31 (divide -> 3.87), + # folded 3.87 < 11.31 (use as-is). + if qk_scale >= sqrt_hd: + return qk_scale / sqrt_hd + return qk_scale + + +class MuseGlimmerRMSNorm(nn.Module): + """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end). + + ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the + input dtype. When ``with_scale`` is False the layer is weightless (used for + QK-norm and the token-embedding norm). + """ + + def __init__( + self, + dim: int | None = None, + eps: float = 1e-6, + with_scale: bool = True, + weight_offset: int = 0, + ) -> None: + super().__init__() + self.eps = eps + self.with_scale = with_scale + self.weight_offset = weight_offset + if with_scale: + assert dim is not None + self.weight = nn.Parameter(torch.zeros(dim)) + else: + self.register_parameter("weight", None) + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + out = self._norm(hidden_states.float()) + if self.with_scale: + out = out * (self.weight.float() + self.weight_offset) + return out.type_as(hidden_states) + + +class MuseGlimmerVisionAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + prefix: str = "", + ) -> None: + super().__init__() + if hidden_size % num_heads: + raise ValueError("MuseGlimmer vision hidden size must divide num heads") + + self.total_num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.use_data_parallel = is_vit_use_data_parallel(num_heads) + self.tp_size = ( + 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() + ) + self.num_heads = divide(num_heads, self.tp_size) + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=num_heads, + total_num_kv_heads=num_heads, + bias=True, + prefix=f"{prefix}.qkv_proj", + disable_tp=self.use_data_parallel, + ) + self.o_proj = RowParallelLinear( + input_size=hidden_size, + output_size=hidden_size, + bias=True, + prefix=f"{prefix}.o_proj", + disable_tp=self.use_data_parallel, + ) + self.attn = MMEncoderAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + scale=self.head_dim**-0.5, + prefix=f"{prefix}.attn", + ) + self.apply_rotary_emb = ApplyRotaryEmb( + enforce_enable=True, + is_neox_style=True, + enable_fp32_compute=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + sequence_lengths: torch.Tensor | None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + qkv = qkv.view( + hidden_states.shape[0], + 3, + self.num_heads, + self.head_dim, + ) + query, key, value = qkv.unbind(1) + query, key = self.apply_rotary_emb( + torch.stack([query, key]).contiguous(), + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ).unbind(0) + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + + output = self.attn( + query=query, + key=key, + value=value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + output = output.reshape(hidden_states.shape[0], -1) + output, _ = self.o_proj(output) + return output + + +class MuseGlimmerVisionMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True) + self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.c_proj(F.gelu(self.c_fc(hidden_states))) + + +class MuseGlimmerVisionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + intermediate_size: int, + prefix: str = "", + ) -> None: + super().__init__() + self.ln_1 = nn.LayerNorm(hidden_size) + self.attn = MuseGlimmerVisionAttention( + hidden_size, + num_heads, + prefix=f"{prefix}.attn", + ) + self.ln_2 = nn.LayerNorm(hidden_size) + self.mlp = MuseGlimmerVisionMLP(hidden_size, intermediate_size) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + sequence_lengths: torch.Tensor | None, + ) -> torch.Tensor: + batch_size, seq_len, hidden_size = hidden_states.shape + flattened = hidden_states.view(batch_size * seq_len, hidden_size) + flattened = flattened + self.attn( + self.ln_1(flattened), + rotary_pos_emb_cos, + rotary_pos_emb_sin, + cu_seqlens, + max_seqlen, + sequence_lengths, + ) + flattened = flattened + self.mlp(self.ln_2(flattened)) + return flattened.view(batch_size, seq_len, hidden_size) + + +class MuseGlimmerVisionEncoder(nn.Module): + def __init__(self, config, prefix: str = "") -> None: + super().__init__() + hidden_size = config.hidden_size + patch_dim = config.patch_temporal * 3 * config.patch_size**2 + self.hidden_size = hidden_size + self.patch_size = config.patch_size + self.patch_temporal = config.patch_temporal + self.merge_kernel_size = config.merge_kernel_size + self.pos_emb_height = config.pos_emb_height + self.pos_emb_width = config.pos_emb_width + self.head_dim = hidden_size // config.num_attention_heads + self.layer_types = list(config.layer_types) + if len(self.layer_types) != config.num_hidden_layers: + raise ValueError( + "MuseGlimmer vision layer_types must match num_hidden_layers" + ) + if self.head_dim % 4: + raise ValueError("MuseGlimmer vision head dimension must be divisible by 4") + + self.conv1_linear = nn.Linear(patch_dim, hidden_size, bias=False) + self.positional_embedding_vlm = nn.Parameter( + torch.zeros(config.pos_emb_height * config.pos_emb_width, hidden_size) + ) + self.ln_pre = nn.LayerNorm(hidden_size) + self.transformer = nn.ModuleList( + [ + MuseGlimmerVisionBlock( + hidden_size, + config.num_attention_heads, + config.intermediate_size, + prefix=f"{prefix}.transformer.{layer_idx}", + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + vision_attention = self.transformer[0].attn + self.tp_size = vision_attention.tp_size + self.attn_backend = vision_attention.attn.attn_backend + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size( + config.num_attention_heads, self.head_dim + ) + self.ln_post = nn.LayerNorm(hidden_size) + + expected_output = hidden_size * self.merge_kernel_size**2 + if config.output_dim != expected_output: + raise ValueError( + f"MuseGlimmer vision output_dim={config.output_dim} does not match " + f"pixel-shuffle output {expected_output}" + ) + + def _make_2d_rope( + self, grid_height: int, grid_width: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + spatial_dim = self.head_dim // 2 + inv_freq = 1.0 / ( + 10000.0 + ** ( + torch.arange(0, spatial_dim, 2, dtype=torch.float32, device=device) + / spatial_dim + ) + ) + height = torch.arange(1, grid_height + 1, dtype=torch.float32, device=device) + width = torch.arange(1, grid_width + 1, dtype=torch.float32, device=device) + height = height.unsqueeze(1).expand(-1, grid_width).reshape(-1) + width = width.unsqueeze(0).expand(grid_height, -1).reshape(-1) + freq_w = torch.outer(width, inv_freq) + freq_h = torch.outer(height, inv_freq) + freqs = torch.cat([freq_w, freq_h], dim=-1) + return torch.cos(freqs), torch.sin(freqs) + + def _get_pos_emb( + self, grid_height: int, grid_width: int, device: torch.device + ) -> torch.Tensor: + h_grid = ( + torch.arange(grid_height, device=device, dtype=torch.float32) + 0.5 + ) * (self.pos_emb_height / grid_height) - 0.5 + w_grid = ( + torch.arange(grid_width, device=device, dtype=torch.float32) + 0.5 + ) * (self.pos_emb_width / grid_width) - 0.5 + h_floor = torch.floor(h_grid).long() + w_floor = torch.floor(w_grid).long() + h_ceil = h_floor + 1 + w_ceil = w_floor + 1 + h_frac = h_grid - h_floor.float() + w_frac = w_grid - w_floor.float() + + h_floor_valid = (h_floor >= 0) & (h_floor < self.pos_emb_height) + h_ceil_valid = (h_ceil >= 0) & (h_ceil < self.pos_emb_height) + w_floor_valid = (w_floor >= 0) & (w_floor < self.pos_emb_width) + w_ceil_valid = (w_ceil >= 0) & (w_ceil < self.pos_emb_width) + h_floor = h_floor.clamp(0, self.pos_emb_height - 1) + h_ceil = h_ceil.clamp(0, self.pos_emb_height - 1) + w_floor = w_floor.clamp(0, self.pos_emb_width - 1) + w_ceil = w_ceil.clamp(0, self.pos_emb_width - 1) + + h_floor_offset = h_floor * self.pos_emb_width + h_ceil_offset = h_ceil * self.pos_emb_width + indices = torch.stack( + [ + (h_floor_offset[:, None] + w_floor[None, :]).flatten(), + (h_floor_offset[:, None] + w_ceil[None, :]).flatten(), + (h_ceil_offset[:, None] + w_floor[None, :]).flatten(), + (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(), + ] + ) + weights = torch.stack( + [ + ( + (1 - h_frac)[:, None] + * (1 - w_frac)[None, :] + * (h_floor_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + (1 - h_frac)[:, None] + * w_frac[None, :] + * (h_floor_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * (1 - w_frac)[None, :] + * (h_ceil_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * w_frac[None, :] + * (h_ceil_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ] + ) + return (self.positional_embedding_vlm[indices] * weights[..., None]).sum(0) + + def _pixel_shuffle_downsample( + self, hidden_states: torch.Tensor, grid_height: int, grid_width: int + ) -> torch.Tensor: + factor = self.merge_kernel_size + output_tokens = (grid_height // factor) * (grid_width // factor) + permutation = torch.arange( + grid_height * grid_width, device=hidden_states.device + ) + permutation = permutation.view( + grid_height // factor, factor, grid_width // factor, factor + ) + permutation = permutation.permute(0, 2, 1, 3).reshape(-1) + hidden_states = hidden_states.squeeze(0)[permutation] + hidden_size = hidden_states.shape[-1] + hidden_states = ( + hidden_states.view(output_tokens, factor * factor, hidden_size) + .permute(0, 2, 1) + .contiguous() + .view(output_tokens, hidden_size * factor * factor) + ) + return hidden_states.unsqueeze(0) + + def _get_sparse_permutation( + self, grid_height: int, grid_width: int, device: torch.device + ) -> tuple[torch.Tensor, list[int]]: + block_height = self.pos_emb_height + block_width = self.pos_emb_width + padded_height = math.ceil(grid_height / block_height) * block_height + padded_width = math.ceil(grid_width / block_width) * block_width + indices = torch.arange(grid_height * grid_width, device=device).view( + grid_height, grid_width + ) + indices = F.pad( + indices, + (0, padded_width - grid_width, 0, padded_height - grid_height), + value=-1, + ).flatten() + indices = indices.view( + padded_height // block_height, + block_height, + padded_width // block_width, + block_width, + ) + indices = indices.permute(0, 2, 1, 3).reshape(-1) + valid = (indices != -1).view(-1, block_height * block_width) + return indices[indices != -1], valid.sum(dim=1).tolist() + + def _get_attention_metadata( + self, + seq_lens: Sequence[int], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + cu_seqlens_np = np.concatenate( + [ + np.zeros(1, dtype=np.int32), + np.asarray(seq_lens, dtype=np.int32).cumsum(dtype=np.int32), + ] + ) + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( + self.attn_backend, + cu_seqlens_np, + device, + ) + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen( + self.attn_backend, + cu_seqlens_np, + ), + dtype=torch.int32, + ) + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( + self.attn_backend, + cu_seqlens_np, + self.hidden_size, + self.tp_size, + device, + fp8_padded_hidden_size=self.fp8_padded_hidden_size, + ) + return cu_seqlens, max_seqlen, sequence_lengths + + def _patchify(self, pixels: torch.Tensor) -> torch.Tensor: + patch_size = self.patch_size + _, channels, height, width = pixels.shape + grid_height = height // patch_size + grid_width = width // patch_size + if channels == 3: + patches = pixels.unfold(2, patch_size, patch_size).unfold( + 3, patch_size, patch_size + ) + patches = patches.contiguous().view( + 1, channels, grid_height, grid_width, patch_size, patch_size + ) + patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous() + patches = patches.unsqueeze(3).expand( + -1, -1, -1, self.patch_temporal, -1, -1, -1 + ) + elif channels == self.patch_temporal * 3: + frame_patches = [] + for frame_idx in range(self.patch_temporal): + frame = pixels[:, frame_idx * 3 : (frame_idx + 1) * 3] + frame = frame.unfold(2, patch_size, patch_size).unfold( + 3, patch_size, patch_size + ) + frame = frame.contiguous().view( + 1, 3, grid_height, grid_width, patch_size, patch_size + ) + frame_patches.append(frame.permute(0, 2, 3, 1, 4, 5).contiguous()) + patches = torch.stack(frame_patches, dim=3) + else: + raise ValueError( + f"MuseGlimmer vision input has {channels} channels; expected 3 or " + f"{self.patch_temporal * 3}" + ) + return patches.reshape(1, grid_height * grid_width, -1) + + def forward(self, pixel_values: Sequence[torch.Tensor]) -> torch.Tensor: + device = self.conv1_linear.weight.device + dtype = self.conv1_linear.weight.dtype + has_sparse_layers = any( + layer_type != "full_attention" for layer_type in self.layer_types + ) + + all_hidden_states = [] + all_rotary_pos_emb_cos = [] + all_rotary_pos_emb_sin = [] + sparse_seq_lens: list[int] = [] + global_seq_lens: list[int] = [] + metadata = [] + for pixels in pixel_values: + pixels = pixels.to(device=device, dtype=dtype) + if pixels.ndim == 3: + pixels = pixels.unsqueeze(0) + if pixels.ndim != 4 or pixels.shape[0] != 1: + raise ValueError( + "Each MuseGlimmer vision input must have shape [C, H, W]" + ) + spatial_stride = self.patch_size * self.merge_kernel_size + if pixels.shape[-2] % spatial_stride or pixels.shape[-1] % spatial_stride: + raise ValueError( + "MuseGlimmer vision input dimensions must divide the " + "merged patch stride" + ) + grid_height = pixels.shape[-2] // self.patch_size + grid_width = pixels.shape[-1] // self.patch_size + num_tokens = grid_height * grid_width + hidden_states = self.conv1_linear(self._patchify(pixels)) + hidden_states = hidden_states + self._get_pos_emb( + grid_height, grid_width, device + ).unsqueeze(0).to(dtype) + hidden_states = self.ln_pre(hidden_states.view(-1, self.hidden_size)).view( + 1, -1, self.hidden_size + ) + rotary_pos_emb_cos, rotary_pos_emb_sin = self._make_2d_rope( + grid_height, grid_width, device + ) + + permutation = None + if has_sparse_layers: + permutation, seq_lens = self._get_sparse_permutation( + grid_height, grid_width, device + ) + hidden_states = hidden_states[:, permutation] + rotary_pos_emb_cos = rotary_pos_emb_cos[permutation] + rotary_pos_emb_sin = rotary_pos_emb_sin[permutation] + sparse_seq_lens.extend(seq_lens) + + all_hidden_states.append(hidden_states.squeeze(0)) + all_rotary_pos_emb_cos.append(rotary_pos_emb_cos) + all_rotary_pos_emb_sin.append(rotary_pos_emb_sin) + global_seq_lens.append(num_tokens) + metadata.append((grid_height, grid_width, num_tokens, permutation)) + + hidden_states = torch.cat(all_hidden_states).unsqueeze(0) + rotary_pos_emb_cos = torch.cat(all_rotary_pos_emb_cos) + rotary_pos_emb_sin = torch.cat(all_rotary_pos_emb_sin) + global_attention_metadata = self._get_attention_metadata( + global_seq_lens, + device, + ) + sparse_attention_metadata = ( + self._get_attention_metadata(sparse_seq_lens, device) + if sparse_seq_lens + else None + ) + for layer_type, block in zip(self.layer_types, self.transformer): + attention_metadata = ( + global_attention_metadata + if layer_type == "full_attention" + else sparse_attention_metadata + ) + if attention_metadata is None: + raise ValueError("MuseGlimmer sparse attention metadata is missing") + hidden_states = block( + hidden_states, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + *attention_metadata, + ) + + features = [] + offset = 0 + for grid_height, grid_width, num_tokens, permutation in metadata: + item = hidden_states[:, offset : offset + num_tokens] + offset += num_tokens + if permutation is not None: + inverse = torch.empty_like(permutation) + inverse[permutation] = torch.arange(len(permutation), device=device) + item = item[:, inverse] + item = self.ln_post(item.view(-1, self.hidden_size)).view( + 1, -1, self.hidden_size + ) + features.append( + self._pixel_shuffle_downsample(item, grid_height, grid_width).squeeze(0) + ) + return torch.cat(features) + + +class MuseGlimmerVisionAdapter(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.c_fc = nn.Linear(config.output_dim, config.adapter_dim, bias=False) + self.c_proj = nn.Linear(config.adapter_dim, config.adapter_dim, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.gelu(self.c_proj(F.gelu(self.c_fc(hidden_states)))) + + +class MuseGlimmerMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_activation: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_activation != "silu": + raise ValueError( + f"MuseGlimmer uses `silu` as the hidden activation; " + f"got `{hidden_activation}`." + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MuseGlimmerAttention(nn.Module): + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.layer_idx = extract_layer_index(prefix) + + tp_size = get_tensor_model_parallel_world_size() + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + # MuseGlimmer overrides Gemma2's query_pre_attn_scalar scaling with the + # standard 1/sqrt(head_dim). The query is *additionally* pre-scaled by + # scale_query_by after QK-norm (see forward). + self.scaling = self.head_dim**-0.5 + + # iRoPE: NoPE layers (no_rope_layers[i] == 0) run full attention; RoPE + # layers run sliding-window attention. + self.use_rope = config.no_rope_layers[self.layer_idx] == 1 + + self.use_qk_norm = _muse_glimmer_use_qk_norm(config) + if self.use_qk_norm: + # Weightless, computed in fp32, applied per head over head_dim. + self.qk_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) + # Post-QK-norm query pre-scale, normalized across the native (raw + # ~43.78) and modular (pre-folded ~3.87) config schemas. + self.scale_query_by = _muse_glimmer_query_prescale(config) + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.use_output_gate = _muse_glimmer_use_attn_output_gate(config) + if self.use_output_gate: + self.output_gate_proj = ColumnParallelLinear( + self.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate_proj", + ) + + self.rotary_emb = ( + get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters, + # HF converter permutes q/k to NEOX (half-split) + # layout via _permute_for_rope + is_neox_style=True, + ) + if self.use_rope + else None + ) + + # Full attention on NoPE layers, sliding window otherwise. + sliding_window = None if not self.use_rope else config.sliding_window + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + logits_soft_cap=None, # MuseGlimmer sets attn_logit_softcapping = None + per_layer_sliding_window=sliding_window, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.use_qk_norm: + # QK-norm over head_dim, fp32, applied BEFORE RoPE; then pre-scale q. + q = q.reshape(-1, self.head_dim) + q = self.qk_norm(q).reshape(-1, self.q_size) * self.scale_query_by + k = k.reshape(-1, self.head_dim) + k = self.qk_norm(k).reshape(-1, self.kv_size) + q = q.to(v.dtype) + k = k.to(v.dtype) + + if self.rotary_emb is not None: + q, k = self.rotary_emb(positions, q, k) + + attn_output = self.attn(q, k, v) + + if self.use_output_gate: + # Gate reads the layer input hidden states (not the attn output). + gate, _ = self.output_gate_proj(hidden_states) + attn_output = torch.sigmoid(gate) * attn_output + + output, _ = self.o_proj(attn_output) + return output + + +class MuseGlimmerDecoderLayer(nn.Module): + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = MuseGlimmerAttention( + config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = MuseGlimmerMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + # Sandwich norms with baked +1 offset. Pre-norms use rms_norm_eps; the + # post-norms use the (typically smaller) post_norm_eps. + self.input_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 + ) + self.post_attention_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.post_norm_eps, weight_offset=1 + ) + self.pre_feedforward_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 + ) + self.post_feedforward_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.post_norm_eps, weight_offset=1 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Gemma2-style sandwich, replicated explicitly (matches HF MuseGlimmer). + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + return hidden_states, residual + + +@support_torch_compile +class MuseGlimmerModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = _text_config(vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + # MuseGlimmer normalizes token embeddings with a weightless RMSNorm instead of + # Gemma's sqrt(hidden_size) multiplier. + self.embed_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: MuseGlimmerDecoderLayer( + config, cache_config, quant_config, prefix=prefix + ), + prefix=f"{prefix}.layers", + ) + # Final norm: weight-as-scale, no offset. + self.norm = MuseGlimmerRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_norm(self.embed_tokens(input_ids)) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, None + ) + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + hidden_states, residual = layer(positions, hidden_states, residual) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, None + ) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + +@MULTIMODAL_REGISTRY.register_processor( + MuseGlimmerMultiModalProcessor, + info=MuseGlimmerProcessingInfo, + dummy_inputs=MuseGlimmerDummyInputsBuilder, +) +class MuseGlimmerForCausalLM( + nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + # Weight-name normalization. Two checkpoint conventions are supported: + # + # * HF MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``): the + # multimodal ``MuseGlimmerConfig`` prefixes the language model with + # ``model.language_model.``; the per-layer sandwich norms are already + # named ``input_layernorm`` / ``post_attention_layernorm`` / + # ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm``. + # + # * Legacy HF export (an earlier checkpoint convention): uses ``model.`` + # and a different sandwich-norm naming where ``post_attn_norm`` is the + # true post-attention norm and ``post_attention_layernorm`` is actually + # the pre-feedforward norm. We remap those to MuseGlimmer's names. + # + # CONVENTION DISAMBIGUATION (critical): the two checkpoint families use + # DIFFERENT sandwich-norm names, and they must not be conflated: + # + # * Canonical MuseGlimmer export (current + # ``convert_muse_glimmer_weights_to_hf.py`` — what partners ship): + # keys are ``model.language_model.layers.N.*`` and + # the norms are ALREADY named ``input_layernorm`` / + # ``post_attention_layernorm`` / ``pre_feedforward_layernorm`` / + # ``post_feedforward_layernorm``. No norm rename needed — pass through. + # + # * Legacy HF export (an earlier checkpoint convention): keys are + # ``model.layers.N.*`` and the sandwich norms are + # named ``input_layernorm`` / ``post_attention_layernorm`` (this one is + # actually the PRE-feedforward norm) / ``post_attn_norm`` (the true + # post-attention norm) / ``post_ffn_norm``. These must be remapped. + # + # The unambiguous discriminator is the PREFIX: legacy keys start with + # ``model.layers.`` while canonical keys start with + # ``model.language_model.layers.``. ``orig_to_new_regex`` runs BEFORE the + # prefix strip (see WeightsMapper._map_name_with_shard), so we anchor the + # legacy renames on ``^model\.layers\.`` — they fire ONLY on legacy keys and + # leave canonical/partner checkpoints untouched. Rule order within the regex + # dict matters: the ``post_attention_layernorm`` -> ``pre_feedforward_...`` + # rule must precede the ``post_attn_norm`` -> ``post_attention_layernorm`` + # rule so the latter's output is not re-captured by the former (regex rules + # apply as a single forward pass). + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "model.vision_tower.patch_embedder.position_embedding_table.weight": ( + "model.vision_tower.positional_embedding_vlm" + ), + "model.vision_tower.layers.": "model.vision_tower.transformer.", + ".norm1.": ".ln_1.", + ".norm2.": ".ln_2.", + ".attn.proj.": ".attn.o_proj.", + ".mlp.fc1.": ".mlp.c_fc.", + ".mlp.fc2.": ".mlp.c_proj.", + ".self_attn.gate_proj": ".self_attn.output_gate_proj", + }, + orig_to_new_prefix={ + "model.rotary_emb.": None, + "model.language_model.": "model.", + "language_model.": "model.", + "model.vision_tower.patch_embedder.patch_embedding.": ( + "model.vision_tower.conv1_linear." + ), + "model.vision_tower.": "vision_encoder.", + "vision_tower.": "vision_encoder.", + "model.vision_encoder.": "vision_encoder.", + "model.vision_adapter.fc1.": "model.vision_adapter.c_fc.", + "model.vision_adapter.fc2.": "model.vision_adapter.c_proj.", + "model.vision_adapter.": "vision_adapter.", + "model.vision_projection.": "vision_projection.", + "model.perception_emb_norm.": "perception_emb_norm.", + }, + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + }, + ) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return IMAGE_TOKEN + if modality.startswith("video"): + return VIDEO_TOKEN + raise ValueError(f"Unsupported MuseGlimmer modality: {modality}") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + text_config = _text_config(config) + vision_config = _vision_config(config) + quant_config = vllm_config.quant_config + self.config = config + self.text_config = text_config + self.quant_config = quant_config + self.has_vision = _muse_glimmer_has_vision(config) + + with self._mark_language_model(vllm_config): + self.model = MuseGlimmerModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + if self.has_vision: + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_encoder = MuseGlimmerVisionEncoder( + vision_config, + prefix=maybe_prefix(prefix, "vision_encoder"), + ) + self.vision_adapter = MuseGlimmerVisionAdapter(vision_config) + self.vision_projection = nn.Linear( + vision_config.adapter_dim, + text_config.hidden_size, + bias=False, + ) + self.perception_emb_norm = ( + MuseGlimmerRMSNorm(eps=text_config.rms_norm_eps, with_scale=False) + if text_config.normalize_tok_embeddings + else nn.Identity() + ) + else: + self.vision_encoder = None + self.vision_adapter = None + self.vision_projection = None + self.perception_emb_norm = None + + self.lm_head = ParallelLMHead( + text_config.vocab_size, + text_config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.output_multiplier = text_config.output_multiplier + self.final_logit_softcapping = text_config.final_logit_softcapping + self.logits_processor = LogitsProcessor(text_config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + image_token_id = getattr( + config, "image_token_id", getattr(config, "patch_token_id", 200092) + ) + video_token_id = getattr(config, "video_token_id", 200091) + self.configure_mm_token_handling( + text_config.vocab_size, [image_token_id, video_token_id] + ) + + def _parse_and_validate_image_input( + self, + **kwargs: object, + ) -> MuseGlimmerImagePixelInputs | None: + pixel_values = kwargs.pop("image_pixel_values", None) + feature_sizes = kwargs.pop("image_feature_sizes", None) + if pixel_values is None and feature_sizes is None: + return None + if pixel_values is None or feature_sizes is None: + raise ValueError( + "MuseGlimmer image_pixel_values and image_feature_sizes " + "must be provided together" + ) + if not isinstance(feature_sizes, torch.Tensor): + raise ValueError("MuseGlimmer image_feature_sizes must be a tensor") + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + return MuseGlimmerImagePixelInputs( + type="image_pixels", + pixel_values=pixel_values, + feature_sizes=feature_sizes.reshape(-1), + ) + + def _parse_and_validate_video_input( + self, + **kwargs: object, + ) -> MuseGlimmerVideoPixelInputs | None: + pixel_values = kwargs.pop("video_pixel_values", None) + feature_sizes = kwargs.pop("video_feature_sizes", None) + if pixel_values is None and feature_sizes is None: + return None + if pixel_values is None or feature_sizes is None: + raise ValueError( + "MuseGlimmer video_pixel_values and video_feature_sizes " + "must be provided together" + ) + if not isinstance(feature_sizes, torch.Tensor): + raise ValueError("MuseGlimmer video_feature_sizes must be a tensor") + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 4: + pixel_values = pixel_values.unsqueeze(0) + patch_temporal = int(_vision_config(self.config).patch_temporal) + return MuseGlimmerVideoPixelInputs( + type="video_pixels", + pixel_values=pixel_values, + feature_sizes=feature_sizes.reshape(-1), + resolve_bindings={"c": patch_temporal * 3}, + ) + + def _encode_pixel_groups( + self, + pixel_groups: Sequence[torch.Tensor], + feature_sizes: Sequence[int], + ) -> tuple[torch.Tensor, ...]: + if ( + self.vision_encoder is None + or self.vision_adapter is None + or self.vision_projection is None + or self.perception_emb_norm is None + ): + raise ValueError("This MuseGlimmer checkpoint has no vision tower") + features = self.vision_encoder(pixel_groups) + features = self.vision_adapter(features) + features = self.vision_projection(features) + features = self.perception_emb_norm(features) + if features.shape[0] != sum(feature_sizes): + raise ValueError( + f"MuseGlimmer produced {features.shape[0]} vision features for " + f"{sum(feature_sizes)} placeholder tokens" + ) + return features.split(list(feature_sizes)) + + def _process_image_input( + self, + image_input: MuseGlimmerImagePixelInputs, + ) -> tuple[torch.Tensor, ...]: + images = list(image_input["pixel_values"]) + sizes = [int(size) for size in image_input["feature_sizes"].tolist()] + if len(images) != len(sizes): + raise ValueError("MuseGlimmer image batch metadata does not match pixels") + return self._encode_pixel_groups(images, sizes) + + def _process_video_input( + self, + video_input: MuseGlimmerVideoPixelInputs, + ) -> tuple[torch.Tensor, ...]: + videos = list(video_input["pixel_values"]) + sizes = [int(size) for size in video_input["feature_sizes"].tolist()] + if len(videos) != len(sizes): + raise ValueError("MuseGlimmer video batch metadata does not match pixels") + groups = [group for video in videos for group in video] + return self._encode_pixel_groups(groups, sizes) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + video_input = self._parse_and_validate_video_input(**kwargs) + + embeddings: MultiModalEmbeddings = [] + for key in kwargs: + if key == "image_pixel_values" and image_input is not None: + embeddings.extend(self._process_image_input(image_input)) + elif key == "video_pixel_values" and video_input is not None: + embeddings.extend(self._process_video_input(video_input)) + return embeddings + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + logits = logits * self.output_multiplier + if self.final_logit_softcapping is not None: + cap = self.final_logit_softcapping + logits = cap * torch.tanh(logits / cap) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 90b2f48faa..35c04331b0 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -72,6 +72,26 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool: ) +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: + """The target's RoPE layout, from its first attention layer. + + A DFlash head must rotate Q/K the way the target it was distilled against + does, and a mismatch is silent — acceptance collapses but nothing errors and + the output stays correct. Draft checkpoints do not carry this, so take it + from the target. None if the target uses no RoPE. + """ + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + for module in language_model.modules(): + style = getattr(module, "is_neox_style", None) + if isinstance(style, bool): + return style + return None + + def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: spec_config = vllm_config.speculative_config config = spec_config.draft_model_config.hf_config @@ -166,6 +186,7 @@ class DFlashQwen3Attention(nn.Module): add_swa_attention_sink_bias: bool = False, sliding_window: int | None = None, causal: bool = False, + is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -209,6 +230,7 @@ class DFlashQwen3Attention(nn.Module): self.rotary_emb = get_rope( self.head_dim, max_position=max_position, + is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) @@ -291,6 +313,13 @@ class DFlashQwen3DecoderLayer(nn.Module): # non-causal) from the draft config. sliding_window, causal = _resolve_layer_attention(config, layer_idx) + # RoPE layout, copied off the target at load time by the draft loader + # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it: + # a head distilled from an interleaved-RoPE target must rotate the way + # that target does, or every drafted Q/K is wrong and acceptance + # collapses with no error raised. + is_neox_style = getattr(config, "is_neox_style", True) + self.self_attn = DFlashQwen3Attention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, @@ -301,6 +330,7 @@ class DFlashQwen3DecoderLayer(nn.Module): add_swa_attention_sink_bias=add_swa_attention_sink_bias, sliding_window=sliding_window, causal=causal, + is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, @@ -345,7 +375,14 @@ class DFlashQwen3DecoderLayer(nn.Module): @support_torch_compile class DFlashQwen3Model(nn.Module): hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={"midlayer.": "layers.0."}, + orig_to_new_substr={ + "midlayer.": "layers.0.", + # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder + # `encoder.fc` / `encoder.output_norm_enc`; this head calls them + # `fc` / `hidden_norm`. Same tensors and shapes, different names. + "encoder.output_norm_enc.": "hidden_norm.", + "encoder.fc.": "fc.", + }, orig_to_new_stacked={ ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 87545b0925..1e4bf6731c 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -183,6 +183,7 @@ _TEXT_GENERATION_MODELS = { "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"), + "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), "OPTForCausalLM": ("opt", "OPTForCausalLM"), @@ -530,6 +531,7 @@ _MULTIMODAL_MODELS = { "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NVLM_D": ("nvlm_d", "NVLM_D_Model"), + "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), "OpenPanguVLForConditionalGeneration": ( "openpangu_vl", @@ -623,6 +625,14 @@ _SPECULATIVE_DECODING_MODELS = { "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash + # implementation. EAGLEConfig rewrites a dflash draft's architecture to + # DFlash{arch} unless it already starts or ends with "DFlash" (see + # transformers_utils/configs/eagle.py), so the name the registry is asked + # for is DFlashMuseGlimmerAssistantModel -- same convention as + # DFlashLagunaForCausalLM. The bare name is kept as a defensive alias. + "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "K3DSparkModel": ( diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 6c761147c9..8b52c4fca8 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -120,6 +120,10 @@ _REASONING_PARSERS_TO_REGISTER = { "olmo3_reasoning_parser", "Olmo3ReasoningParser", ), + "muse_glimmer": ( + "muse_glimmer_reasoning_parser", + "MuseGlimmerReasoningParser", + ), "qwen3": ( "qwen3_engine_reasoning_parser", "Qwen3ParserReasoningAdapter", diff --git a/vllm/reasoning/muse_glimmer_reasoning_parser.py b/vllm/reasoning/muse_glimmer_reasoning_parser.py new file mode 100644 index 0000000000..d3abedca4b --- /dev/null +++ b/vllm/reasoning/muse_glimmer_reasoning_parser.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning-content parser for MuseGlimmer. +Port of the ``reasoning_content`` rule from the HuggingFace MuseGlimmer +``MUSE_GLIMMER_RESPONSE_SCHEMA`` (synced with internal master). MuseGlimmer emits +chain-of-thought in ``to=self`` channels delimited by ``<|message|>`` ... ``<|eom|>``: + to=self<|message|>...reasoning...<|eom|> +A turn may contain several ``to=self`` blocks interleaved with tool calls, and a +tool call or final answer follows in its own channel. +Because MuseGlimmer's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed +to be single vocab tokens across every checkpoint's tokenizer, this parser works +on the decoded text with regexes rather than the single start/end-token base class. +Usage: ``--reasoning-parser muse_glimmer`` +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence + +import regex as re + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser + +_EOM = "<|eom|>" +_EOT = "<|eot|>" +_FUNCTION_CALLS_OPEN = "" +_REASONING_OPEN = "to=self<|message|>" +_ASSISTANT_TURN_OPEN = "<|start|>assistant" +# A channel header: ``to=<|message|>`` where recipient is ``self`` +# (reasoning), ``user`` (final answer) or ``[.]`` (tool call). +_CHANNEL_HEADER_RE = re.compile(r"to=(?P[^\s<]+)<\|message\|>") +_HEADER_PAT = r"to=[^\s<]+<\|message\|>" +# Collapse the gap between reasoning blocks so multiple to=self spans join. +_COLLAPSE_RE = re.compile( + r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL +) +_REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL) +_CONTENT_RE = re.compile( + r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL +) +# Strip a CLOSED reasoning span (header .. <|eom|>). +_STRIP_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL +) +# An UNTERMINATED trailing reasoning span. The model sometimes leaves the +# analysis channel WITHOUT emitting <|eom|>, writing a bare +# ``to=<|message|>`` header instead (observed deterministically for a call +# with EMPTY arguments on a tool that has optional parameters; reproduced on +# other engines too, so it is a model-side defect, not engine-specific). +# +# These two patterns MUST therefore stop at the next channel header rather than +# running to end-of-text. An unbounded ``...$`` version consumes the real tool +# call along with the reasoning: `is_reasoning_end` then never fires, the parser +# never leaves the reasoning phase, the tool parser is never invoked, and the +# entire generation is dropped (empty reasoning, empty content, no tool call). +_STRIP_OPEN_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>" + r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +_OPEN_REASONING_RE = re.compile( + r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +# Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body. +_HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>") +# A trailing fragment that could still grow into a channel header (" t", " to", +# " to=", " to=skill"). Without this the recipient name leaks into reasoning and +# then has to be un-emitted once ``<|message|>`` arrives. +_OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$") + + +def _current_assistant_turn(text: str) -> str: + """Return only the text generated in the current assistant turn. + ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start, + and an MuseGlimmer prompt legitimately contains ATEM markers (``render_tool_defs`` + writes a literal ```` example into the system message, + and prior assistant turns may carry real tool calls). Anchoring on the last + channel-open keeps prompt text from deciding the phase. + """ + idx = text.rfind(_ASSISTANT_TURN_OPEN) + return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text + + +def _trim_open_body(body: str) -> str: + """Hold back any tail of a still-growing body that could still be framing. + Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs + the partial-marker trim (``<``) before the partial-header trim can see + `" to=skill"`. Trimming only once leaks the recipient name as reasoning. + """ + while True: + trimmed = body + for marker in _HOLDBACK_MARKERS: + for k in range(min(len(marker) - 1, len(trimmed)), 0, -1): + if trimmed.endswith(marker[:k]): + trimmed = trimmed[:-k] + break + else: + continue + break + header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed) + if header_tail is not None: + trimmed = trimmed[: header_tail.start()] + if trimmed == body: + return body + body = trimmed + + +class MuseGlimmerReasoningParser(ReasoningParser): + def __init__(self, tokenizer, *args, **kwargs) -> None: + super().__init__(tokenizer, *args, **kwargs) + # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified + # `previous_text` is unsafe: a classified body legitimately shrinks when + # a partial header becomes recognisable, and diffing against the shrunken + # value re-emits text that already went out. + self._emitted_reasoning: str = "" + self._emitted_content: str = "" + self._tool_handoff_done: bool = False + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Preserve MuseGlimmer's ATEM framing tokens in the decoded output. + vLLM's serving default is ``skip_special_tokens=True``, which strips + ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the + parsers run, collapsing reasoning into content and breaking channel + scoping. Unlike the base tool-parser hook we do NOT touch + ``structured_outputs`` -- MuseGlimmer emits native ATEM, not JSON. + """ + request.skip_special_tokens = False + return request + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + """Whether the model has left reasoning and opened a TOOL channel. + A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this + parser surfaces that content itself. Only a real tool channel switches + the ``DelegatingParser`` phase machine over to the tool parser. + Both closed and unterminated reasoning spans are stripped before the + check, so an ```` the model merely echoes inside its CoT + never flips the phase. + """ + try: + text = self.model_tokenizer.decode(input_ids) + except Exception: + return False + remainder = self._tool_channel_remainder(text) + return _FUNCTION_CALLS_OPEN in remainder or " bool: + return self.is_reasoning_end(list(input_ids)) + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + # Content-id slicing is unreliable for multi-token markers; the serving + # path uses extract_reasoning() for the final split. + return [] + + @classmethod + def _scoped_turn(cls, text: str) -> str: + """Current assistant turn with reasoning spans removed.""" + scoped = _current_assistant_turn(text) + scoped = _STRIP_REASONING_RE.sub("", scoped) + return _STRIP_OPEN_REASONING_RE.sub("", scoped) + + @classmethod + def _tool_channel_remainder(cls, text: str) -> str: + """Text from the first tool-channel header onward, framing INCLUDED. + ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever + this parser returns as ``.content`` on the transition delta and commits + it; anything not returned is destroyed. It must start AT the + ``to=<|message|>`` header -- handing over the text after the header + loses the recipient, and the tool parser then sees a bare ``<|message|>``, + classifies it as the content channel, and leaks the ATEM markup. + """ + scoped = cls._scoped_turn(text) + for match in _CHANNEL_HEADER_RE.finditer(scoped): + if match.group("recipient") not in ("self", "user"): + return scoped[match.start() :] + return "" + + @staticmethod + def _classify_bodies(text: str) -> tuple[str, str]: + """Split ``text`` into (reasoning_body, content_body), channel-aware. + Framing markers and tool channels contribute nothing -- the tool parser + owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel + header, or at end-of-text (an OPEN body, which is held back). + """ + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + pos = 0 + n = len(text) + while pos < n: + match = _CHANNEL_HEADER_RE.search(text, pos) + if not match: + break + recipient = match.group("recipient") + body_start = match.end() + eom = text.find(_EOM, body_start) + eot = text.find(_EOT, body_start) + terminators = [p for p in (eom, eot) if p != -1] + next_header = _CHANNEL_HEADER_RE.search(text, body_start) + if next_header is not None: + terminators.append(next_header.start()) + body_end = min(terminators) if terminators else n + body = text[body_start:body_end] + if not terminators: + body = _trim_open_body(body) + if recipient == "self": + reasoning_parts.append(body) + elif ( + # Never surface tool XML echoed into a user channel. + recipient == "user" + and _FUNCTION_CALLS_OPEN not in body + and " str | None: + """Promote un-surfaced content when the stream ends mid-reasoning. + ``DelegatingParser.finalize_generation`` calls this when + ``reasoning_ended`` is still False. Returns only the channel-classified + ``to=user`` body, and only the portion not already streamed. + """ + _, content_body = self._classify_bodies(previous_text) + remainder = content_body[len(self._emitted_content) :] + return remainder or None + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + collapsed = _COLLAPSE_RE.sub("\n", model_output) + matches = _REASONING_RE.findall(collapsed) + reasoning = "\n".join(matches) if matches else None + # Truncation fallback: generation stopped inside a to=self block, so + # there is no closing <|eom|>. Bounded at the next channel header so a + # real tool call that follows a header-less channel switch is not + # absorbed into the reasoning field. + open_match = _OPEN_REASONING_RE.search(model_output) + if open_match and open_match.group(1): + partial = open_match.group(1) + reasoning = f"{reasoning}\n{partial}" if reasoning else partial + # Content is everything that is not a reasoning block. In a + # reasoning+tool-call turn there is no to=user answer, but the tool + # channels MUST be forwarded -- the unified parser runs the tool parser + # on this returned `content`, not on the original model_output. + remainder = _STRIP_REASONING_RE.sub("", model_output) + remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder) + if " DeltaMessage | None: + """Channel-aware streaming split of reasoning vs content. + Classifies the full ``current_text`` and emits only what has not been + emitted yet, so no framing token is ever surfaced and a delta straddling + a channel boundary only contributes the portion inside a real body. + """ + curr_reason, curr_content = self._classify_bodies(current_text) + reasoning_delta = "" + if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len( + self._emitted_reasoning + ): + reasoning_delta = curr_reason[len(self._emitted_reasoning) :] + self._emitted_reasoning = curr_reason + content_delta = "" + if curr_content.startswith(self._emitted_content) and len(curr_content) > len( + self._emitted_content + ): + content_delta = curr_content[len(self._emitted_content) :] + self._emitted_content = curr_content + # Hand the tool channel to the tool parser exactly once, starting at its + # header. parse_delta discards anything not returned here. + # + # This MUST fire on the same delta where is_reasoning_end() flips, i.e. + # only once the tool channel actually contains ATEM. Emitting it earlier + # -- when only the bare `to=<|message|>` header has arrived -- keeps + # the parser in the reasoning phase, so parse_delta never replaces this + # DeltaMessage with the tool parser's and the header is delivered to the + # client as visible content. + handoff = "" + if not self._tool_handoff_done: + remainder = self._tool_channel_remainder(current_text) + if _FUNCTION_CALLS_OPEN in remainder or "assistant to=self<|message|>...reasoning...<|eom|> + <|start|>assistant to=.<|message|> + + + value + + <|eom|> # non-final call + <|start|>assistant to=user<|message|>...final answer...<|eot|> + +Channel scoping is essential: an ```` echoed inside a ``to=self`` +reasoning block or a ``to=user`` final answer must NOT be parsed as a real tool +call. + +Rather than *subtracting* reasoning/answer spans with regex substitutions (the +approach the HF ``response_schema`` uses, which is safe only on a complete, +well-formed turn), this parser *segments* the output into messages and then +*selects* the tool-channel bodies. On a complete turn the two are equivalent; +on a truncated or damaged turn, subtraction can delete a valid tool call +(an unterminated ``to=self`` block makes the non-greedy strip run to the next +``<|eom|>``, which belongs to the tool-call message) whereas selection cannot. + +Usage: ``--enable-auto-tool-choice --tool-call-parser muse_glimmer`` +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Sequence + +import regex as re +from openai.types.responses import ToolChoiceFunction +from transformers import PreTrainedTokenizerBase + +from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.logger import init_logger +from vllm.tool_parsers.abstract_tool_parser import ( + Tool, + ToolParser, +) + +logger = init_logger(__name__) + +# --- Message framing ------------------------------------------------------- +# An assistant message header. All three parts are optional except the +# <|message|> terminator: +# "<|start|>assistant to=get_weather<|message|>" -- after an <|eom|> boundary +# " to=self<|message|>" -- first message of a turn +# (the prompt already ended +# with "<|start|>assistant") +# "<|message|>" -- bare recipient (public CoT +# / untagged content) +_MSG_HEADER_RE = re.compile( + r"(?:<\|start\|>\s*assistant)?[^\S\n]*(?:to=(?P[A-Za-z0-9_.\-]+))?<\|message\|>" +) +_MSG_END_RE = re.compile(r"<\|eom\|>|<\|eot\|>") + +# Recipients whose bodies are NOT tool calls. +_REASONING_RECIPIENT = "self" +_USER_RECIPIENT = "user" + +# Structural markers that must never reach the client. A streamed body is held +# back by up to len(marker)-1 characters so a marker split across two chunks is +# not emitted as content. +_STRUCTURAL_MARKERS = ("<|eom|>", "<|eot|>", "<|start|>", "<|message|>") +_MAX_MARKER_LEN = max(len(m) for m in _STRUCTURAL_MARKERS) +# A trailing " to=NAME" that could still grow into a bare message header (the +# first message of a turn has no <|start|> prefix -- the prompt ends with +# "<|start|>assistant", so the model's first emitted text is " to=self<|message|>"). +_OPEN_TAIL_TO_RE = re.compile(r"[^\S\n]+to=[A-Za-z0-9_.\-]*$") + +# --- Tool-call extraction (unchanged from MUSE_GLIMMER_RESPONSE_SCHEMA) ------------- +_INVOKE_RE = re.compile(r"()", re.DOTALL) +_NAME_RE = re.compile(r']*?\bname="([^"]+)"') +_PARAM_RE = re.compile( + r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)', + re.DOTALL, +) +_FUNCTION_CALLS_OPEN = "" + + +def _decode_value(raw: str): + """JSON-decode a parameter value when possible, else keep the raw string. + + Mirrors the schema's ``x-parser: json`` with ``allow_non_json: True``. + """ + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw + + +def _iter_messages(text: str) -> Iterator[tuple[str | None, str, bool]]: + """Segment *text* into assistant messages. + + Yields ``(recipient, body, closed)`` per message, where ``recipient`` is + ``None`` for a bare ``<|message|>`` header and ``closed`` is False for a + message that has not (yet) seen ``<|eom|>`` / ``<|eot|>``. + + A message is also terminated by the start of the NEXT header. Without that, + a reasoning block whose ``<|eom|>`` is missing (truncation, or a chunk + dropped at the reasoning -> tool transition) would absorb the tool-call + message that follows it and the call would be lost -- the same defect the + subtractive regexes have. + """ + pos = 0 + while pos < len(text): + header = _MSG_HEADER_RE.search(text, pos) + if header is None: + return + body_start = header.end() + end = _MSG_END_RE.search(text, body_start) + nxt = _MSG_HEADER_RE.search(text, body_start) + body_end = end.start() if end is not None else len(text) + closed = end is not None + if nxt is not None and nxt.start() < body_end: + body_end = nxt.start() + closed = False + next_pos = nxt.start() + else: + next_pos = end.end() if end is not None else len(text) + body = text[body_start:body_end] + # A body can never legitimately contain <|start|>. Seeing one means the + # next header is only partially generated (its <|message|> has not + # arrived), so the regex above could not recognise it yet. Cut there, + # otherwise the streamed body would grow to include the next header and + # then shrink back once it completes. + start_tok = body.find("<|start|>") + if start_tok != -1: + body = body[:start_tok] + closed = False + yield header.group("rcpt"), body, closed + pos = next_pos + + +def _trailing_partial_marker_len(text: str) -> int: + """Length of the longest suffix of *text* that prefixes a structural marker.""" + max_overlap = min(len(text), _MAX_MARKER_LEN - 1) + for overlap in range(max_overlap, 0, -1): + suffix = text[-overlap:] + if any(marker.startswith(suffix) for marker in _STRUCTURAL_MARKERS): + return overlap + return 0 + + +def _safe_open_body(body: str) -> str: + """Trim the tail of a still-growing body to what is safe to emit now. + + Holds back anything that could still turn out to be structural, so the + emitted prefix only ever grows. Chunks under speculative decoding are large + enough that markers routinely straddle them. + """ + tail_to = _OPEN_TAIL_TO_RE.search(body) + if tail_to is not None: + return body[: tail_to.start()] + partial = _trailing_partial_marker_len(body) + return body[: len(body) - partial] if partial else body + + +class MuseGlimmerToolParser(ToolParser): + # MuseGlimmer emits ATEM markup around tool-call arguments. The generic + # named/required tool_choice path in vllm/parser/abstract_parser.py assigns + # the raw model text straight to FunctionCall.arguments, which leaks that + # framing and yields invalid JSON; the "required" branch then silently + # swallows the ValidationError and returns tool_calls=null. Opting out + # routes named/required through extract_tool_calls / + # extract_tool_calls_streaming -- the same path "auto" already uses. + supports_required_and_named = False + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + tools: list[Tool] | None = None, + ) -> None: + super().__init__(tokenizer, tools) + # Streaming cursors. vLLM constructs one ToolParser per request, so + # instance state is per-stream. + self._streamed_content_len: int = 0 + self._streamed_reasoning_len: int = 0 + self._emitted_tool_calls: int = 0 + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Force special tokens through, and keep JSON guided decoding off. + + ``skip_special_tokens`` defaults to True on both ChatCompletionRequest + and ResponsesRequest. Every rule in this parser keys off ``<|message|>`` + / ``<|eom|>`` / ``<|eot|>`` / ``<|start|>``, so with the default the + channel framing is stripped before we see it: message segmentation + finds nothing, reasoning-channel invokes are indistinguishable from real + ones, and the raw ATEM markup falls through to the client as content. + Set it unconditionally and FIRST -- not only for tools requests, and not + relying on the reasoning parser to have set it. + + For required/named ``tool_choice`` the base hook installs a JSON schema + constraint (ToolParser.adjust_request -> get_json_schema_from_tools). + MuseGlimmer emits ATEM XML, so under that constraint it writes JSON *inside* + the tool channel -- ``[{"name": ...`` -- with no + ```` for this parser to find. Skip the base hook so those + choices decode natively, the same as "auto". + """ + request.skip_special_tokens = False + + tool_choice = getattr(request, "tool_choice", None) + if request.tools and ( + tool_choice == "required" + or isinstance( + tool_choice, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + ) + ): + return request + return super().adjust_request(request) + + # ---------------- channel selection ---------------- + + @classmethod + def _tool_channel_text(cls, text: str) -> str: + """Concatenate the bodies of messages addressed to a tool. + + Falls back to the whole text when no message header is present at all -- + that means the framing never reached us (``skip_special_tokens`` was on, + or the chunk carrying the header was dropped upstream), and scanning + everything is strictly better than returning nothing. + """ + bodies = [ + body + for rcpt, body, _closed in _iter_messages(text) + if rcpt is not None + and rcpt != _REASONING_RECIPIENT + and rcpt != _USER_RECIPIENT + ] + if bodies: + return "\n".join(bodies) + if _MSG_HEADER_RE.search(text) is None and ( + _FUNCTION_CALLS_OPEN in text or " tuple[str, str, bool, bool]: + """Return ``(content, reasoning, content_open, reasoning_open)``. + + The ``*_open`` flags say whether that channel's LAST message is still + being generated; only then must the caller hold back a partial + structural marker. Tracking them per channel matters: a closed + reasoning block whose text happens to end in ``<`` would otherwise stay + permanently truncated while a later content message is open. + """ + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + content_open = False + reasoning_open = False + for rcpt, body, closed in _iter_messages(text): + if rcpt == _REASONING_RECIPIENT: + reasoning_parts.append(body) + reasoning_open = not closed + elif rcpt is None or rcpt == _USER_RECIPIENT: + content_parts.append(body) + content_open = not closed + return ( + "".join(content_parts), + "".join(reasoning_parts), + content_open, + reasoning_open, + ) + + # ---------------- tool name binding ---------------- + + @staticmethod + def _registered_names(request: ChatCompletionRequest | None) -> set[str]: + """Names of the tools the client registered on this request.""" + names: set[str] = set() + tools = getattr(request, "tools", None) if request is not None else None + for t in tools or []: + fn = getattr(t, "function", None) or t + name = getattr(fn, "name", None) + if name is None and isinstance(fn, dict): + name = fn.get("name") + if name: + names.add(name) + return names + + @staticmethod + def _normalize_name(emitted: str, registered: set[str]) -> str: + """Map an emitted ATEM invoke name back to a registered tool name. + + When a client registers a BARE name (e.g. ``get_weather``) the shipped + chat template renders the valid recipient as ``"get_weather.*"``, and + the model duly emits ``get_weather.get_weather``. Collapsing that + doubled form is safe: head and tail are identical and the collapsed + name is registered. + + Anything else is passed through unchanged. Matching on the trailing + segment alone is NOT safe -- an emitted ``weather.get`` against a + registered ``{calendar.get}`` has a unique leaf match and would silently + dispatch the wrong tool. + """ + if not registered or emitted in registered: + return emitted + head, sep, tail = emitted.partition(".") + if sep and head == tail and head in registered: + return head + logger.warning( + "MuseGlimmer: emitted tool name %r does not match any registered tool; " + "passing through unchanged.", + emitted, + ) + return emitted + + @classmethod + def _parse_tool_calls( + cls, text: str, registered: set[str] | None = None + ) -> list[ToolCall]: + registered = registered or set() + scoped = cls._tool_channel_text(text) + tool_calls: list[ToolCall] = [] + for invoke in _INVOKE_RE.findall(scoped): + name_m = _NAME_RE.search(invoke) + if not name_m: + continue + name = cls._normalize_name(name_m.group(1), registered) + args: dict = {} + for pm in _PARAM_RE.finditer(invoke): + args[pm.group("key")] = _decode_value(pm.group("value")) + tool_calls.append( + ToolCall( + function=FunctionCall( + name=name, + arguments=json.dumps(args, ensure_ascii=False), + ) + ) + ) + return tool_calls + + @classmethod + def _extract_content(cls, text: str) -> str | None: + """Return the user-facing body, or the raw text when unframed.""" + content, reasoning, _c_open, _r_open = cls._visible_channels(text) + if content: + return content + # No framing at all -> the whole thing is plain content. + if not reasoning and _MSG_HEADER_RE.search(text) is None: + return text or None + return None + + # ---------------- non-streaming ---------------- + + def extract_tool_calls( + self, model_output: str, request: ChatCompletionRequest + ) -> ExtractedToolCallInformation: + if ( + _FUNCTION_CALLS_OPEN not in model_output + and "... parsed -- typically a truncated + # call (finish_reason='length'/abort). Log it: silently + # returning "no tool call" here is indistinguishable from the + # model choosing not to call one, which makes this failure mode + # invisible in production. + logger.warning( + "MuseGlimmer: tool channel opened but no complete " + "parsed (truncated tool call?); returning content only." + ) + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._extract_content(model_output), + ) + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=self._extract_content(model_output), + ) + except Exception: + logger.exception("Error extracting MuseGlimmer ATEM tool calls.") + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + # ---------------- streaming ---------------- + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + """Incremental ATEM streaming for tool calls AND content. + + This parser owns every delta once reasoning has ended: in + vllm/parser/abstract_parser.py::parse_delta the "pass through as + content" fallback is guarded by ``not self._in_tool_call_phase(state)``, + and ``_in_tool_call_phase`` is simply ``tool_parser is not None and + state.reasoning_ended``. So with a tool parser loaded that fallback is + dead code, and returning None here DISCARDS the delta. Anything we do + not emit -- including the ``to=user`` final answer -- never reaches the + client. Hence content is emitted here, not left to the reasoning parser. + + Tool calls are surfaced only when an ```` block becomes + complete: the XML is opaque until closed and MuseGlimmer parameters are not + incremental JSON, so there is nothing meaningful to stream before then. + """ + if not previous_text: + # First delta of the tool phase (parse_delta resets previous_text + # to "" when it hands the stream over). Reset the cursors. + self._streamed_content_len = 0 + self._streamed_reasoning_len = 0 + self._emitted_tool_calls = 0 + + try: + registered = self._registered_names(request) + calls = self._parse_tool_calls(current_text, registered) + content, reasoning, content_open, reasoning_open = self._visible_channels( + current_text + ) + + # Trim the tail of a channel that is still growing, so the emitted + # prefix never shrinks between deltas. + if content_open: + content = _safe_open_body(content) + if reasoning_open: + reasoning = _safe_open_body(reasoning) + + content_delta = content[self._streamed_content_len :] + reasoning_delta = reasoning[self._streamed_reasoning_len :] + + tool_deltas: list[DeltaToolCall] = [] + for i in range(self._emitted_tool_calls, len(calls)): + fn = calls[i].function + tool_deltas.append( + DeltaToolCall( + index=i, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=fn.name, + arguments=fn.arguments, + ).model_dump(exclude_none=True), + ) + ) + + if not content_delta and not reasoning_delta and not tool_deltas: + return None + + self._streamed_content_len = len(content) + self._streamed_reasoning_len = len(reasoning) + self._emitted_tool_calls = len(calls) + + message = DeltaMessage() + if content_delta: + message.content = content_delta + if reasoning_delta: + message.reasoning = reasoning_delta + if tool_deltas: + message.tool_calls = tool_deltas + return message + except Exception: + logger.exception("Error extracting MuseGlimmer ATEM streaming tool calls.") + return None diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 8bbba8df24..87271c2588 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -102,6 +102,10 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", kimi_k25="KimiK25Config", + muse_glimmer="MuseGlimmerConfig", + muse_glimmer_text="MuseGlimmerTextConfig", + muse_glimmer_vision="MuseGlimmerVisionConfig", + muse_glimmer_assistant="MuseGlimmerAssistantConfig", kimi_k3="KimiK3Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index a6d477831d..c38cc4b4b8 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -74,6 +74,10 @@ _CLASS_TO_MODULE: dict[str, str] = { "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", + "MuseGlimmerConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerTextConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerVisionConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer", "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", "NemotronConfig": "vllm.transformers_utils.configs.nemotron", @@ -162,6 +166,10 @@ __all__ = [ "KimiLinearConfig", "KimiVLConfig", "KimiK25Config", + "MuseGlimmerConfig", + "MuseGlimmerTextConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerAssistantConfig", "KimiK3Config", "KimiK3VisionConfig", "NemotronConfig", diff --git a/vllm/transformers_utils/configs/muse_glimmer.py b/vllm/transformers_utils/configs/muse_glimmer.py new file mode 100644 index 0000000000..9022c84cf0 --- /dev/null +++ b/vllm/transformers_utils/configs/muse_glimmer.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MuseGlimmer model configuration for vLLM. + +Native vLLM copy of the MuseGlimmer HuggingFace configs +(``configuration_muse_glimmer.py``). MuseGlimmer's +``model_type`` is not yet registered in released transformers, so vLLM ships +this config so partners can serve MuseGlimmer checkpoints *without* trust_remote_code. + +The field set and defaults are kept byte-for-byte in sync with the HF reference +so a checkpoint's ``config.json`` deserializes identically here. Both the text +and vision configs are consumed by the native multimodal serving path. +""" + +from __future__ import annotations + +from transformers import Qwen3Config +from transformers.configuration_utils import PretrainedConfig + + +def _default_no_rope_layers(num_hidden_layers: int) -> list[int]: + # iRoPE mask: NoPE every 4 layers, counted backward from the last layer. + stride = 4 + return [ + 0 if (num_hidden_layers - 1 - i) % stride == 0 else 1 + for i in range(num_hidden_layers) + ] + + +class MuseGlimmerTextConfig(PretrainedConfig): + model_type = "muse_glimmer_text" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 202_048, + hidden_size: int = 6656, + intermediate_size: int = 19968, + num_hidden_layers: int = 52, + num_attention_heads: int = 32, + num_key_value_heads: int = 2, + head_dim: int = 128, + hidden_activation: str = "silu", + max_position_embeddings: int = 16_384, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-5, + use_cache: bool = True, + pad_token_id: int | None = None, + eos_token_id: int | list[int] | None = 200_001, + bos_token_id: int | None = 200_000, + tie_word_embeddings: bool = False, + rope_parameters: dict | None = None, + rope_theta: float | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + query_pre_attn_scalar: int = 256, + sliding_window: int | None = 2048, + layer_types: list[str] | None = None, + final_logit_softcapping: float | None = 20.0, + attn_logit_softcapping: float | None = None, + use_bidirectional_attention: bool | None = None, + # MuseGlimmer-specific + qk_scale_factor: float = 43.7840518911, + use_qk_norm: bool = True, + use_attn_output_gate: bool = True, + output_multiplier: float = 0.19611613513818404, + normalize_tok_embeddings: bool = True, + post_norm_eps: float = 1e-8, + no_rope_layers: list[int] | None = None, + **kwargs, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_activation = hidden_activation + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.query_pre_attn_scalar = query_pre_attn_scalar + self.sliding_window = sliding_window + self.final_logit_softcapping = final_logit_softcapping + self.attn_logit_softcapping = attn_logit_softcapping + self.use_bidirectional_attention = use_bidirectional_attention + + # RoPE: accept either an explicit rope_parameters dict (HF 5.x) or a + # bare rope_theta; normalize to rope_parameters for vLLM's get_rope. + if rope_parameters is None: + theta = rope_theta if rope_theta is not None else 500_000.0 + rope_parameters = {"rope_type": "default", "rope_theta": theta} + self.rope_parameters = rope_parameters + # vLLM reads rope_theta off the config in some codepaths. + self.rope_theta = rope_parameters.get("rope_theta", 500_000.0) + + # MuseGlimmer-specific fields + self.qk_scale_factor = qk_scale_factor + self.use_qk_norm = use_qk_norm + self.use_attn_output_gate = use_attn_output_gate + self.output_multiplier = output_multiplier + self.normalize_tok_embeddings = normalize_tok_embeddings + self.post_norm_eps = post_norm_eps + + self.no_rope_layers = ( + no_rope_layers + if no_rope_layers is not None + else _default_no_rope_layers(num_hidden_layers) + ) + if layer_types is None: + layer_types = [ + "full_attention" if self.no_rope_layers[i] == 0 else "sliding_attention" + for i in range(num_hidden_layers) + ] + self.layer_types = layer_types + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class MuseGlimmerVisionConfig(PretrainedConfig): + model_type = "muse_glimmer_vision" + + def __init__( + self, + patch_size: int = 14, + pos_emb_height: int = 32, + pos_emb_width: int = 32, + num_attention_heads: int = 16, + num_hidden_layers: int = 50, + hidden_size: int = 1536, + intermediate_size: int = 8960, + hidden_act: str = "gelu", + merge_kernel_size: int = 2, + rope_parameters: dict | None = None, + max_position_embeddings: int = 32 * 32, + output_dim: int = 6144, + patch_temporal: int = 2, + adapter_dim: int = 4096, + layer_norm_eps: float = 1e-5, + layer_types: list[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.patch_size = patch_size + self.pos_emb_height = pos_emb_height + self.pos_emb_width = pos_emb_width + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.merge_kernel_size = merge_kernel_size + self.rope_parameters = rope_parameters + self.max_position_embeddings = max_position_embeddings + self.output_dim = output_dim + self.patch_temporal = patch_temporal + self.adapter_dim = adapter_dim + self.layer_norm_eps = layer_norm_eps + if layer_types is None: + stride = 4 + layer_types = [ + "full_attention" + if (i + 1) % stride == 0 or i == num_hidden_layers - 1 + else "sliding_attention" + for i in range(num_hidden_layers) + ] + self.layer_types = layer_types + + +class MuseGlimmerConfig(PretrainedConfig): + model_type = "muse_glimmer" + sub_configs = { + "text_config": MuseGlimmerTextConfig, + "vision_config": MuseGlimmerVisionConfig, + } + + # --- Flat (legacy-converter) -> canonical normalization ------------------- + # MuseGlimmer checkpoints exist in two config layouts in the wild: + # + # * CANONICAL (current HF converter, transformers 5.15): nested + # ``text_config`` / ``vision_config`` sub-dicts with canonical field + # names (``hidden_activation``, ``final_logit_softcapping``, + # ``rope_parameters``, ``vision_config.hidden_size`` ...). + # + # * FLAT (older converter, e.g. Ruan's ``rl_v1/hf``, transformers 5.9): + # every field is a top-level key, with different names + # (``hidden_act``, ``output_soft_cap_temp``, ``rope_theta``, + # ``vision_latent_dim`` ...) and NO ``text_config`` nesting. + # + # Without normalization a flat config silently deserializes to an + # ALL-DEFAULT text config (every checkpoint value ignored) — a dangerous + # correctness bug: a checkpoint whose arch differs from the defaults would + # load into a wrong-shaped model with no error. So when we detect a flat + # config we hoist the flat fields into ``text_config`` / ``vision_config`` + # and rename them to canonical names. + + # flat text field name -> canonical MuseGlimmerTextConfig field name + _FLAT_TEXT_RENAMES = { + "hidden_act": "hidden_activation", + "output_soft_cap_temp": "final_logit_softcapping", + } + # canonical MuseGlimmerTextConfig constructor params that may appear flat + _FLAT_TEXT_KEYS = frozenset( + { + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_activation", + "hidden_act", + "max_position_embeddings", + "initializer_range", + "rms_norm_eps", + "use_cache", + "tie_word_embeddings", + "rope_parameters", + "rope_theta", + "attention_bias", + "attention_dropout", + "query_pre_attn_scalar", + "sliding_window", + "layer_types", + "final_logit_softcapping", + "output_soft_cap_temp", + "attn_logit_softcapping", + "use_bidirectional_attention", + "qk_scale_factor", + "use_qk_norm", + "use_attn_output_gate", + "output_multiplier", + "normalize_tok_embeddings", + "post_norm_eps", + "no_rope_layers", + } + ) + # flat vision field name -> canonical MuseGlimmerVisionConfig field name + _FLAT_VISION_RENAMES = { + "vision_latent_dim": "hidden_size", + "vision_heads": "num_attention_heads", + "vision_layers": "num_hidden_layers", + "vision_output_dim": "output_dim", + "vision_patch_size": "patch_size", + "vision_patch_temporal": "patch_temporal", + "vision_adapter_dim": "adapter_dim", + "vision_pos_emb_grid_h": "pos_emb_height", + "vision_pos_emb_grid_w": "pos_emb_width", + "vision_downsample_factor": "merge_kernel_size", + } + + @classmethod + def _looks_flat(cls, kwargs: dict) -> bool: + # Flat if there is no explicit text_config but there ARE text-level + # fields at the top level (e.g. hidden_size / num_hidden_layers). + if kwargs.get("text_config") is not None: + return False + return any(k in kwargs for k in ("hidden_size", "num_hidden_layers")) + + @classmethod + def _normalize_flat(cls, kwargs: dict) -> dict: + kwargs = dict(kwargs) + text: dict = {} + for key in list(kwargs.keys()): + if key in cls._FLAT_TEXT_KEYS: + canon = cls._FLAT_TEXT_RENAMES.get(key, key) + text.setdefault(canon, kwargs.pop(key)) + + vision_mlp_ratio = kwargs.pop("vision_mlp_ratio", None) + sparse_factor = kwargs.pop("vision_sparse_attention_factor", None) + vision: dict = {} + for key in list(kwargs.keys()): + if key in cls._FLAT_VISION_RENAMES: + vision.setdefault(cls._FLAT_VISION_RENAMES[key], kwargs.pop(key)) + + if vision_mlp_ratio is not None: + hidden_size = vision.get("hidden_size", 1536) + vision["intermediate_size"] = int(vision_mlp_ratio * hidden_size) + if sparse_factor is not None: + sparse_factor = int(sparse_factor) + if sparse_factor <= 0: + raise ValueError("vision_sparse_attention_factor must be positive") + num_layers = vision.get("num_hidden_layers", 50) + vision["layer_types"] = [ + "full_attention" + if (layer_idx + 1) % sparse_factor == 0 or layer_idx == num_layers - 1 + else "sliding_attention" + for layer_idx in range(num_layers) + ] + + if text: + kwargs["text_config"] = text + if vision: + kwargs["vision_config"] = vision + return kwargs + + def __init__( + self, + text_config: dict | MuseGlimmerTextConfig | None = None, + vision_config: dict | MuseGlimmerVisionConfig | None = None, + image_token_id: int = 200092, + video_token_id: int = 200091, + **kwargs, + ) -> None: + # Detect + fold a flat config into nested text/vision before building + # the sub-configs. (image/video token ids are handled below; a flat + # ``patch_token_id`` alias is mapped too.) + if text_config is None and self._looks_flat(kwargs): + kwargs = self._normalize_flat(kwargs) + text_config = kwargs.pop("text_config", None) + vision_config = kwargs.pop("vision_config", vision_config) + # flat image-token alias + if "patch_token_id" in kwargs and "image_token_id" not in kwargs: + image_token_id = kwargs.pop("patch_token_id") + + if text_config is None: + self.text_config = MuseGlimmerTextConfig() + elif isinstance(text_config, dict): + self.text_config = MuseGlimmerTextConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = MuseGlimmerVisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = MuseGlimmerVisionConfig(**vision_config) + else: + self.vision_config = vision_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.patch_token_id = image_token_id + super().__init__(**kwargs) + + +class MuseGlimmerAssistantConfig(Qwen3Config): + """Config for the Muse Glimmer DFlash draft head. + + The head is Qwen3-shaped and runs on vLLM's generic ``qwen3_dflash`` + implementation, so this derives from ``Qwen3Config``. It cannot BE + ``Qwen3Config``, because two of the checkpoint's values do not survive that + class: + + * ``sliding_window`` is gated behind ``use_sliding_window``, which defaults + to False -- ``Qwen3Config(sliding_window=2048).sliding_window`` is None. + The checkpoint declares ``sliding_window: 2048`` and five + ``sliding_attention`` layers, so the window silently disappears and the + DFlash path then raises "sliding attention requires a window size". + * ``vocab_size`` is absent from the checkpoint, so Qwen3's default of + 151936 applies instead of Muse Glimmer's 202048. That one is *silent*: it + builds an all-zero ``draft_id_to_target_id`` remap, and also puts + pad/bos/eos/mask token ids out of range. + + Both defaults are set here so an unmodified checkpoint loads correctly. + """ + + model_type = "muse_glimmer_assistant" + + def __init__( + self, + vocab_size: int = 202048, + use_sliding_window: bool = True, + dflash_config: dict | None = None, + **kwargs, + ) -> None: + super().__init__( + vocab_size=vocab_size, use_sliding_window=use_sliding_window, **kwargs + ) + # muse_glimmer_assistant uses non-causal attention. + self.dflash_config = {"causal": False} | (dflash_config or {}) + + +__all__ = [ + "MuseGlimmerTextConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerConfig", + "MuseGlimmerAssistantConfig", +] diff --git a/vllm/transformers_utils/processors/muse_glimmer.py b/vllm/transformers_utils/processors/muse_glimmer.py new file mode 100644 index 0000000000..c1a5326ea1 --- /dev/null +++ b/vllm/transformers_utils/processors/muse_glimmer.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace processor for the MuseGlimmer multimodal model (text + image + video). + +This packages MuseGlimmer's vision preprocessing and token-span expansion the standard +HF way, so you can do: + + processor = AutoProcessor.from_pretrained(hf_dir, trust_remote_code=True) + messages = [{"role": "user", "content": [ + {"type": "image"}, + {"type": "text", "text": "Describe this image."}, + ]}] + text = processor.apply_chat_template(messages, add_generation_prompt=True) + inputs = processor(text=text, images=[pil_image], return_tensors="pt") + out = model.generate(**inputs, max_new_tokens=256) + +It reproduces the exact token layout the model expects: + + image -> <|image_start|> + <|patch|> * N + <|image_end|> + video -> <|vid_start|> + ( "Time: X.Xs" + <|video|> * P [+ <|vid_frame_separator|>] )* + + <|vid_end|> + +where N = floor(H/(patch*ds)) * floor(W/(patch*ds)) (<= 4096) for images and +P = the same grid count per temporal frame-group (<= 144) for video, with the +target (H, W) chosen by MuseGlimmerImageProcessor.compute_image_size / +MuseGlimmerVideoProcessor.compute_video_frame_size (the same grid logic as +modeling_muse_glimmer.MuseGlimmerVisionEncoder). + +The chat template emits one sentinel per media item -- ``<|image|>`` for images, +``<|video|>`` for videos -- which ``MuseGlimmerProcessor.__call__`` expands into the +spans above using each media item's computed token count. Vision features are +returned in ``pixel_values`` (a list of per-image / per-frame-group tensors) in +the order the sentinels appear, matching how ``MuseGlimmerModel`` merges them at the +``<|patch|>`` / ``<|video|>`` positions. +""" + +from __future__ import annotations + +import itertools +import math +from pathlib import Path + +import torch +from PIL import Image +from torchvision import transforms as T +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_processing_utils import BaseImageProcessor +from transformers.processing_utils import ProcessorMixin +from transformers.video_processing_utils import BaseVideoProcessor + +# Single-token sentinels emitted by the chat template, one per media item. +IMAGE_SENTINEL = "<|image|>" # id 200090; absent from the image/video spans +VIDEO_SENTINEL = "<|video|>" # id 200091; also the per-frame token (see __call__) + +# Jinja chat template (multimodal superset of convert_muse_glimmer_to_hf.CHAT_TEMPLATE). +# Message content may be a plain string OR a list of parts +# ({"type": "text"|"image"|"video", ...}); image/video parts render as a single +# sentinel that __call__ expands. +MUSE_GLIMMER_MM_CHAT_TEMPLATE = ( + "{{- bos_token -}}" + "{%- macro render_parts(content) -%}" + "{%- if content is string -%}{{- content -}}" + "{%- else -%}" + "{%- for part in content -%}" + "{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}" + "{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}" + "{%- elif part['type'] == 'text' -%}{{- part['text'] -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- endif -%}" + "{%- endmacro -%}" + # At inference (add_generation_prompt), inject a default system message when + # the caller supplied none. Full-transcript rendering + # (add_generation_prompt=False) is left unchanged. + "{%- set ns = namespace(has_system=false) -%}" + "{%- for m in messages -%}" + "{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt and not ns.has_system -%}" + "{{- '<|start|>system<|message|>You are a helpful assistant.<|eot|>' -}}" + "{%- endif -%}" + "{%- for message in messages -%}" + "{%- set role = message['role'] -%}" + "{%- if role == 'assistant' -%}" + "{%- set recipient = message.get('recipient') -%}" + "{%- 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_parts(message['content']) -}}" + "{{- ('<|eot|>' if end_turn else '<|eom|>') -}}" + "{%- elif role == 'tool' -%}" + "{%- set name = message.get('name', '') -%}" + # Tool content is emitted as-is (string body or interleaved image/text parts). + # Any wrapper is baked into the SFT data content; the + # tokenizer does not add it, so the template must not either. + "{{- '<|start|>tool ' + name + '<|message|>' -}}" + "{{- render_parts(message['content']) -}}" + "{{- '<|eot|>' -}}" + "{%- else -%}" + "{%- set header = role -%}" + "{%- if message.get('name') -%}" + "{%- set header = role + ' ' + message['name'] -%}" + "{%- endif -%}" + "{{- '<|start|>' + header + '<|message|>' -}}" + "{{- render_parts(message['content']) -}}" + "{{- '<|eot|>' -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}" +) + + +def _grid_size( + img_w: int, img_h: int, patch_hw: int, max_tokens: int +) -> tuple[int, int, int]: + """Pick the integer (H, W) grid closest to the aspect ratio under the token cap. + + Replicates MuseGlimmerVisionEncoder._compute_grid_size + (modeling_muse_glimmer.py) so the processor needs no torch model import. + Returns (target_h, target_w, n_tokens). + """ + i_nph = img_h / patch_hw + i_npw = img_w / patch_hw + ratio = i_npw / i_nph if i_nph > 0 else 1.0 + if i_nph * i_npw > max_tokens: + i_nph = (max_tokens / ratio) ** 0.5 + i_npw = i_nph * ratio + candidates = list( + set( + itertools.product( + [math.floor(i_nph), math.ceil(i_nph)], + [math.floor(i_npw), math.ceil(i_npw)], + ) + ) + ) + candidates = [ + (nph, npw) + for nph, npw in candidates + if nph >= 1 and npw >= 1 and nph * npw <= max_tokens + ] + if not candidates: + candidates = [(max(1, round(i_nph)), max(1, round(i_npw)))] + nph, npw = min(candidates, key=lambda c: abs(c[0] / c[1] - img_h / img_w)) + return nph * patch_hw, npw * patch_hw, nph * npw + + +class MuseGlimmerImageProcessor(BaseImageProcessor): + """Resize + normalize MuseGlimmer images and compute patch-token counts. + + Variable-resolution: each image is resized to the grid that best matches its + aspect ratio under the per-image token cap, then normalized with mean/std 0.5. + Returns per-image tensors (not stacked) because MuseGlimmer consumes a list of + variable-size images. Video frames are handled by MuseGlimmerVideoProcessor. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + patch_size: int = 14, + downsample_factor: int = 2, + max_image_tokens: int = 4096, + image_mean: float = 0.5, + image_std: float = 0.5, + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = patch_size + self.downsample_factor = downsample_factor + self.max_image_tokens = max_image_tokens + self.image_mean = image_mean + self.image_std = image_std + + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: + # Functional normalize -- no stored Normalize object, so the processor + # stays JSON-serializable for to_dict() / save_pretrained. + return T.functional.normalize( + T.functional.to_tensor(image), + [self.image_mean] * 3, + [self.image_std] * 3, + ) + + # -- size computation (mirrors modeling_muse_glimmer.MuseGlimmerVisionEncoder) -- + def compute_image_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: + ph = self.patch_size * self.downsample_factor + return _grid_size(img_w, img_h, ph, self.max_image_tokens) + + # -- preprocessing --------------------------------------------------------- + def preprocess_image(self, image: Image.Image) -> tuple[torch.Tensor, int]: + """Return (pixel tensor [3, H, W], n_patch_tokens) for one image.""" + image = image.convert("RGB") + target_h, target_w, n_tokens = self.compute_image_size( + image.width, image.height + ) + image = image.resize((target_w, target_h), Image.LANCZOS) + return self._to_norm_tensor(image), n_tokens + + +# torchcodec is the training-faithful video decoder; other decoders (torchvision / +# PyAV) seek + color-convert differently and diverge from training. Optional at +# import time so text/image paths load without it; decode_video raises if used. +try: + import torchcodec +except Exception: + torchcodec = None + + +class MuseGlimmerVideoProcessor(BaseVideoProcessor): + """MuseGlimmer video preprocessing behind the standard HF + ``AutoVideoProcessor`` API. + + Exposes MuseGlimmer's training-faithful video handling as a first-class video + processor so you can call ``processor(videos="clip.mp4")`` and downstream + users can discover it via the standard HF ``AutoVideoProcessor`` API plus a + ``video_preprocessor_config.json``. It wraps: + + * torchcodec decode (matches training; other decoders diverge), + * uniform frame sampling to a whole multiple of ``patch_temporal``, + * real per-group PTS timestamps (rendered as ``Time: X.Xs`` by + MuseGlimmerProcessor), + * ``patch_temporal`` frame-grouping (frames cat on the channel axis -> + ``[patch_temporal * 3, H, W]``; the encoder detects video by channel count). + + It deliberately overrides ``preprocess`` instead of the BaseVideoProcessor fast + pipeline (group_videos_by_shape / smart_resize / single stacked tensor): MuseGlimmer + consumes a LIST of variable-size group tensors and needs the real per-group PTS, + neither of which the stacked fast path models. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + patch_size: int = 14, + downsample_factor: int = 2, + patch_temporal: int = 2, + max_video_frame_tokens: int = 144, + image_mean: float = 0.5, + image_std: float = 0.5, + video_num_frames: int = 96, + video_sampling_fps: float = 2.0, + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = patch_size + self.downsample_factor = downsample_factor + self.patch_temporal = patch_temporal + self.max_video_frame_tokens = max_video_frame_tokens + self.image_mean = image_mean + self.image_std = image_std + self.video_num_frames = video_num_frames + self.video_sampling_fps = video_sampling_fps + + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: + return T.functional.normalize( + T.functional.to_tensor(image), + [self.image_mean] * 3, + [self.image_std] * 3, + ) + + def decode_video(self, video_path: str) -> tuple[list[Image.Image], list[float]]: + """Sample frames + per-group PTS with torchcodec (the training decode path). + + ``timestamps[g]`` is the ACTUAL decoded PTS of the first frame in temporal + group ``g``; ``len(frames)`` is a whole multiple of ``patch_temporal``. + """ + if torchcodec is None: + raise RuntimeError( + "torchcodec is required for video decoding (it matches the training " + "decode path). See the Environment Setup steps in hf/README.md." + ) + pt = self.patch_temporal + reader = torchcodec.decoders.VideoDecoder(video_path) + total = len(reader) + assert reader.metadata.average_fps is not None, ( + f"Video has no FPS metadata: {video_path}" + ) + fps = reader.metadata.average_fps + assert self.video_sampling_fps and self.video_sampling_fps > 0, ( + f"video_sampling_fps must be positive, got {self.video_sampling_fps}" + ) + n = min( + int(total * self.video_sampling_fps / fps), self.video_num_frames, total + ) + n = max(pt, (n // pt) * pt) + n = min(n, total) + if n < pt: + raise ValueError( + f"Video has only {total} decodable frame(s) but needs at least " + f"{pt} (one temporal patch): {video_path}" + ) + indices = torch.linspace(0, total - 1, n).long().tolist() + frames: list[Image.Image] = [] + timestamps: list[float] = [] + for j, i in enumerate(indices): + fr = reader[i] + frames.append( + Image.fromarray(fr.data.permute(1, 2, 0).numpy()).convert("RGB") + ) + if j % pt == 0: + pts = getattr(fr, "pts_seconds", None) + timestamps.append(float(pts) if pts is not None else i / fps) + return frames, timestamps + + def compute_video_frame_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: + """Pick (H, W, tokens_per_frame) for a video frame under the per-frame cap. + + Mirrors MuseGlimmerImageProcessor.compute_image_size but uses the (smaller) + per-frame video token budget; shares the same ``_grid_size`` grid logic as + modeling_muse_glimmer.MuseGlimmerVisionEncoder. + """ + ph = self.patch_size * self.downsample_factor + return _grid_size(img_w, img_h, ph, self.max_video_frame_tokens) + + def _group_frames( + self, frames: list[Image.Image] + ) -> tuple[list[torch.Tensor], int, int]: + """Resize/normalize frames and cat ``patch_temporal`` frames per group. + + Returns (group_tensors, n_groups, tokens_per_group). + """ + pt = self.patch_temporal + if pt <= 0: + raise ValueError(f"patch_temporal must be positive, got {pt}") + if not frames: + raise ValueError("video must contain at least one frame") + if padding := (-len(frames)) % pt: + frames = [*frames, *([frames[-1]] * padding)] + first = frames[0].convert("RGB") + target_h, target_w, n_tokens = self.compute_video_frame_size( + first.width, first.height + ) + groups: list[torch.Tensor] = [] + for i in range(0, len(frames), pt): + grp = [ + self._to_norm_tensor( + frames[i + j] + .convert("RGB") + .resize((target_w, target_h), Image.LANCZOS) + ) + for j in range(pt) + ] + groups.append(torch.cat(grp, dim=0)) + return groups, len(groups), n_tokens + + @staticmethod + def _normalize_videos(videos) -> list: + """Normalize the ``videos`` arg to a list of per-video items. + + A video item is a path (str/Path) or a list of PIL frames. Accepts a + single video (one path, or one list of frames) or a list of those. + """ + if isinstance(videos, (str, Path)): + return [videos] + if videos and not isinstance(videos[0], (list, tuple, str, Path)): + return [videos] # a single list of PIL frames + return list(videos or []) + + def preprocess_one( + self, + video: str | Path | list[Image.Image], + timestamps: list[float] | None = None, + ) -> tuple[list[torch.Tensor], int, int, list[float]]: + """Decode (if a path) + group ONE video. + + Returns (group_tensors, n_groups, tokens_per_group, group_timestamps). This + is the single per-video implementation shared by ``preprocess`` (the HF + AutoVideoProcessor batch API) and ``MuseGlimmerProcessor.__call__`` (prompt + building), so the two paths cannot drift. For a path, frames + real + per-group PTS come from ``decode_video``; for pre-decoded frames the given + ``timestamps`` are used verbatim (empty if not supplied). + """ + if isinstance(video, (str, Path)): + frames, ts = self.decode_video(str(video)) + else: + frames = [f.convert("RGB") for f in video] + ts = list(timestamps) if timestamps is not None else [] + groups, n_groups, tokens_per_group = self._group_frames(frames) + return groups, n_groups, tokens_per_group, ts + + def preprocess( + self, + videos, + video_timestamps: list[list[float]] | None = None, + return_tensors: str | None = "pt", + **kwargs, + ) -> BatchFeature: + """Preprocess video file path(s) or pre-decoded frame list(s). + + Accepts a single video (path str, or list of PIL frames) or a list of + those. Returns a BatchFeature carrying, in video order: + * ``pixel_values`` -- flat list of [pt*3, H, W] group tensors, + * ``video_num_groups`` -- groups per video, + * ``video_tokens_per_group`` -- patch-token count per group per video, + * ``video_timestamps`` -- per-group PTS per video, + i.e. exactly what MuseGlimmerProcessor needs to build the + ``<|vid_start|> ( Time: X.Xs <|video|>*P [sep] )* <|vid_end|>`` block. + """ + videos = self._normalize_videos(videos) + pixel_values: list[torch.Tensor] = [] + num_groups: list[int] = [] + tokens_per_group: list[int] = [] + out_ts: list[list[float]] = [] + for idx, v in enumerate(videos): + ts_in = video_timestamps[idx] if video_timestamps else None + groups, ng, tpg, ts = self.preprocess_one(v, ts_in) + pixel_values += groups + num_groups.append(ng) + tokens_per_group.append(tpg) + out_ts.append(ts) + batch = BatchFeature( + data={ + "video_num_groups": num_groups, + "video_tokens_per_group": tokens_per_group, + "video_timestamps": out_ts, + }, + tensor_type=None, + ) + # Variable-size list -- BatchFeature would fail to stack into one tensor. + batch["pixel_values"] = pixel_values + return batch + + +class MuseGlimmerProcessor(ProcessorMixin): + """Bundle MuseGlimmerImageProcessor + MuseGlimmerVideoProcessor + tokenizer; + expand media sentinels into spans. + + Images go through ``image_processor`` (MuseGlimmerImageProcessor) and videos through + ``video_processor`` (MuseGlimmerVideoProcessor) -- one preprocessing implementation + each, no overlap. + """ + + attributes = ["image_processor", "video_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + video_processor_class = "AutoVideoProcessor" + tokenizer_class = "PreTrainedTokenizerFast" + + def __init__( + self, + image_processor=None, + video_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + if image_processor is None: + image_processor = MuseGlimmerImageProcessor() + if video_processor is None: + video_processor = MuseGlimmerVideoProcessor() + super().__init__( + image_processor, + video_processor, + tokenizer, + chat_template=chat_template or MUSE_GLIMMER_MM_CHAT_TEMPLATE, + **kwargs, + ) + + def _sid(self, token: str) -> int: + return self.tokenizer.convert_tokens_to_ids(token) + + def _image_block(self, n_tokens: int) -> list[int]: + return ( + [self._sid("<|image_start|>")] + + [self._sid("<|patch|>")] * n_tokens + + [self._sid("<|image_end|>")] + ) + + def _video_block( + self, + n_groups: int, + tokens_per_group: int, + timestamps: list[float] | None = None, + ) -> list[int]: + """Per-group ``Time: X.Xs`` + <|video|>*P, separated/terminated. + + ``timestamps`` (one per temporal group, in seconds) should be the ACTUAL + decoded frame times -- training renders the real factored frame PTS, not a + uniform grid. Falls back to ``g*patch_temporal/fps`` only when timestamps + are not supplied (e.g. callers that don't track frame times); that + approximation can differ from training by the per-frame jitter. + """ + vid = self._sid("<|video|>") + sep = self._sid("<|vid_frame_separator|>") + pt = self.video_processor.patch_temporal + fps = self.video_processor.video_sampling_fps + block = [self._sid("<|vid_start|>")] + for g in range(n_groups): + ts = timestamps[g] if timestamps is not None else g * pt / fps + block += self.tokenizer.encode(f"Time: {ts:.1f}s", add_special_tokens=False) + block += [vid] * tokens_per_group + block.append(sep if g < n_groups - 1 else self._sid("<|vid_end|>")) + return block + + def __call__( + self, + text: str | list[str] | None = None, + images: list[Image.Image] | None = None, + videos: list[list[Image.Image] | str] | None = None, + video_timestamps: list[list[float]] | None = None, + return_tensors: str | None = "pt", + **kwargs, + ) -> BatchFeature: + if text is None: + raise ValueError( + "`text` is required (use apply_chat_template to build it)." + ) + if isinstance(text, (list, tuple)): + if len(text) != 1: + raise ValueError( + "MuseGlimmerProcessor supports a single text sample per call." + ) + text = text[0] + + images = list(images or []) + videos = list(videos or []) + image_sentinel = self._sid(IMAGE_SENTINEL) + video_sentinel = self._sid(VIDEO_SENTINEL) + + # Preprocess media up front so we can expand sentinels in document order. + # Images via image_processor; videos via video_processor (each video item + # is a path -- decoded here -- or a pre-decoded list of PIL frames). + prepped_images = [self.image_processor.preprocess_image(im) for im in images] + prepped_videos = [ + self.video_processor.preprocess_one( + v, video_timestamps[i] if video_timestamps else None + ) + for i, v in enumerate(videos) + ] + + ids = self.tokenizer.encode(text, add_special_tokens=False) + n_img = sum(1 for t in ids if t == image_sentinel) + n_vid = sum(1 for t in ids if t == video_sentinel) + if n_img != len(prepped_images): + raise ValueError( + f"{n_img} image sentinel(s) in text but " + f"{len(prepped_images)} image(s) given." + ) + if n_vid != len(prepped_videos): + raise ValueError( + f"{n_vid} video sentinel(s) in text but " + f"{len(prepped_videos)} video(s) given." + ) + + out_ids: list[int] = [] + pixel_values: list[torch.Tensor] = [] + img_i = vid_i = 0 + for tid in ids: + if tid == image_sentinel: + tensor, n_tokens = prepped_images[img_i] + img_i += 1 + out_ids += self._image_block(n_tokens) + pixel_values.append(tensor) + elif tid == video_sentinel: + groups, n_groups, tokens_per_group, ts = prepped_videos[vid_i] + vid_i += 1 + out_ids += self._video_block(n_groups, tokens_per_group, ts or None) + pixel_values += groups + else: + out_ids.append(tid) + + data: dict = { + "input_ids": [out_ids], + "attention_mask": [[1] * len(out_ids)], + } + batch = BatchFeature(data=data, tensor_type=return_tensors) + # Keep pixel_values as a list of variable-size tensors (MuseGlimmerModel + # consumes a list); BatchFeature would fail to stack them into one tensor. + if pixel_values: + batch["pixel_values"] = pixel_values + return batch + + +__all__ = [ + "MuseGlimmerImageProcessor", + "MuseGlimmerVideoProcessor", + "MuseGlimmerProcessor", + "MUSE_GLIMMER_MM_CHAT_TEMPLATE", +] diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 626dd36dea..e773ecf305 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -80,6 +80,21 @@ class DFlashProposer(SpecDecodeBaseProposer): self.draft_model_config.hf_config ) + @override + def load_model(self, target_model: torch.nn.Module) -> None: + from vllm.model_executor.models.qwen3_dflash import ( + dflash_target_rope_is_neox_style, + ) + + draft_model_config = self.speculative_config.draft_model_config + assert draft_model_config is not None + # The drafter must rotate Q/K the way its target does. Take that from the + # built target before super() constructs the draft. + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style + super().load_model(target_model) + @override def _create_draft_vllm_config(self) -> VllmConfig: base = super()._create_draft_vllm_config() diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 478274a056..b9e6ea02f4 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -13,11 +13,19 @@ from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + from vllm.model_executor.models.qwen3_dflash import ( + dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, + ) speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + # The drafter must rotate Q/K the way its target does. Take that from the + # built target before super() constructs the draft. + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style # Select an attention backend that supports the drafter's attention: mixing # a non-causal layer onto a causal-only backend would fail. draft_vllm_config = replace( @@ -46,7 +54,10 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if hasattr(target_model, "get_language_model") else target_model ) - target_inner = target_language_model.model + # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language + # model, so get_language_model() already returns the inner module and has + # no .model of its own. + target_inner = getattr(target_language_model, "model", target_language_model) draft_inner = dflash_model.model # Skip embedding sharing under PP — each rank owns its own embedding.