@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 08ed2b9688
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 14:45 UTC

feat(models): implement Google Gemma 4 architecture support (MoE, Multimodal, Reasoning, Tool-Use) (#38826)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Signed-off-by: Luciano Martins <lucianomartins@google.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com>

Luciano Martins committed Apr 2, 2026 at 15:13 UTC 08ed2b9688b49b1d28c0d78edd71048d2b88c82b
20 files changed +5051 -1
tests/models/multimodal/generation/test_common.py
+16
@@ -394,6 +394,22 @@ VLM_TEST_SETTINGS = {
394 vllm_runner_kwargs={"mm_processor_kwargs": {"do_pan_and_scan": True}},
395 patch_hf_runner=model_utils.gemma3_patch_hf_runner,
396 ),
397 + "gemma4": VLMTestInfo(
398 + models=["google/gemma-4-E2B-it"],
399 + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
400 + prompt_formatter=lambda img_prompt: f"<bos><start_of_turn>user\n{img_prompt}<end_of_turn>\n<start_of_turn>model\n", # noqa: E501
401 + single_image_prompts=IMAGE_ASSETS.prompts(
402 + {
403 + "stop_sign": "What's the content in the center of the image?",
404 + "cherry_blossom": "What is the season?",
405 + }
406 + ),
407 + multi_image_prompt="Describe the two images in detail.",
408 + max_model_len=4096,
409 + max_num_seqs=2,
410 + auto_cls=AutoModelForImageTextToText,
411 + vllm_runner_kwargs={"limit_mm_per_prompt": {"image": 4}},
412 + ),
413 "granite_vision": VLMTestInfo(
414 models=["ibm-granite/granite-vision-3.3-2b"],
415 test_type=(VLMTestType.IMAGE),
tests/models/multimodal/processing/test_gemma4.py new
+44
@@ -0,0 +1,44 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.multimodal import MULTIMODAL_REGISTRY
7 +
8 +from ....conftest import ImageTestAssets
9 +from ...utils import build_model_context
10 +
11 +# TODO: to be updated to "google/gemma-4-e2b-it" once the models are available
12 +GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
13 +
14 +
15 +@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
16 +def test_limit_mm_per_prompt(
17 + image_assets: ImageTestAssets,
18 + model_id: str,
19 +):
20 + """Test that limit_mm_per_prompt accurately restricts multiple images."""
21 + # We only allow 1 image
22 + ctx = build_model_context(
23 + model_id,
24 + mm_processor_kwargs={},
25 + limit_mm_per_prompt={"image": 1},
26 + )
27 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
28 +
29 + # Provide 2 images in the prompt
30 + prompt = "<image><image>"
31 + # image_assets usually has multiple images
32 + images = [asset.pil_image for asset in image_assets][:2]
33 + if len(images) < 2:
34 + images = [images[0], images[0]]
35 +
36 + mm_data = {"image": images}
37 +
38 + # Expect ValueError when exceeding limit
39 + with pytest.raises(ValueError, match="At most 1 image"):
40 + processor(
41 + prompt,
42 + mm_items=processor.info.parse_mm_data(mm_data),
43 + hf_processor_mm_kwargs={},
44 + )
tests/models/registry.py
+8
@@ -277,6 +277,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
277 "google/gemma-2-9b", extras={"tiny": "google/gemma-2-2b-it"}
278 ),
279 "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),
280 + "Gemma4ForCausalLM": _HfExamplesInfo(
281 + "google/gemma-4-E2B-it",
282 + min_transformers_version="5.0.0",
283 + ),
284 "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"),
285 "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),
286 "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),
@@ -813,6 +817,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
817 ),
818 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
819 "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
820 + "Gemma4ForConditionalGeneration": _HfExamplesInfo(
821 + "google/gemma-4-E2B-it",
822 + min_transformers_version="5.5.0",
823 + ),
824 "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"),
825 "GlmAsrForConditionalGeneration": _HfExamplesInfo(
826 "zai-org/GLM-ASR-Nano-2512",
tests/reasoning/test_gemma4_reasoning_parser.py new
+196
@@ -0,0 +1,196 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from tests.reasoning.utils import run_reasoning_extraction
7 +from vllm.reasoning import ReasoningParser, ReasoningParserManager
8 +
9 +# Using mistral tokenizer as a generic mock since the actual model is not on HF
10 +from vllm.tokenizers.registry import get_tokenizer
11 +
12 +parser_name = "gemma4"
13 +
14 +
15 +@pytest.fixture(scope="module")
16 +def generic_tokenizer():
17 + return get_tokenizer("google/gemma-4-E2B-it")
18 +
19 +
20 +INVALID_SIMPLE_NONSTREAMING = {
21 + "output": "This is a reasoning section<channel|>This is the rest",
22 + "reasoning": "This is a reasoning section",
23 + "content": "This is the rest",
24 + "is_reasoning_end": True,
25 +}
26 +INVALID_SIMPLE_STREAMING = {
27 + "output": "This is a reasoning section<channel|>This is the rest",
28 + "reasoning": None,
29 + "content": "This is a reasoning sectionThis is the rest",
30 + "is_reasoning_end": True,
31 +}
32 +INVALID_COMPLETE_NONSTREAMING = {
33 + "output": "This is a reasoning section<channel|>",
34 + "reasoning": "This is a reasoning section",
35 + "content": None,
36 + "is_reasoning_end": True,
37 +}
38 +INVALID_COMPLETE_STREAMING = {
39 + "output": "This is a reasoning section<channel|>",
40 + "reasoning": None,
41 + "content": "This is a reasoning section",
42 + "is_reasoning_end": True,
43 +}
44 +NO_CONTENT = {
45 + "output": "<|channel>This is reasoning",
46 + "reasoning": "This is reasoning",
47 + "content": None,
48 + "is_reasoning_end": False,
49 +}
50 +NO_REASONING = {
51 + "output": "This is content",
52 + "reasoning": None,
53 + "content": "This is content",
54 + "is_reasoning_end": False,
55 +}
56 +REASONING_WITH_CHANNEL = {
57 + "output": "<|channel>This is a reasoning section<channel|>This is the rest",
58 + "reasoning": "This is a reasoning section",
59 + "content": "This is the rest",
60 + "is_reasoning_end": True,
61 +}
62 +COMPLETE_REASONING_WITH_CHANNEL = {
63 + "output": "<|channel>This is a reasoning section<channel|>",
64 + "reasoning": "This is a reasoning section",
65 + "content": None,
66 + "is_reasoning_end": True,
67 +}
68 +MULTIPLE_LINES_WITH_CHANNEL = {
69 + "output": "<|channel>This\nThat<channel|>This is the rest\nThat",
70 + "reasoning": "This\nThat",
71 + "content": "This is the rest\nThat",
72 + "is_reasoning_end": True,
73 +}
74 +CHANNEL_NO_END = {
75 + "output": "<|channel>This is a reasoning section",
76 + "reasoning": "This is a reasoning section",
77 + "content": None,
78 + "is_reasoning_end": False,
79 +}
80 +EMPTY = {
81 + "output": "",
82 + "reasoning": None,
83 + "content": "",
84 + "is_reasoning_end": False,
85 +}
86 +NEW_LINE_NONSTREAMING = {
87 + "output": (
88 + "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
89 + ),
90 + "reasoning": "This is a reasoning section",
91 + "content": "\nThis is the rest",
92 + "is_reasoning_end": True,
93 +}
94 +NEW_LINE_STREAMING = {
95 + "output": (
96 + "Before\n<|channel>This is a reasoning section<channel|>\nThis is the rest"
97 + ),
98 + "reasoning": "This is a reasoning section",
99 + "content": "Before\n\nThis is the rest",
100 + "is_reasoning_end": True,
101 +}
102 +
103 +TEST_CASES = [
104 + pytest.param(False, INVALID_SIMPLE_NONSTREAMING, id="invalid_simple"),
105 + pytest.param(True, INVALID_SIMPLE_STREAMING, id="invalid_simple_streaming"),
106 + pytest.param(False, INVALID_COMPLETE_NONSTREAMING, id="invalid_complete"),
107 + pytest.param(True, INVALID_COMPLETE_STREAMING, id="invalid_complete_streaming"),
108 + pytest.param(False, NO_CONTENT, id="no_content"),
109 + pytest.param(False, NO_REASONING, id="no_reasoning"),
110 + pytest.param(False, REASONING_WITH_CHANNEL, id="reasoning"),
111 + pytest.param(True, REASONING_WITH_CHANNEL, id="reasoning_streaming"),
112 + pytest.param(False, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning"),
113 + pytest.param(
114 + True, COMPLETE_REASONING_WITH_CHANNEL, id="complete_reasoning_streaming"
115 + ),
116 + pytest.param(False, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines"),
117 + pytest.param(True, MULTIPLE_LINES_WITH_CHANNEL, id="multiple_lines_streaming"),
118 + pytest.param(False, CHANNEL_NO_END, id="no_end"),
119 + pytest.param(True, CHANNEL_NO_END, id="no_end_streaming"),
120 + pytest.param(False, EMPTY, id="empty"),
121 + pytest.param(False, NEW_LINE_NONSTREAMING, id="new_line"),
122 + pytest.param(True, NEW_LINE_STREAMING, id="new_line_streaming"),
123 +]
124 +
125 +
126 +@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
127 +def test_gemma4_reasoning(
128 + streaming: bool,
129 + param_dict: dict,
130 + generic_tokenizer,
131 +):
132 + output = param_dict["output"]
133 +
134 + # Resolve token IDs dynamically from the real tokenizer
135 + vocab = generic_tokenizer.get_vocab()
136 + start_token_id = vocab["<|channel>"]
137 + end_token_id = vocab["<channel|>"]
138 +
139 + index_start = output.find("<|channel>")
140 + len_start = len("<|channel>")
141 + index_end = output.find("<channel|>")
142 + len_end = len("<channel|>")
143 +
144 + output_tokens = []
145 +
146 + def _encode(text: str) -> list[int]:
147 + if not text:
148 + return []
149 + # Handle both raw transformers and vLLM wrappers
150 + enc = getattr(generic_tokenizer, "tokenizer", generic_tokenizer)
151 + try:
152 + return enc.encode(text, add_special_tokens=False)
153 + except TypeError:
154 + return enc.encode(text)
155 +
156 + if index_start != -1:
157 + output_before = output[:index_start]
158 + output_tokens += _encode(output_before)
159 + output_tokens += [start_token_id]
160 +
161 + if index_end != -1:
162 + output_middle = output[index_start + len_start : index_end]
163 + output_after = output[index_end + len_end :]
164 + output_tokens += _encode(output_middle)
165 + output_tokens += [end_token_id]
166 + output_tokens += _encode(output_after)
167 + else:
168 + output_middle = output[index_start + len_start :]
169 + output_tokens += _encode(output_middle)
170 + elif index_end != -1:
171 + output_before = output[:index_end]
172 + output_after = output[index_end + len_end :]
173 + output_tokens += _encode(output_before)
174 + output_tokens += [end_token_id]
175 + output_tokens += _encode(output_after)
176 + else:
177 + output_tokens += _encode(output)
178 +
179 + parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
180 + generic_tokenizer
181 + )
182 +
183 + # We use the generic run_reasoning_extraction from utils
184 + # Use decode per token to get standard spaces instead of
185 + # SentencePiece space characters
186 + output_token_strings = [generic_tokenizer.decode([t]) for t in output_tokens]
187 + reasoning, content = run_reasoning_extraction(
188 + parser, output_token_strings, streaming=streaming
189 + )
190 +
191 + assert reasoning == param_dict["reasoning"]
192 + assert content == param_dict["content"]
193 +
194 + # Test is_reasoning_end
195 + is_reasoning_end = parser.is_reasoning_end(output_tokens)
196 + assert is_reasoning_end == param_dict["is_reasoning_end"]
tests/tool_parsers/test_gemma4_tool_parser.py new
+504
@@ -0,0 +1,504 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from typing import Any
6 +from unittest.mock import MagicMock
7 +
8 +import pytest
9 +
10 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
11 +from vllm.tool_parsers.gemma4_tool_parser import (
12 + TOOL_CALL_END,
13 + TOOL_CALL_START,
14 + Gemma4ToolParser,
15 + _parse_gemma4_args,
16 + _parse_gemma4_array,
17 +)
18 +
19 +# ---------------------------------------------------------------------------
20 +# Fixtures
21 +# ---------------------------------------------------------------------------
22 +
23 +
24 +@pytest.fixture
25 +def mock_tokenizer():
26 + tokenizer = MagicMock()
27 + tokenizer.encode.return_value = [1, 2, 3]
28 + # Include the tool call start token in the vocab for the parser
29 + tokenizer.get_vocab.return_value = {TOOL_CALL_START: 48, TOOL_CALL_END: 49}
30 + return tokenizer
31 +
32 +
33 +@pytest.fixture
34 +def parser(mock_tokenizer):
35 + return Gemma4ToolParser(mock_tokenizer)
36 +
37 +
38 +@pytest.fixture
39 +def mock_request():
40 + request = MagicMock(spec=ChatCompletionRequest)
41 + request.tools = []
42 + request.tool_choice = "auto"
43 + return request
44 +
45 +
46 +# ---------------------------------------------------------------------------
47 +# Unit tests for _parse_gemma4_args (shared parser logic)
48 +# ---------------------------------------------------------------------------
49 +
50 +
51 +class TestParseGemma4Args:
52 + def test_empty_string(self):
53 + assert _parse_gemma4_args("") == {}
54 +
55 + def test_whitespace_only(self):
56 + assert _parse_gemma4_args(" ") == {}
57 +
58 + def test_single_string_value(self):
59 + result = _parse_gemma4_args('location:<|"|>Paris<|"|>')
60 + assert result == {"location": "Paris"}
61 +
62 + def test_string_value_with_comma(self):
63 + result = _parse_gemma4_args('location:<|"|>Paris, France<|"|>')
64 + assert result == {"location": "Paris, France"}
65 +
66 + def test_multiple_string_values(self):
67 + result = _parse_gemma4_args(
68 + 'location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>'
69 + )
70 + assert result == {"location": "San Francisco", "unit": "celsius"}
71 +
72 + def test_integer_value(self):
73 + result = _parse_gemma4_args("count:42")
74 + assert result == {"count": 42}
75 +
76 + def test_float_value(self):
77 + result = _parse_gemma4_args("score:3.14")
78 + assert result == {"score": 3.14}
79 +
80 + def test_boolean_true(self):
81 + result = _parse_gemma4_args("flag:true")
82 + assert result == {"flag": True}
83 +
84 + def test_boolean_false(self):
85 + result = _parse_gemma4_args("flag:false")
86 + assert result == {"flag": False}
87 +
88 + def test_mixed_types(self):
89 + result = _parse_gemma4_args(
90 + 'name:<|"|>test<|"|>,count:42,active:true,score:3.14'
91 + )
92 + assert result == {
93 + "name": "test",
94 + "count": 42,
95 + "active": True,
96 + "score": 3.14,
97 + }
98 +
99 + def test_nested_object(self):
100 + result = _parse_gemma4_args('nested:{inner:<|"|>value<|"|>}')
101 + assert result == {"nested": {"inner": "value"}}
102 +
103 + def test_array_of_strings(self):
104 + result = _parse_gemma4_args('items:[<|"|>a<|"|>,<|"|>b<|"|>]')
105 + assert result == {"items": ["a", "b"]}
106 +
107 + def test_unterminated_string(self):
108 + """Unterminated strings should take everything after the delimiter."""
109 + result = _parse_gemma4_args('key:<|"|>unterminated')
110 + assert result == {"key": "unterminated"}
111 +
112 + def test_empty_value(self):
113 + """Key with no value after colon."""
114 + result = _parse_gemma4_args("key:")
115 + assert result == {"key": ""}
116 +
117 +
118 +class TestParseGemma4Array:
119 + def test_string_array(self):
120 + result = _parse_gemma4_array('<|"|>a<|"|>,<|"|>b<|"|>')
121 + assert result == ["a", "b"]
122 +
123 + def test_empty_array(self):
124 + result = _parse_gemma4_array("")
125 + assert result == []
126 +
127 + def test_bare_values(self):
128 + result = _parse_gemma4_array("42,true,3.14")
129 + assert result == [42, True, 3.14]
130 +
131 +
132 +# ---------------------------------------------------------------------------
133 +# Non-streaming extraction tests
134 +# ---------------------------------------------------------------------------
135 +
136 +
137 +class TestExtractToolCalls:
138 + def test_no_tool_calls(self, parser, mock_request):
139 + model_output = "Hello, how can I help you today?"
140 + result = parser.extract_tool_calls(model_output, mock_request)
141 +
142 + assert result.tools_called is False
143 + assert result.tool_calls == []
144 + assert result.content == model_output
145 +
146 + def test_single_tool_call(self, parser, mock_request):
147 + model_output = (
148 + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}<tool_call|>'
149 + )
150 + result = parser.extract_tool_calls(model_output, mock_request)
151 +
152 + assert result.tools_called is True
153 + assert len(result.tool_calls) == 1
154 + assert result.tool_calls[0].function.name == "get_weather"
155 + args = json.loads(result.tool_calls[0].function.arguments)
156 + assert args == {"location": "London"}
157 +
158 + def test_multiple_arguments(self, parser, mock_request):
159 + model_output = (
160 + "<|tool_call>call:get_weather{"
161 + 'location:<|"|>San Francisco<|"|>,'
162 + 'unit:<|"|>celsius<|"|>}'
163 + "<tool_call|>"
164 + )
165 + result = parser.extract_tool_calls(model_output, mock_request)
166 +
167 + assert result.tools_called is True
168 + assert len(result.tool_calls) == 1
169 + assert result.tool_calls[0].function.name == "get_weather"
170 + args = json.loads(result.tool_calls[0].function.arguments)
171 + assert args == {"location": "San Francisco", "unit": "celsius"}
172 +
173 + def test_text_before_tool_call(self, parser, mock_request):
174 + model_output = (
175 + "Let me check the weather for you. "
176 + '<|tool_call>call:get_weather{location:<|"|>Paris<|"|>}'
177 + "<tool_call|>"
178 + )
179 + result = parser.extract_tool_calls(model_output, mock_request)
180 +
181 + assert result.tools_called is True
182 + assert result.content == "Let me check the weather for you."
183 + assert len(result.tool_calls) == 1
184 + assert result.tool_calls[0].function.name == "get_weather"
185 +
186 + def test_multiple_tool_calls(self, parser, mock_request):
187 + model_output = (
188 + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}'
189 + "<tool_call|>"
190 + '<|tool_call>call:get_time{location:<|"|>London<|"|>}'
191 + "<tool_call|>"
192 + )
193 + result = parser.extract_tool_calls(model_output, mock_request)
194 +
195 + assert result.tools_called is True
196 + assert len(result.tool_calls) == 2
197 + assert result.tool_calls[0].function.name == "get_weather"
198 + assert result.tool_calls[1].function.name == "get_time"
199 +
200 + def test_nested_arguments(self, parser, mock_request):
201 + model_output = (
202 + "<|tool_call>call:complex_function{"
203 + 'nested:{inner:<|"|>value<|"|>},'
204 + 'list:[<|"|>a<|"|>,<|"|>b<|"|>]}'
205 + "<tool_call|>"
206 + )
207 + result = parser.extract_tool_calls(model_output, mock_request)
208 +
209 + assert result.tools_called is True
210 + assert len(result.tool_calls) == 1
211 + assert result.tool_calls[0].function.name == "complex_function"
212 + args = json.loads(result.tool_calls[0].function.arguments)
213 + assert args == {"nested": {"inner": "value"}, "list": ["a", "b"]}
214 +
215 + def test_tool_call_with_number_and_boolean(self, parser, mock_request):
216 + model_output = (
217 + "<|tool_call>call:set_status{"
218 + "is_active:true,"
219 + "count:42,"
220 + "score:3.14}"
221 + "<tool_call|>"
222 + )
223 + result = parser.extract_tool_calls(model_output, mock_request)
224 +
225 + assert result.tools_called is True
226 + assert len(result.tool_calls) == 1
227 + assert result.tool_calls[0].function.name == "set_status"
228 + args = json.loads(result.tool_calls[0].function.arguments)
229 + assert args == {"is_active": True, "count": 42, "score": 3.14}
230 +
231 + def test_incomplete_tool_call(self, parser, mock_request):
232 + model_output = '<|tool_call>call:get_weather{location:<|"|>London'
233 + result = parser.extract_tool_calls(model_output, mock_request)
234 +
235 + # Incomplete — no <tool_call|> end marker, regex won't match
236 + assert result.tools_called is False
237 + assert result.content == model_output
238 +
239 + def test_hyphenated_function_name(self, parser, mock_request):
240 + """Ensure function names with hyphens are parsed correctly."""
241 + model_output = (
242 + '<|tool_call>call:get-weather{location:<|"|>London<|"|>}<tool_call|>'
243 + )
244 + result = parser.extract_tool_calls(model_output, mock_request)
245 +
246 + assert result.tools_called is True
247 + assert result.tool_calls[0].function.name == "get-weather"
248 +
249 + def test_dotted_function_name(self, parser, mock_request):
250 + """Ensure function names with dots are parsed correctly."""
251 + model_output = (
252 + '<|tool_call>call:weather.get{location:<|"|>London<|"|>}<tool_call|>'
253 + )
254 + result = parser.extract_tool_calls(model_output, mock_request)
255 +
256 + assert result.tools_called is True
257 + assert result.tool_calls[0].function.name == "weather.get"
258 +
259 + def test_no_arguments(self, parser, mock_request):
260 + """Tool calls with empty arguments."""
261 + model_output = "<|tool_call>call:get_status{}<tool_call|>"
262 + result = parser.extract_tool_calls(model_output, mock_request)
263 +
264 + assert result.tools_called is True
265 + assert result.tool_calls[0].function.name == "get_status"
266 + args = json.loads(result.tool_calls[0].function.arguments)
267 + assert args == {}
268 +
269 +
270 +# ---------------------------------------------------------------------------
271 +# Streaming extraction tests
272 +# ---------------------------------------------------------------------------
273 +
274 +
275 +class TestStreamingExtraction:
276 + """Tests for the streaming tool call extraction.
277 +
278 + These simulate the token-by-token streaming that vLLM performs,
279 + feeding incremental text to extract_tool_calls_streaming() and
280 + verifying that the accumulated argument deltas form valid JSON.
281 + """
282 +
283 + def _simulate_streaming(
284 + self, parser: Gemma4ToolParser, mock_request: Any, chunks: list[str]
285 + ) -> list[tuple[Any, str]]:
286 + """Feed chunks through the streaming parser and collect results.
287 +
288 + Returns a list of (delta_message, accumulated_text) tuples.
289 + """
290 + results: list[tuple[Any, str]] = []
291 + previous_text: str = ""
292 + previous_token_ids: list[int] = []
293 +
294 + for chunk in chunks:
295 + current_text = previous_text + chunk
296 + # Use token ID 48 for tool_call start, 49 for end, 0 otherwise
297 + delta_token_ids: list[int] = []
298 + if TOOL_CALL_START in chunk:
299 + delta_token_ids.append(48)
300 + elif TOOL_CALL_END in chunk:
301 + delta_token_ids.append(49)
302 + else:
303 + delta_token_ids.append(0)
304 +
305 + current_token_ids = previous_token_ids + delta_token_ids
306 +
307 + delta = parser.extract_tool_calls_streaming(
308 + previous_text=previous_text,
309 + current_text=current_text,
310 + delta_text=chunk,
311 + previous_token_ids=tuple(previous_token_ids),
312 + current_token_ids=tuple(current_token_ids),
313 + delta_token_ids=tuple(delta_token_ids),
314 + request=mock_request,
315 + )
316 + results.append((delta, current_text))
317 + previous_text = current_text
318 + previous_token_ids = list(current_token_ids)
319 +
320 + return results
321 +
322 + def _collect_arguments(self, results):
323 + """Collect all argument deltas from streaming results into one string."""
324 + args_text = ""
325 + for delta, _ in results:
326 + if delta and delta.tool_calls:
327 + for tc in delta.tool_calls:
328 + func = tc.function if isinstance(tc.function, dict) else tc.function
329 + if isinstance(func, dict):
330 + arg = func.get("arguments", "")
331 + else:
332 + arg = getattr(func, "arguments", "") or ""
333 + if arg:
334 + args_text += arg
335 + return args_text
336 +
337 + def _collect_function_name(self, results):
338 + """Extract the function name from streaming results."""
339 + for delta, _ in results:
340 + if delta and delta.tool_calls:
341 + for tc in delta.tool_calls:
342 + func = tc.function if isinstance(tc.function, dict) else tc.function
343 + if isinstance(func, dict):
344 + name = func.get("name")
345 + else:
346 + name = getattr(func, "name", None)
347 + if name:
348 + return name
349 + return None
350 +
351 + def test_basic_streaming_single_tool(self, parser, mock_request):
352 + """Simulate the exact streaming scenario from the bug report.
353 +
354 + Model generates:
355 + <|tool_call>call:get_weather{location:<|"|>Paris, France<|"|>}<tool_call|>
356 +
357 + Expected: arguments should be valid JSON {"location": "Paris, France"}
358 + """
359 + chunks = [
360 + "<|tool_call>",
361 + "call:get_weather{",
362 + 'location:<|"|>Paris',
363 + ", France",
364 + '<|"|>}',
365 + "<tool_call|>",
366 + ]
367 +
368 + results = self._simulate_streaming(parser, mock_request, chunks)
369 +
370 + # Verify function name
371 + name = self._collect_function_name(results)
372 + assert name == "get_weather", f"Expected 'get_weather', got '{name}'"
373 +
374 + # Verify arguments form valid JSON
375 + args_text = self._collect_arguments(results)
376 + assert args_text, "No arguments were streamed"
377 + parsed_args = json.loads(args_text)
378 + assert parsed_args == {"location": "Paris, France"}
379 +
380 + def test_streaming_multi_arg(self, parser, mock_request):
381 + """Streaming with multiple arguments."""
382 + chunks = [
383 + "<|tool_call>",
384 + "call:get_weather{",
385 + 'location:<|"|>Tokyo<|"|>,',
386 + 'unit:<|"|>celsius<|"|>}',
387 + "<tool_call|>",
388 + ]
389 +
390 + results = self._simulate_streaming(parser, mock_request, chunks)
391 +
392 + name = self._collect_function_name(results)
393 + assert name == "get_weather"
394 +
395 + args_text = self._collect_arguments(results)
396 + assert args_text
397 + parsed_args = json.loads(args_text)
398 + assert parsed_args == {"location": "Tokyo", "unit": "celsius"}
399 +
400 + def test_streaming_no_extra_brace(self, parser, mock_request):
401 + """Verify the closing } is NOT leaked into arguments (Bug #2)."""
402 + chunks = [
403 + "<|tool_call>",
404 + "call:get_weather{",
405 + 'location:<|"|>London<|"|>}',
406 + "<tool_call|>",
407 + ]
408 +
409 + results = self._simulate_streaming(parser, mock_request, chunks)
410 + args_text = self._collect_arguments(results)
411 + assert args_text
412 +
413 + # The args text must be valid JSON (no extra })
414 + parsed = json.loads(args_text)
415 + assert parsed == {"location": "London"}
416 +
417 + # Specifically assert no double-brace
418 + assert args_text.count("}") <= 1, (
419 + f"Arguments contain extra closing brace: {args_text!r}"
420 + )
421 +
422 + def test_streaming_no_unquoted_keys(self, parser, mock_request):
423 + """Verify keys are properly quoted in JSON (Bug #1)."""
424 + chunks = [
425 + "<|tool_call>",
426 + "call:get_weather{",
427 + 'location:<|"|>Paris<|"|>}',
428 + "<tool_call|>",
429 + ]
430 +
431 + results = self._simulate_streaming(parser, mock_request, chunks)
432 + args_text = self._collect_arguments(results)
433 +
434 + # Must start with { and contain quoted key
435 + assert args_text.lstrip().startswith("{"), (
436 + f"Arguments don't start with '{{': {args_text!r}"
437 + )
438 + assert '"location"' in args_text, (
439 + f"Key 'location' not properly quoted: {args_text!r}"
440 + )
441 +
442 + def test_streaming_name_no_call_prefix(self, parser, mock_request):
443 + """Verify function name has no 'call:' prefix."""
444 + chunks = [
445 + "<|tool_call>",
446 + "call:get_weather{",
447 + 'location:<|"|>Paris<|"|>}',
448 + "<tool_call|>",
449 + ]
450 +
451 + results = self._simulate_streaming(parser, mock_request, chunks)
452 + name = self._collect_function_name(results)
453 + assert name == "get_weather"
454 + assert not name.startswith("call:"), f"Name has 'call:' prefix: {name!r}"
455 +
456 + def test_streaming_text_before_tool_call(self, parser, mock_request):
457 + """Text before tool call should be emitted as content."""
458 + chunks = [
459 + "Let me check ",
460 + "the weather. ",
461 + "<|tool_call>",
462 + "call:get_weather{",
463 + 'location:<|"|>London<|"|>}',
464 + "<tool_call|>",
465 + ]
466 +
467 + results = self._simulate_streaming(parser, mock_request, chunks)
468 +
469 + # First chunks should be content
470 + content_parts = []
471 + for delta, _ in results:
472 + if delta and delta.content:
473 + content_parts.append(delta.content)
474 +
475 + assert "".join(content_parts).strip().startswith("Let me check")
476 +
477 + def test_streaming_numeric_args(self, parser, mock_request):
478 + """Streaming with numeric and boolean argument values."""
479 + chunks = [
480 + "<|tool_call>",
481 + "call:set_config{",
482 + "count:42,",
483 + "active:true}",
484 + "<tool_call|>",
485 + ]
486 +
487 + results = self._simulate_streaming(parser, mock_request, chunks)
488 + args_text = self._collect_arguments(results)
489 + if args_text:
490 + parsed_args = json.loads(args_text)
491 + assert parsed_args["count"] == 42
492 + assert parsed_args["active"] is True
493 +
494 + def test_streaming_empty_args(self, parser, mock_request):
495 + """Tool call with no arguments."""
496 + chunks = [
497 + "<|tool_call>",
498 + "call:get_status{}",
499 + "<tool_call|>",
500 + ]
501 +
502 + results = self._simulate_streaming(parser, mock_request, chunks)
503 + name = self._collect_function_name(results)
504 + assert name == "get_status"
vllm/model_executor/layers/rotary_embedding/__init__.py
+12
@@ -12,6 +12,7 @@ from .dual_chunk_rope import DualChunkRotaryEmbedding
12 from .dynamic_ntk_alpha_rope import DynamicNTKAlphaRotaryEmbedding
13 from .dynamic_ntk_scaling_rope import DynamicNTKScalingRotaryEmbedding
14 from .fope import FourierRotaryEmbedding
15 +from .gemma4_rope import Gemma4RotaryEmbedding
16 from .linear_scaling_rope import LinearScalingRotaryEmbedding
17 from .llama3_rope import Llama3RotaryEmbedding
18 from .llama4_vision_rope import Llama4VisionRotaryEmbedding
@@ -134,6 +135,17 @@ def get_rope(
135 is_neox_style,
136 dtype,
137 )
138 + elif scaling_type == "proportional":
139 + # Proportional RoPE is used by Gemma4 for global (full) attention.
140 + # Gemma4 uses a sparse/fractional RoPE with cross-mixing between halves.
141 + rotary_emb = Gemma4RotaryEmbedding(
142 + head_size,
143 + rotary_dim,
144 + max_position,
145 + base,
146 + is_neox_style,
147 + dtype,
148 + )
149 elif scaling_type == "llama3":
150 scaling_factor = rope_parameters["factor"]
151 low_freq_factor = rope_parameters["low_freq_factor"]
vllm/model_executor/layers/rotary_embedding/gemma4_rope.py new
+84
@@ -0,0 +1,84 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gemma4-specific Rotary Positional Embeddings (proportional scaling).
4 +
5 +Gemma4 uses "proportional" RoPE which computes inv_freq frequencies scaled
6 +by head_dim (not rotary_dim), and zero-pads for non-rotated dimensions when
7 +partial_rotary_factor < 1. The actual rotation uses standard neox-style
8 +rotate_half, matching HF transformers' apply_rotary_pos_emb.
9 +"""
10 +
11 +import torch
12 +
13 +from .base import RotaryEmbedding
14 +
15 +
16 +class Gemma4RotaryEmbedding(RotaryEmbedding):
17 + """Gemma4 proportional RoPE.
18 +
19 + Extends RotaryEmbedding (which provides standard neox-style rotation
20 + via ops.rotary_embedding CUDA kernel) but overrides the inv_freq
21 + computation to match HF's _compute_proportional_rope_parameters:
22 + - Frequency exponents use head_dim (not rotary_dim) as denominator
23 + - Non-rotated dims are zero-padded (cos=1, sin=0 = identity rotation)
24 +
25 + When partial_rotary_factor=1.0 (the default for some variants), ALL dims are
26 + rotated and this is equivalent to standard RotaryEmbedding with
27 + head_dim-scaled frequencies.
28 + """
29 +
30 + def __init__(
31 + self,
32 + head_size: int,
33 + rotary_dim: int,
34 + max_position_embeddings: int,
35 + base: float,
36 + is_neox_style: bool,
37 + dtype: torch.dtype,
38 + ) -> None:
39 + # Number of rotation angle pairs (from partial_rotary_factor)
40 + self.rope_angles = rotary_dim // 2
41 + # Non-rotated angle pairs per half
42 + self.nope_angles = (head_size // 2) - self.rope_angles
43 +
44 + # Important: set rotary_dim = head_size so the base class's
45 + # forward_static applies rotation to ALL dims of the cos/sin cache.
46 + # The non-rotated dims will have cos=1, sin=0 (identity) thanks
47 + # to our _compute_inv_freq zero-padding.
48 + super().__init__(
49 + head_size,
50 + head_size, # rotary_dim = head_size (full application)
51 + max_position_embeddings,
52 + base,
53 + is_neox_style,
54 + dtype,
55 + )
56 +
57 + def _compute_inv_freq(self, base: float) -> torch.Tensor:
58 + """Compute frequencies matching HF proportional RoPE.
59 +
60 + Key difference from base: exponent denominator is head_size (not
61 + rotary_dim), and non-rotated dims are zero-padded.
62 + """
63 + # HF formula: base ** (arange(0, 2*rope_angles, 2) / head_dim)
64 + freq_exponents = (
65 + torch.arange(0, 2 * self.rope_angles, 2, dtype=torch.float) / self.head_size
66 + )
67 + inv_freq = 1.0 / (base**freq_exponents)
68 +
69 + # Zero-pad for non-rotated dims (identity rotation: cos=1, sin=0)
70 + if self.nope_angles > 0:
71 + inv_freq = torch.cat(
72 + [
73 + inv_freq,
74 + torch.zeros(self.nope_angles, dtype=torch.float),
75 + ]
76 + )
77 + return inv_freq
78 +
79 + def extra_repr(self) -> str:
80 + s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
81 + s += f", rope_angles={self.rope_angles}, nope_angles={self.nope_angles}"
82 + s += f", max_position_embeddings={self.max_position_embeddings}"
83 + s += f", base={self.base}, is_neox_style={self.is_neox_style}"
84 + return s
vllm/model_executor/models/config.py
+55
@@ -9,6 +9,7 @@ from vllm.utils.math_utils import round_up
9 if TYPE_CHECKING:
10 from vllm.config import ModelConfig, VllmConfig
11
12 +
13 logger = init_logger(__name__)
14
15
@@ -52,6 +53,58 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
53 hf_config.is_causal = not hf_config.use_bidirectional_attention
54
55
56 +class Gemma4Config(VerifyAndUpdateConfig):
57 + @staticmethod
58 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
59 + """Force unified attention backend for models with heterogeneous
60 + head dimensions.
61 +
62 + Some Gemma4 variants use different head dimensions for
63 + sliding window (head_dim) vs full attention (global_head_dim) layers.
64 + When global_head_dim > 256, FlashAttention rejects those layers
65 + (head_size <= 256 kernel limit), causing vLLM to select a different
66 + backend for each layer type. This mixed-backend execution produces
67 + numerical divergence and output corruption.
68 +
69 + The fix detects heterogeneous head dimensions from the model config
70 + and forces TRITON_ATTN (which has no head_size ceiling) for all
71 + layers when the user hasn't explicitly chosen a backend.
72 +
73 + TODO: Heterogeneous head_sizes (head_dim != global_head_dim)
74 + require NixlConnector changes to support per-layer KV transfer
75 + with different head dimensions for prefill-decode disaggregation.
76 + """
77 + hf_text_config = vllm_config.model_config.hf_text_config
78 + head_dim = getattr(hf_text_config, "head_dim", None)
79 + global_head_dim = getattr(hf_text_config, "global_head_dim", None)
80 +
81 + # Only force Triton when head dimensions actually differ AND the
82 + # larger one exceeds FlashAttention's kernel limit (head_size <= 256).
83 + # This avoids unnecessary backend forcing on smaller models where
84 + # the config carries global_head_dim but all layers can still use
85 + # the same FA backend.
86 + max_head_dim = max(head_dim or 0, global_head_dim or 0)
87 + if (
88 + head_dim is not None
89 + and global_head_dim is not None
90 + and head_dim != global_head_dim
91 + and max_head_dim > 256
92 + and vllm_config.attention_config.backend is None
93 + ):
94 + from vllm.v1.attention.backends.registry import (
95 + AttentionBackendEnum,
96 + )
97 +
98 + vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
99 + logger.info(
100 + "Gemma4 model has heterogeneous head dimensions "
101 + "(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN "
102 + "backend to prevent mixed-backend numerical divergence.",
103 + head_dim,
104 + global_head_dim,
105 + )
106 +
107 +
108 class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
109 @staticmethod
110 def verify_and_update_config(vllm_config: "VllmConfig") -> None:
@@ -533,6 +586,8 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
586 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
587 "FalconMambaForCausalLM": MambaModelConfig,
588 "Gemma3TextModel": Gemma3TextModelConfig,
589 + "Gemma4ForCausalLM": Gemma4Config,
590 + "Gemma4ForConditionalGeneration": Gemma4Config,
591 "GptOssForCausalLM": GptOssForCausalLMConfig,
592 "GteModel": SnowflakeGteNewModelConfig,
593 "GteNewForSequenceClassification": GteNewModelConfig,
vllm/model_executor/models/gemma4.py new
+1239
@@ -0,0 +1,1239 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 The vLLM team.
4 +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
5 +#
6 +#
7 +
8 +# Licensed under the Apache License, Version 2.0 (the "License");
9 +# you may not use this file except in compliance with the License.
10 +# You may obtain a copy of the License at
11 +#
12 +# http://www.apache.org/licenses/LICENSE-2.0
13 +#
14 +# Unless required by applicable law or agreed to in writing, software
15 +# distributed under the License is distributed on an "AS IS" BASIS,
16 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 +# See the License for the specific language governing permissions and
18 +# limitations under the License.
19 +"""Gemma 4 model implementation for vLLM."""
20 +
21 +from collections.abc import Iterable
22 +from itertools import islice
23 +
24 +import regex as re
25 +import torch
26 +from torch import nn
27 +
28 +from vllm.compilation.decorators import support_torch_compile
29 +from vllm.config import CacheConfig, VllmConfig
30 +from vllm.distributed import (
31 + get_pp_group,
32 + get_tensor_model_parallel_rank,
33 + get_tensor_model_parallel_world_size,
34 +)
35 +from vllm.logger import init_logger
36 +from vllm.model_executor.layers.activation import GeluAndMul
37 +from vllm.model_executor.layers.attention import Attention
38 +from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
39 +from vllm.model_executor.layers.layernorm import RMSNorm
40 +from vllm.model_executor.layers.linear import (
41 + ColumnParallelLinear,
42 + MergedColumnParallelLinear,
43 + QKVParallelLinear,
44 + ReplicatedLinear,
45 + RowParallelLinear,
46 +)
47 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
48 +from vllm.model_executor.layers.quantization import QuantizationConfig
49 +from vllm.model_executor.layers.rotary_embedding import get_rope
50 +from vllm.model_executor.layers.vocab_parallel_embedding import (
51 + ParallelLMHead,
52 + VocabParallelEmbedding,
53 +)
54 +from vllm.model_executor.model_loader.weight_utils import (
55 + default_weight_loader,
56 + maybe_remap_kv_scale_name,
57 +)
58 +from vllm.sequence import IntermediateTensors
59 +
60 +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
61 +from .utils import (
62 + AutoWeightsLoader,
63 + extract_layer_index,
64 + is_pp_missing_parameter,
65 + make_layers,
66 + maybe_prefix,
67 +)
68 +
69 +logger = init_logger(__name__)
70 +
71 +
72 +def _get_text_config(config):
73 + """Dereference text_config if config is a nested Gemma4Config.
74 +
75 + Gemma4 checkpoints use architectures=["Gemma4ForConditionalGeneration"]
76 + which yields a Gemma4Config with nested text_config. This function
77 + transparently returns the text config regardless of nesting.
78 + """
79 + if hasattr(config, "text_config"):
80 + return config.text_config
81 + return config
82 +
83 +
84 +class Gemma4MLP(nn.Module):
85 + def __init__(
86 + self,
87 + hidden_size: int,
88 + intermediate_size: int,
89 + hidden_activation: str,
90 + quant_config: QuantizationConfig | None = None,
91 + prefix: str = "",
92 + ) -> None:
93 + super().__init__()
94 + self.gate_up_proj = MergedColumnParallelLinear(
95 + hidden_size,
96 + [intermediate_size] * 2,
97 + bias=False,
98 + quant_config=quant_config,
99 + prefix=f"{prefix}.gate_up_proj",
100 + )
101 + self.down_proj = RowParallelLinear(
102 + intermediate_size,
103 + hidden_size,
104 + bias=False,
105 + quant_config=quant_config,
106 + prefix=f"{prefix}.down_proj",
107 + )
108 + if hidden_activation != "gelu_pytorch_tanh":
109 + raise ValueError(
110 + "Gemma4 uses `gelu_pytorch_tanh` as the hidden activation "
111 + "function. Please set `hidden_act` and `hidden_activation` to "
112 + "`gelu_pytorch_tanh`."
113 + )
114 + self.act_fn = GeluAndMul(approximate="tanh")
115 +
116 + def forward(self, x: torch.Tensor) -> torch.Tensor:
117 + gate_up, _ = self.gate_up_proj(x)
118 + x = self.act_fn(gate_up)
119 + x, _ = self.down_proj(x)
120 + return x
121 +
122 +
123 +class Gemma4Router(nn.Module):
124 + """Router for Gemma4 MoE that preprocesses input before projection.
125 +
126 + Applies RMSNorm (no learned weight), root_size scaling
127 + (hidden_size^{-0.5}), then a learned per-dimension scale before
128 + projecting to expert logits.
129 +
130 + This preprocessing is applied ONLY to the router's input, not to
131 + the expert MLPs' input.
132 + """
133 +
134 + def __init__(
135 + self,
136 + config,
137 + quant_config: QuantizationConfig | None = None,
138 + prefix: str = "",
139 + ) -> None:
140 + super().__init__()
141 + self.hidden_size = config.hidden_size
142 +
143 + # RMSNorm without learned weight — pure normalization only
144 + self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps, has_weight=False)
145 + # Per-dimension learned scale, applied after norm + root_size
146 + self.scale = nn.Parameter(torch.ones(self.hidden_size))
147 + # Constant 1/sqrt(hidden_size) scaling factor
148 + self.register_buffer(
149 + "root_size",
150 + torch.tensor(self.hidden_size**-0.5),
151 + persistent=False,
152 + )
153 + # Project to expert logits; replicated across TP for consistent routing
154 + # GateLinear supports bf16 W/A → fp32 output, which is important
155 + # because the topk kernel often needs fp32 for stable routing.
156 + self.proj = GateLinear(
157 + self.hidden_size,
158 + config.num_experts,
159 + bias=False,
160 + out_dtype=torch.float32,
161 + prefix=f"{prefix}.proj",
162 + )
163 +
164 + def forward(self, x: torch.Tensor) -> torch.Tensor:
165 + """Returns raw router logits [T, E]."""
166 + x = self.norm(x)
167 + x = x * self.root_size.to(x.dtype)
168 + x = x * self.scale.to(x.dtype)
169 + router_logits, _ = self.proj(x)
170 + return router_logits
171 +
172 +
173 +class Gemma4MoE(nn.Module):
174 + """Mixture of Experts for Gemma4 using vLLM's FusedMoE.
175 +
176 + Wraps FusedMoE with custom routing. The router projection is
177 + external (Gemma4Router) — this class only handles expert dispatch.
178 +
179 + Gemma4 routing: softmax over ALL experts → top-k → renormalize.
180 + per_expert_scale is folded into routing weights for mathematical
181 + correctness with FusedMoE's fused kernel.
182 + """
183 +
184 + def __init__(
185 + self,
186 + config,
187 + quant_config: QuantizationConfig | None = None,
188 + prefix: str = "",
189 + ) -> None:
190 + super().__init__()
191 + self.hidden_size = config.hidden_size
192 + self.num_experts = config.num_experts
193 +
194 + # Per-expert output scale folded into routing weights so that
195 + # FusedMoE's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
196 + self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts))
197 +
198 + # Gemma4 routing: softmax over ALL experts → top-k → renormalize.
199 + # FusedMoE's built-in fused_topk scopes softmax differently, so
200 + # a custom routing function is needed for numerical correctness.
201 + per_expert_scale = self.per_expert_scale
202 +
203 + def routing_function(
204 + hidden_states: torch.Tensor,
205 + gating_output: torch.Tensor,
206 + topk: int,
207 + renormalize: bool,
208 + ) -> tuple[torch.Tensor, torch.Tensor]:
209 + _, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
210 + router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1)
211 + indicator = torch.nn.functional.one_hot(
212 + topk_ids, num_classes=gating_output.size(-1)
213 + ).sum(dim=-2)
214 + gate_weights = indicator * router_probabilities
215 + renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True)
216 + renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0)
217 + dispatch_weights = gate_weights / renorm_factor
218 +
219 + topk_weights = dispatch_weights.gather(1, topk_ids)
220 +
221 + # Fold per_expert_scale into routing weights
222 + expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype)
223 + topk_weights = topk_weights * expert_scales
224 + return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
225 +
226 + # FusedMoE experts with custom Gemma4 routing
227 + self.experts = FusedMoE(
228 + num_experts=config.num_experts,
229 + top_k=config.top_k_experts,
230 + hidden_size=config.hidden_size,
231 + intermediate_size=getattr(
232 + config,
233 + "moe_intermediate_size",
234 + getattr(config, "expert_intermediate_size", None),
235 + ),
236 + reduce_results=True,
237 + renormalize=True,
238 + quant_config=quant_config,
239 + prefix=f"{prefix}.experts",
240 + custom_routing_function=routing_function,
241 + activation="gelu",
242 + )
243 +
244 + def forward(self, x: torch.Tensor, router_logits: torch.Tensor) -> torch.Tensor:
245 + return self.experts(x, router_logits)
246 +
247 +
248 +class Gemma4Attention(nn.Module):
249 + def __init__(
250 + self,
251 + config,
252 + hidden_size: int,
253 + num_heads: int,
254 + num_kv_heads: int,
255 + head_dim: int,
256 + max_position_embeddings: int,
257 + use_k_eq_v: bool = False,
258 + cache_config: CacheConfig | None = None,
259 + quant_config: QuantizationConfig | None = None,
260 + attn_logits_soft_cap: float | None = None,
261 + prefix: str = "",
262 + ) -> None:
263 + super().__init__()
264 + self.config = config
265 + self.hidden_size = hidden_size
266 + self.use_k_eq_v = use_k_eq_v
267 +
268 + tp_size = get_tensor_model_parallel_world_size()
269 + self.tp_rank = get_tensor_model_parallel_rank()
270 + self.total_num_heads = num_heads
271 + assert self.total_num_heads % tp_size == 0
272 + self.num_heads = self.total_num_heads // tp_size
273 + self.total_num_kv_heads = num_kv_heads
274 + if self.total_num_kv_heads >= tp_size:
275 + assert self.total_num_kv_heads % tp_size == 0
276 + else:
277 + assert tp_size % self.total_num_kv_heads == 0
278 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
279 + self.head_dim = head_dim
280 + self.q_size = self.num_heads * self.head_dim
281 + self.kv_size = self.num_kv_heads * self.head_dim
282 + # Gemma4 uses scaling=1.0.
283 + # Unlike Gemma2/3, query_pre_attn_scalar is NOT used here;
284 + # Q/K norms with learnable weights handle scaling implicitly.
285 + self.scaling = 1.0
286 +
287 + # QKVParallelLinear handles GQA correctly for all layer types.
288 + # k_eq_v layers load K weights into both K and V slots via
289 + # _weight_iterator remapping — no structural difference needed.
290 + self.qkv_proj = QKVParallelLinear(
291 + hidden_size,
292 + self.head_dim,
293 + self.total_num_heads,
294 + self.total_num_kv_heads,
295 + bias=config.attention_bias,
296 + quant_config=quant_config,
297 + prefix=f"{prefix}.qkv_proj",
298 + )
299 + self.o_proj = RowParallelLinear(
300 + self.total_num_heads * self.head_dim,
301 + hidden_size,
302 + bias=config.attention_bias,
303 + quant_config=quant_config,
304 + prefix=f"{prefix}.o_proj",
305 + )
306 +
307 + # Q/K norms: output = norm(x) * weight (learnable per-head scale)
308 + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
309 + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
310 + # V norm: no learnable scale (pure normalization only)
311 + self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False)
312 +
313 + # Determine layer type and sliding window
314 + layer_idx = extract_layer_index(prefix)
315 + layer_type = config.layer_types[layer_idx]
316 + self.is_sliding = layer_type == "sliding_attention"
317 + sliding_window = config.sliding_window if self.is_sliding else None
318 +
319 + # Initialize RoPE based on layer type.
320 + # Gemma4 uses different RoPE parameters for sliding vs full attention.
321 + if layer_type in config.rope_parameters:
322 + # Per-layer-type rope config (dict format).
323 + # rope_parameters already contains the correct
324 + # partial_rotary_factor per layer type (1.0 for full
325 + # attention, 1.0 for sliding). Do NOT override with
326 + # global_partial_rotary_factor — that config key is
327 + # not needed for Gemma4 — config uses per-layer rope_parameters.
328 + rope_parameters = dict(config.rope_parameters[layer_type])
329 + else:
330 + # Legacy config format fallback.
331 + rope_parameters = dict(config.rope_parameters.copy())
332 + if self.is_sliding:
333 + rope_parameters["rope_theta"] = getattr(
334 + config, "rope_local_base_freq", 10000.0
335 + )
336 +
337 + # KV sharing: layers in the last `num_kv_shared_layers` share KV
338 + # cache with earlier layers of the same type.
339 + kv_sharing_target_layer_name = None
340 + self.is_kv_shared_layer = False
341 + num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0)
342 + if num_kv_shared_layers > 0:
343 + first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers
344 + if layer_idx >= first_kv_shared_layer_idx:
345 + self.is_kv_shared_layer = True
346 + # Find the last non-shared layer of the same attention type
347 + prev_layers = config.layer_types[:first_kv_shared_layer_idx]
348 + current_layer_type = config.layer_types[layer_idx]
349 + kv_shared_layer_index = (
350 + len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type)
351 + )
352 + if kv_shared_layer_index >= 0:
353 + if ".layers." in prefix:
354 + param_name_before_layers = prefix.split(".layers.")[0]
355 + else:
356 + raise ValueError(
357 + "Unexpected prefix format for Gemma4Attention: "
358 + f"'{prefix}'. Expected to contain '.layers.'."
359 + )
360 + kv_sharing_target_layer_name = (
361 + f"{param_name_before_layers}.layers."
362 + f"{kv_shared_layer_index}.self_attn.attn"
363 + )
364 +
365 + self.rotary_emb = get_rope(
366 + self.head_dim,
367 + max_position=max_position_embeddings,
368 + rope_parameters=rope_parameters,
369 + is_neox_style=True,
370 + )
371 +
372 + self.attn = Attention(
373 + self.num_heads,
374 + self.head_dim,
375 + self.scaling,
376 + num_kv_heads=self.num_kv_heads,
377 + cache_config=cache_config,
378 + quant_config=quant_config,
379 + logits_soft_cap=attn_logits_soft_cap,
380 + per_layer_sliding_window=sliding_window,
381 + kv_sharing_target_layer_name=kv_sharing_target_layer_name,
382 + prefix=f"{prefix}.attn",
383 + )
384 +
385 + def forward(
386 + self,
387 + positions: torch.Tensor,
388 + hidden_states: torch.Tensor,
389 + **kwargs,
390 + ) -> torch.Tensor:
391 + # Unified QKV path (works for both k_eq_v and standard layers).
392 + # For k_eq_v, K weights are loaded into both K and V slots of
393 + # qkv_proj, so V == K automatically.
394 + qkv, _ = self.qkv_proj(hidden_states)
395 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
396 +
397 + # Q norm (always applied)
398 + q = q.unflatten(-1, (self.num_heads, self.head_dim))
399 + q = self.q_norm(q)
400 + q = q.flatten(-2, -1)
401 +
402 + if not self.is_kv_shared_layer:
403 + # Non-shared: apply K norm + RoPE, V norm
404 + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
405 + k = self.k_norm(k)
406 + k = k.flatten(-2, -1)
407 + q, k = self.rotary_emb(positions, q, k)
408 +
409 + v = v.unflatten(-1, (self.num_kv_heads, self.head_dim))
410 + v = self.v_norm(v)
411 + v = v.flatten(-2, -1)
412 + else:
413 + # Shared: only apply RoPE to Q
414 + q = self.rotary_emb(positions, q, k)[0]
415 +
416 + attn_output = self.attn(q, k, v)
417 + output, _ = self.o_proj(attn_output)
418 +
419 + return output
420 +
421 +
422 +class Gemma4DecoderLayer(nn.Module):
423 + def __init__(
424 + self,
425 + config,
426 + cache_config: CacheConfig | None = None,
427 + quant_config: QuantizationConfig | None = None,
428 + prefix: str = "",
429 + ) -> None:
430 + super().__init__()
431 + self.hidden_size = config.hidden_size
432 + self.hidden_size_per_layer_input = getattr(
433 + config, "hidden_size_per_layer_input", 0
434 + )
435 +
436 + layer_idx = extract_layer_index(prefix)
437 + self.layer_idx = layer_idx
438 +
439 + # Gemma4 uses different head dimensions for sliding vs full attention
440 + layer_type = config.layer_types[layer_idx]
441 + self.is_full_attention = layer_type == "full_attention"
442 + if self.is_full_attention:
443 + head_dim = getattr(config, "global_head_dim", config.head_dim)
444 + else:
445 + head_dim = config.head_dim
446 +
447 + # Determine if this full-attention layer uses k_eq_v
448 + # (laptop variant: no v_proj, K reused as V on full attention layers)
449 + use_k_eq_v = self.is_full_attention and getattr(
450 + config, "attention_k_eq_v", False
451 + )
452 +
453 + # For k_eq_v full-attention layers, use num_global_key_value_heads
454 + # as the KV head count when k_eq_v is enabled.
455 + if use_k_eq_v:
456 + num_kv_heads = getattr(
457 + config, "num_global_key_value_heads", config.num_key_value_heads
458 + )
459 + else:
460 + num_kv_heads = config.num_key_value_heads
461 +
462 + self.self_attn = Gemma4Attention(
463 + config=config,
464 + hidden_size=self.hidden_size,
465 + num_heads=config.num_attention_heads,
466 + num_kv_heads=num_kv_heads,
467 + head_dim=head_dim,
468 + max_position_embeddings=config.max_position_embeddings,
469 + use_k_eq_v=use_k_eq_v,
470 + cache_config=cache_config,
471 + quant_config=quant_config,
472 + attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
473 + prefix=f"{prefix}.self_attn",
474 + )
475 +
476 + # Compute per-layer intermediate_size from config.
477 + # When use_double_wide_mlp is set, intermediate_size doubles for
478 + # KV-shared layers (layers >= first_kv_shared_layer_idx).
479 + first_kv_shared_layer_idx = config.num_hidden_layers - getattr(
480 + config, "num_kv_shared_layers", 0
481 + )
482 + is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0
483 + use_double_wide_mlp = (
484 + getattr(config, "use_double_wide_mlp", False) and is_kv_shared_layer
485 + )
486 + layer_intermediate_size = config.intermediate_size * (
487 + 2 if use_double_wide_mlp else 1
488 + )
489 +
490 + self.mlp = Gemma4MLP(
491 + hidden_size=self.hidden_size,
492 + intermediate_size=layer_intermediate_size,
493 + hidden_activation=config.hidden_activation,
494 + quant_config=quant_config,
495 + prefix=f"{prefix}.mlp",
496 + )
497 +
498 + # Layer norms: output = norm(x) * weight
499 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
500 + self.post_attention_layernorm = RMSNorm(
501 + config.hidden_size, eps=config.rms_norm_eps
502 + )
503 + self.pre_feedforward_layernorm = RMSNorm(
504 + config.hidden_size, eps=config.rms_norm_eps
505 + )
506 + self.post_feedforward_layernorm = RMSNorm(
507 + config.hidden_size, eps=config.rms_norm_eps
508 + )
509 +
510 + # MoE (Mixture of Experts) — router + expert block parallel to MLP
511 + self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr(
512 + config, "use_second_mlp_block", False
513 + )
514 + if self.enable_moe_block:
515 + self.router = Gemma4Router(
516 + config,
517 + quant_config=quant_config,
518 + prefix=f"{prefix}.router",
519 + )
520 + self.moe = Gemma4MoE(
521 + config,
522 + quant_config=quant_config,
523 + prefix=f"{prefix}.moe",
524 + )
525 + self.post_feedforward_layernorm_1 = RMSNorm(
526 + config.hidden_size, eps=config.rms_norm_eps
527 + )
528 + self.post_feedforward_layernorm_2 = RMSNorm(
529 + config.hidden_size, eps=config.rms_norm_eps
530 + )
531 + self.pre_feedforward_layernorm_2 = RMSNorm(
532 + config.hidden_size, eps=config.rms_norm_eps
533 + )
534 + else:
535 + self.router = None
536 + self.moe = None
537 + self.post_feedforward_layernorm_1 = None
538 + self.post_feedforward_layernorm_2 = None
539 + self.pre_feedforward_layernorm_2 = None
540 +
541 + # Per-Layer Embedding (PLE) components — present in each decoder layer
542 + if (
543 + self.hidden_size_per_layer_input is not None
544 + and self.hidden_size_per_layer_input > 0
545 + ):
546 + # Gate: projects hidden_states → per-layer dim for gating
547 + self.per_layer_input_gate = ReplicatedLinear(
548 + self.hidden_size,
549 + self.hidden_size_per_layer_input,
550 + bias=False,
551 + quant_config=quant_config,
552 + prefix=f"{prefix}.per_layer_input_gate",
553 + return_bias=False,
554 + )
555 + # Projection: projects gated per-layer input back → hidden size
556 + self.per_layer_projection = ReplicatedLinear(
557 + self.hidden_size_per_layer_input,
558 + self.hidden_size,
559 + bias=False,
560 + quant_config=quant_config,
561 + prefix=f"{prefix}.per_layer_projection",
562 + return_bias=False,
563 + )
564 + # Post-PLE norm: output = norm(x) * weight
565 + self.post_per_layer_input_norm = RMSNorm(
566 + config.hidden_size, eps=config.rms_norm_eps
567 + )
568 + else:
569 + self.per_layer_input_gate = None
570 + self.per_layer_projection = None
571 + self.post_per_layer_input_norm = None
572 +
573 + # Layer scalar (loaded from checkpoint) — applies to ALL text layers
574 + self.register_buffer("layer_scalar", torch.ones(1))
575 +
576 + def forward(
577 + self,
578 + positions: torch.Tensor,
579 + hidden_states: torch.Tensor,
580 + residual: torch.Tensor | None,
581 + per_layer_input: torch.Tensor | None = None,
582 + **kwargs,
583 + ) -> tuple[torch.Tensor, torch.Tensor]:
584 + # Gemma4 residual pattern:
585 + # 1. input_norm(x) → attn → post_attn_norm → ADD residual
586 + # 2. pre_ff_norm → mlp → post_ff_norm → ADD residual
587 + residual = hidden_states
588 +
589 + hidden_states = self.input_layernorm(residual)
590 +
591 + hidden_states = self.self_attn(
592 + positions=positions,
593 + hidden_states=hidden_states,
594 + **kwargs,
595 + )
596 +
597 + hidden_states = self.post_attention_layernorm(hidden_states)
598 + hidden_states = hidden_states + residual
599 + residual = hidden_states
600 +
601 + # MLP runs unconditionally (same inputs for MoE and non-MoE)
602 + hidden_states = self.pre_feedforward_layernorm(hidden_states)
603 + hidden_states = self.mlp(hidden_states)
604 +
605 + if self.enable_moe_block:
606 + hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states)
607 +
608 + # Router and MoE experts see the residual (pre-MLP state),
609 + # matching the HF transformers forward path
610 + router_logits = self.router(residual)
611 + hidden_states_2 = self.pre_feedforward_layernorm_2(residual)
612 + hidden_states_2 = self.moe(hidden_states_2, router_logits)
613 + hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2)
614 +
615 + # Combine MLP and MoE outputs
616 + hidden_states = hidden_states_1 + hidden_states_2
617 +
618 + hidden_states = self.post_feedforward_layernorm(hidden_states)
619 + hidden_states = hidden_states + residual
620 +
621 + # Apply PLE (Per-Layer Embedding) if configured
622 + if per_layer_input is not None and self.per_layer_input_gate is not None:
623 + gate = self.per_layer_input_gate(hidden_states)
624 + gate = torch.nn.functional.gelu(gate, approximate="tanh")
625 + gated_per_layer = gate * per_layer_input
626 + per_layer_contribution = self.per_layer_projection(gated_per_layer)
627 + per_layer_contribution = self.post_per_layer_input_norm(
628 + per_layer_contribution
629 + )
630 + hidden_states = hidden_states + per_layer_contribution
631 +
632 + # Apply layer scalar for full-attention layers
633 + # Apply per-layer scalar (all text layers)
634 + hidden_states = hidden_states * self.layer_scalar
635 +
636 + return hidden_states, None
637 +
638 +
639 +@support_torch_compile
640 +class Gemma4Model(nn.Module):
641 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
642 + super().__init__()
643 + config = _get_text_config(vllm_config.model_config.hf_config)
644 + cache_config = vllm_config.cache_config
645 + quant_config = vllm_config.quant_config
646 + self.config = config
647 + self.quant_config = quant_config
648 +
649 + # PLE config values (default to 0 if not present — disables PLE)
650 + self.hidden_size_per_layer_input = getattr(
651 + config, "hidden_size_per_layer_input", 0
652 + )
653 + self.vocab_size_per_layer_input = getattr(
654 + config, "vocab_size_per_layer_input", config.vocab_size
655 + )
656 +
657 + self.embed_tokens = VocabParallelEmbedding(
658 + config.vocab_size,
659 + config.hidden_size,
660 + quant_config=quant_config,
661 + prefix=f"{prefix}.embed_tokens",
662 + )
663 +
664 + # Per-Layer Embedding (PLE) components
665 + if (
666 + self.hidden_size_per_layer_input is not None
667 + and self.hidden_size_per_layer_input > 0
668 + ):
669 + total_ple_dim = self.hidden_size_per_layer_input * config.num_hidden_layers
670 + self.embed_tokens_per_layer = VocabParallelEmbedding(
671 + self.vocab_size_per_layer_input,
672 + total_ple_dim,
673 + quant_config=quant_config,
674 + prefix=f"{prefix}.embed_tokens_per_layer",
675 + )
676 + # Scaled embedding factor (from config, not hardcoded)
677 + # Register as buffer so it moves to GPU with the model
678 + # and interacts correctly with torch.compile AOT caching.
679 + self.register_buffer(
680 + "embed_scale_per_layer",
681 + torch.tensor(self.hidden_size_per_layer_input**0.5),
682 + persistent=False,
683 + )
684 + # Projection: hidden_size → total_ple_dim
685 + # ColumnParallelLinear with gather_output=True
686 + self.per_layer_model_projection = ColumnParallelLinear(
687 + config.hidden_size,
688 + total_ple_dim,
689 + bias=False,
690 + gather_output=True,
691 + return_bias=False,
692 + quant_config=quant_config,
693 + prefix=f"{prefix}.per_layer_model_projection",
694 + )
695 + # PLE projection norm: output = norm(x) * weight
696 + self.per_layer_projection_norm = RMSNorm(
697 + self.hidden_size_per_layer_input,
698 + eps=config.rms_norm_eps,
699 + )
700 + # Scale factor for combining projection + per_layer_inputs
701 + # Register as buffer so it moves to GPU with the model
702 + # and interacts correctly with torch.compile AOT caching.
703 + self.register_buffer(
704 + "per_layer_input_scale",
705 + torch.rsqrt(torch.tensor(2.0)),
706 + persistent=False,
707 + )
708 + # Scaled projection: multiply output by hidden_size**-0.5.
709 + # Register as buffer for GPU placement and torch.compile.
710 + self.register_buffer(
711 + "per_layer_projection_scale",
712 + torch.tensor(config.hidden_size**-0.5),
713 + persistent=False,
714 + )
715 + else:
716 + self.embed_tokens_per_layer = None
717 + self.embed_scale_per_layer = None
718 + self.per_layer_model_projection = None
719 + self.per_layer_projection_norm = None
720 + self.per_layer_input_scale = None
721 + self.per_layer_projection_scale = None
722 +
723 + self.start_layer, self.end_layer, self.layers = make_layers(
724 + config.num_hidden_layers,
725 + lambda prefix: Gemma4DecoderLayer(
726 + config,
727 + cache_config=cache_config,
728 + quant_config=quant_config,
729 + prefix=prefix,
730 + ),
731 + prefix=f"{prefix}.layers",
732 + )
733 + # Final norm: output = norm(x) * weight
734 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
735 +
736 + # Embedding scale = sqrt(hidden_size)
737 + # Downcast to model dtype (bfloat16 etc.) for numerical parity
738 + self.register_buffer(
739 + "normalizer",
740 + torch.tensor(config.hidden_size**0.5),
741 + persistent=False,
742 + )
743 + # Custom factory that includes per_layer_inputs for PLE-enabled PP.
744 + # per_layer_inputs has shape (batch, num_layers, per_layer_dim),
745 + # which differs from the standard (batch, hidden_size) shape,
746 + # so we can't use the default factory.
747 + ple_dim = self.hidden_size_per_layer_input
748 + num_layers = config.num_hidden_layers
749 + hidden_size = config.hidden_size
750 +
751 + def _make_empty_intermediate_tensors(
752 + batch_size: int,
753 + dtype: torch.dtype,
754 + device: torch.device,
755 + ) -> IntermediateTensors:
756 + tensors: dict[str, torch.Tensor] = {
757 + "hidden_states": torch.zeros(
758 + (batch_size, hidden_size),
759 + dtype=dtype,
760 + device=device,
761 + ),
762 + "residual": torch.zeros(
763 + (batch_size, hidden_size),
764 + dtype=dtype,
765 + device=device,
766 + ),
767 + }
768 + if ple_dim and ple_dim > 0:
769 + tensors["per_layer_inputs"] = torch.zeros(
770 + (batch_size, num_layers, ple_dim),
771 + dtype=dtype,
772 + device=device,
773 + )
774 + return IntermediateTensors(tensors)
775 +
776 + self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors
777 +
778 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
779 + return self.embed_tokens(input_ids) * self.normalizer
780 +
781 + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor:
782 + """Get per-layer embeddings from embed_tokens_per_layer.
783 +
784 + Returns:
785 + Per-layer embeddings (num_tokens, num_layers,
786 + hidden_size_per_layer_input)
787 + """
788 + if self.embed_tokens_per_layer is None:
789 + return None
790 +
791 + # Handle out-of-vocab tokens for PLE (vocab_size_per_layer_input may
792 + # be smaller than the main vocab_size).
793 + per_layer_inputs_mask = torch.logical_and(
794 + input_ids >= 0,
795 + input_ids < self.vocab_size_per_layer_input,
796 + )
797 + per_layer_inputs_tokens = torch.where(
798 + per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)
799 + )
800 +
801 + # Get packed per-layer embeddings: (num_tokens, total_ple_dim)
802 + per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)
803 +
804 + # Apply embed_scale (sqrt of per-layer hidden dim)
805 + per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer
806 +
807 + # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)
808 + per_layer_embeds = per_layer_embeds.reshape(
809 + *input_ids.shape,
810 + self.config.num_hidden_layers,
811 + self.hidden_size_per_layer_input,
812 + )
813 + return per_layer_embeds
814 +
815 + def project_per_layer_inputs(
816 + self,
817 + inputs_embeds: torch.Tensor,
818 + per_layer_inputs: torch.Tensor | None,
819 + ) -> torch.Tensor:
820 + """Project inputs_embeds and combine with per_layer_inputs.
821 +
822 + Steps:
823 + 1. Project inputs_embeds: hidden_size → total_ple_dim
824 + 2. Scale by hidden_size^{-0.5}
825 + 3. Reshape to (num_tokens, num_layers, per_layer_dim)
826 + 4. Normalize with per_layer_projection_norm
827 + 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2)
828 + """
829 + if self.per_layer_model_projection is None:
830 + return None
831 +
832 + # Project from hidden_size to total_ple_dim
833 + # Scaled projection: output = linear(input, weight) * scale
834 + per_layer_projection = self.per_layer_model_projection(inputs_embeds)
835 + per_layer_projection = per_layer_projection * self.per_layer_projection_scale
836 +
837 + # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)
838 + per_layer_projection = per_layer_projection.reshape(
839 + *inputs_embeds.shape[:-1],
840 + self.config.num_hidden_layers,
841 + self.hidden_size_per_layer_input,
842 + )
843 +
844 + # Normalize
845 + per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
846 +
847 + if per_layer_inputs is None:
848 + return per_layer_projection
849 +
850 + # Combine: (projection + per_layer_inputs) * scale
851 + return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
852 +
853 + def forward(
854 + self,
855 + input_ids: torch.Tensor | None,
856 + positions: torch.Tensor,
857 + intermediate_tensors: IntermediateTensors | None,
858 + inputs_embeds: torch.Tensor | None = None,
859 + per_layer_inputs: torch.Tensor | None = None,
860 + **kwargs,
861 + ) -> torch.Tensor | IntermediateTensors:
862 + if get_pp_group().is_first_rank:
863 + if inputs_embeds is not None:
864 + hidden_states = inputs_embeds
865 + # When called from the multimodal wrapper, raw PLE
866 + # embeddings are pre-computed and passed explicitly.
867 + # Project them through per_layer_model_projection.
868 + per_layer_inputs = self.project_per_layer_inputs(
869 + hidden_states, per_layer_inputs
870 + )
871 + else:
872 + hidden_states = self.embed_input_ids(input_ids)
873 + # Compute per-layer inputs for PLE
874 + per_layer_embeds = self.get_per_layer_inputs(input_ids)
875 + per_layer_inputs = self.project_per_layer_inputs(
876 + hidden_states, per_layer_embeds
877 + )
878 + residual = None
879 + else:
880 + assert intermediate_tensors is not None
881 + hidden_states = intermediate_tensors["hidden_states"]
882 + residual = intermediate_tensors["residual"]
883 + per_layer_inputs = intermediate_tensors.get("per_layer_inputs")
884 +
885 + for layer_idx, layer in enumerate(
886 + islice(self.layers, self.start_layer, self.end_layer)
887 + ):
888 + # Extract the per-layer embedding for this specific layer
889 + if per_layer_inputs is not None:
890 + actual_layer_idx = self.start_layer + layer_idx
891 + layer_per_input = per_layer_inputs[
892 + :, actual_layer_idx, :
893 + ] # (num_tokens, per_layer_dim)
894 + else:
895 + layer_per_input = None
896 + hidden_states, residual = layer(
897 + positions,
898 + hidden_states,
899 + residual,
900 + per_layer_input=layer_per_input,
901 + **kwargs,
902 + )
903 + if not get_pp_group().is_last_rank:
904 + return IntermediateTensors(
905 + {
906 + "hidden_states": hidden_states,
907 + "residual": residual,
908 + "per_layer_inputs": per_layer_inputs,
909 + }
910 + )
911 + # Gemma4 incorporates residual into hidden_states directly
912 + # Apply norm without residual fusion when possible.
913 + if residual is None:
914 + hidden_states = self.norm(hidden_states)
915 + else:
916 + hidden_states, _ = self.norm(hidden_states, residual)
917 + return hidden_states
918 +
919 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
920 + stacked_params_mapping = [
921 + # (param_name, shard_name, shard_id)
922 + ("qkv_proj", "q_proj", "q"),
923 + ("qkv_proj", "k_proj", "k"),
924 + ("qkv_proj", "v_proj", "v"),
925 + ("gate_up_proj", "gate_proj", 0),
926 + ("gate_up_proj", "up_proj", 1),
927 + ]
928 +
929 + # MoE expert weight mapping: checkpoint 3D packed tensors are
930 + # exploded in _weight_iterator to per-expert 2D weights like:
931 + # moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13)
932 + # moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13)
933 + # moe.experts.{id}.down_proj → FusedMoE w2
934 + # We build the mapping directly since Gemma4 uses bare param
935 + # names (no .weight suffix) unlike standard MoE checkpoints.
936 + num_experts = getattr(self.config, "num_experts", None) or 0
937 + expert_params_mapping = [
938 + # (param_name, weight_name, expert_id, shard_id)
939 + (
940 + "experts.w13_weight"
941 + if proj_name in ["gate_proj", "up_proj"]
942 + else "experts.w2_weight",
943 + f"experts.{expert_id}.{proj_name}",
944 + expert_id,
945 + shard_id,
946 + )
947 + for expert_id in range(num_experts)
948 + for shard_id, proj_name in [
949 + ("w1", "gate_proj"),
950 + ("w2", "down_proj"),
951 + ("w3", "up_proj"),
952 + ]
953 + ]
954 + params_dict = dict(self.named_parameters())
955 + # Include buffers (e.g. layer_scalar) so they can be loaded too
956 + params_dict.update(dict(self.named_buffers()))
957 + loaded_params: set[str] = set()
958 + for name, loaded_weight in weights:
959 + if self.quant_config is not None and (
960 + scale_name := self.quant_config.get_cache_scale(name)
961 + ):
962 + param = params_dict[scale_name]
963 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
964 + loaded_weight = loaded_weight[0]
965 + weight_loader(param, loaded_weight)
966 + loaded_params.add(scale_name)
967 + continue
968 +
969 + if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
970 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)
971 + if remapped_name is not None and remapped_name in params_dict:
972 + param = params_dict[remapped_name]
973 + weight_loader = getattr(
974 + param, "weight_loader", default_weight_loader
975 + )
976 + weight_loader(param, loaded_weight)
977 + loaded_params.add(remapped_name)
978 + continue
979 +
980 + for param_name, shard_name, shard_id in stacked_params_mapping:
981 + if shard_name not in name:
982 + continue
983 + stacked_name = name.replace(shard_name, param_name)
984 + # k_eq_v layers use separate q_proj/k_proj instead of
985 + # packed qkv_proj. If the stacked param doesn't exist,
986 + # skip this mapping and fall through to direct load.
987 + if stacked_name not in params_dict:
988 + continue
989 + if is_pp_missing_parameter(stacked_name, self):
990 + continue
991 + param = params_dict[stacked_name]
992 + weight_loader = param.weight_loader
993 + weight_loader(param, loaded_weight, shard_id)
994 + loaded_params.add(stacked_name)
995 + break
996 + else:
997 + for (
998 + param_name,
999 + weight_name,
1000 + expert_id,
1001 + shard_id,
1002 + ) in expert_params_mapping:
1003 + if weight_name not in name:
1004 + continue
1005 + moe_name = name.replace(weight_name, param_name)
1006 + if moe_name not in params_dict:
1007 + continue
1008 + if is_pp_missing_parameter(moe_name, self):
1009 + continue
1010 + param = params_dict[moe_name]
1011 + # Expert weights are already in the correct
1012 + # orientation for FusedMoE after _weight_iterator:
1013 + # gate/up: [I, H] → w1/w3 expects [I, H]
1014 + # down: [H, I] → w2 expects [H, I]
1015 + assert loaded_weight.dim() == 2, (
1016 + f"Expected 2D expert weight for {weight_name}, "
1017 + f"got shape {loaded_weight.shape}"
1018 + )
1019 + weight_loader = param.weight_loader
1020 + weight_loader(
1021 + param,
1022 + loaded_weight,
1023 + weight_name + ".weight",
1024 + shard_id=shard_id,
1025 + expert_id=expert_id,
1026 + )
1027 + loaded_params.add(moe_name)
1028 + break
1029 + else:
1030 + if name.endswith(".bias") and name not in params_dict:
1031 + continue
1032 + name = maybe_remap_kv_scale_name(name, params_dict)
1033 + if name is None:
1034 + continue
1035 + if is_pp_missing_parameter(name, self):
1036 + continue
1037 + param = params_dict[name]
1038 + weight_loader = getattr(
1039 + param, "weight_loader", default_weight_loader
1040 + )
1041 + weight_loader(param, loaded_weight)
1042 + loaded_params.add(name)
1043 +
1044 + return loaded_params
1045 +
1046 +
1047 +class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
1048 + # Note: qkv_proj packing applies to non-k_eq_v layers (sliding
1049 + # attention and full attention without k_eq_v). k_eq_v layers use
1050 + # separate q_proj + k_proj without packing.
1051 + packed_modules_mapping = {
1052 + "qkv_proj": [
1053 + "q_proj",
1054 + "k_proj",
1055 + "v_proj",
1056 + ],
1057 + "gate_up_proj": [
1058 + "gate_proj",
1059 + "up_proj",
1060 + ],
1061 + }
1062 +
1063 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1064 + config = _get_text_config(vllm_config.model_config.hf_config)
1065 + quant_config = vllm_config.quant_config
1066 +
1067 + super().__init__()
1068 + self.config = config
1069 + self.quant_config = quant_config
1070 + self.model = Gemma4Model(
1071 + vllm_config=vllm_config,
1072 + prefix=maybe_prefix(prefix, "model"),
1073 + )
1074 +
1075 + self.lm_head = ParallelLMHead(
1076 + config.vocab_size,
1077 + config.hidden_size,
1078 + quant_config=quant_config,
1079 + prefix=maybe_prefix(prefix, "lm_head"),
1080 + )
1081 + if config.tie_word_embeddings:
1082 + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
1083 +
1084 + self.logits_processor = LogitsProcessor(
1085 + config.vocab_size,
1086 + soft_cap=getattr(config, "final_logit_softcapping", None),
1087 + )
1088 + self.make_empty_intermediate_tensors = (
1089 + self.model.make_empty_intermediate_tensors
1090 + )
1091 +
1092 + # --- MixtureOfExperts protocol ---
1093 + self.expert_weights: list[list[torch.Tensor]] = []
1094 + self.moe_layers: list[nn.Module] = []
1095 + example_moe: Gemma4MoE | None = None
1096 +
1097 + for layer in self.model.layers:
1098 + if hasattr(layer, "moe") and isinstance(layer.moe, Gemma4MoE):
1099 + example_moe = layer.moe
1100 + self.moe_layers.append(layer.moe.experts)
1101 +
1102 + self.num_moe_layers = len(self.moe_layers)
1103 +
1104 + if example_moe is not None:
1105 + self.num_logical_experts = example_moe.num_experts
1106 + self.num_physical_experts = example_moe.num_experts
1107 + self.num_local_physical_experts = example_moe.num_experts
1108 + self.num_routed_experts = example_moe.num_experts
1109 + else:
1110 + self.num_logical_experts = 0
1111 + self.num_physical_experts = 0
1112 + self.num_local_physical_experts = 0
1113 + self.num_routed_experts = 0
1114 +
1115 + self.num_expert_groups = 1
1116 + self.num_shared_experts = 0
1117 + self.num_redundant_experts = 0
1118 +
1119 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1120 + return self.model.embed_input_ids(input_ids)
1121 +
1122 + def forward(
1123 + self,
1124 + input_ids: torch.Tensor,
1125 + positions: torch.Tensor,
1126 + intermediate_tensors: IntermediateTensors | None = None,
1127 + inputs_embeds: torch.Tensor | None = None,
1128 + **kwargs,
1129 + ) -> torch.Tensor | IntermediateTensors:
1130 + hidden_states = self.model(
1131 + input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
1132 + )
1133 + return hidden_states
1134 +
1135 + def compute_logits(
1136 + self,
1137 + hidden_states: torch.Tensor,
1138 + ) -> torch.Tensor | None:
1139 + return self.logits_processor(self.lm_head, hidden_states)
1140 +
1141 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1142 + # Checkpoint weight names use "language_model." prefix (from the
1143 + # Gemma4ForConditionalGeneration wrapper). Strip it to map to our
1144 + # model tree which is just "model.*".
1145 + def _weight_iterator():
1146 + use_k_eq_v = getattr(self.config, "attention_k_eq_v", False)
1147 + # Build set of k_eq_v layer indices (full_attention layers
1148 + # when attention_k_eq_v is enabled). These layers have k_proj
1149 + # but no v_proj in checkpoint — we duplicate k_proj as v_proj.
1150 + k_eq_v_layer_indices: set[int] = set()
1151 + if use_k_eq_v:
1152 + for idx, lt in enumerate(self.config.layer_types):
1153 + if lt == "full_attention":
1154 + k_eq_v_layer_indices.add(idx)
1155 +
1156 + for name, weight in weights:
1157 + # Remap "language_model." → "" to match our model tree.
1158 + # Checkpoint: model.language_model.layers.X.*
1159 + # Our model: model.layers.X.*
1160 + name = name.replace("language_model.", "")
1161 +
1162 + # Remap new HF checkpoint naming to internal vLLM
1163 + # naming: HF moved per_expert_scale to router and
1164 + # renamed moe → experts in the MoE block.
1165 + name = name.replace(
1166 + ".router.per_expert_scale",
1167 + ".moe.per_expert_scale",
1168 + )
1169 + if ".experts.gate_up_proj" in name:
1170 + name = name.replace(
1171 + ".experts.gate_up_proj",
1172 + ".moe.gate_up_proj",
1173 + )
1174 + elif ".experts.down_proj" in name:
1175 + name = name.replace(
1176 + ".experts.down_proj",
1177 + ".moe.down_proj",
1178 + )
1179 +
1180 + # MoE expert weights: checkpoint stores as 3D packed
1181 + # tensors. Explode into per-expert 2D weights for
1182 + # FusedMoE weight_loader.
1183 + #
1184 + # Checkpoint format:
1185 + # moe.gate_up_proj: [E, 2*I, H] (fused gate + up)
1186 + # moe.down_proj: [E, H, I]
1187 + #
1188 + # FusedMoE expects per-expert:
1189 + # w1 (gate): [I, H] — first half of gate_up
1190 + # w3 (up): [I, H] — second half of gate_up
1191 + # w2 (down): [H, I] — as-is from checkpoint
1192 + #
1193 + # No transpose needed: checkpoint orientation already
1194 + # matches FusedMoE's expected layout.
1195 + if "moe.gate_up_proj" in name and weight.dim() == 3:
1196 + num_experts = weight.size(0)
1197 + intermediate_size = weight.size(1) // 2
1198 + for expert_id in range(num_experts):
1199 + gate_weight = weight[expert_id, :intermediate_size, :]
1200 + up_weight = weight[expert_id, intermediate_size:, :]
1201 + base = name.replace("moe.", f"moe.experts.{expert_id}.")
1202 + yield base.replace("gate_up_proj", "gate_proj"), gate_weight
1203 + yield base.replace("gate_up_proj", "up_proj"), up_weight
1204 + continue
1205 +
1206 + if "moe.down_proj" in name and weight.dim() == 3:
1207 + num_experts = weight.size(0)
1208 + for expert_id in range(num_experts):
1209 + expert_name = name.replace("moe.", f"moe.experts.{expert_id}.")
1210 + yield expert_name, weight[expert_id]
1211 + continue
1212 +
1213 + # k_eq_v layers: checkpoint has k_proj but no v_proj.
1214 + # QKVParallelLinear expects both, so duplicate k_proj
1215 + # as v_proj so V gets identical weights to K.
1216 + # ONLY for full_attention layers — sliding layers have
1217 + # their own real v_proj weights.
1218 + if "self_attn.k_proj" in name and k_eq_v_layer_indices:
1219 + m = re.search(r"layers\.(\d+)\.", name)
1220 + if m and int(m.group(1)) in k_eq_v_layer_indices:
1221 + yield name, weight
1222 + yield name.replace("k_proj", "v_proj"), weight.clone()
1223 + continue
1224 +
1225 + yield name, weight
1226 +
1227 + # Skip multimodal weights — handled by the multimodal wrapper.
1228 + # Also skip lm_head when weights are tied.
1229 + skip = [
1230 + "audio_tower.",
1231 + "vision_tower.",
1232 + "embed_audio.",
1233 + "embed_vision.",
1234 + ]
1235 + if self.config.tie_word_embeddings:
1236 + skip.append("lm_head.")
1237 +
1238 + loader = AutoWeightsLoader(self, skip_substrs=skip)
1239 + return loader.load_weights(_weight_iterator())
vllm/model_executor/models/gemma4_mm.py new
+1341
@@ -0,0 +1,1341 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gemma 4 multimodal model (image + audio + video support).
4 +
5 +Adds vision tower, audio tower, and multimodal embedders on top of the
6 +text-only Gemma4ForCausalLM. The vision/audio encoders are loaded via
7 +AutoModel.from_config and run in eager mode while the language model uses
8 +the vLLM-optimized path.
9 +
10 +Video support: Gemma4 does **not** have a native video tower. Videos are
11 +decomposed into timestamped image frames (up to 32 frames at 70 soft tokens
12 +each) and fed through the same vision tower as regular images. The
13 +processor inserts ``mm:ss`` timestamps between frames so the model can
14 +reason about temporal order.
15 +"""
16 +
17 +import math
18 +import sys
19 +from collections.abc import Iterable, Mapping, Sequence
20 +from typing import Annotated, Any, Literal
21 +
22 +import numpy as np
23 +import torch
24 +from PIL import Image as PILImage
25 +from torch import nn
26 +from transformers import AutoModel, BatchFeature
27 +from transformers.models.gemma4 import (
28 + Gemma4Config,
29 + Gemma4Processor,
30 + Gemma4VisionConfig,
31 +)
32 +from transformers.models.gemma4.configuration_gemma4 import (
33 + Gemma4AudioConfig,
34 + Gemma4TextConfig,
35 +)
36 +
37 +from vllm.config import VllmConfig
38 +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
39 +from vllm.inputs import MultiModalDataDict
40 +from vllm.logger import init_logger
41 +from vllm.model_executor.layers.layernorm import RMSNorm
42 +from vllm.model_executor.layers.linear import ReplicatedLinear
43 +from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
44 +from vllm.model_executor.models.module_mapping import MultiModelKeys
45 +from vllm.multimodal import MULTIMODAL_REGISTRY
46 +from vllm.multimodal.inputs import (
47 + MultiModalFieldConfig,
48 + MultiModalKwargsItems,
49 + VideoItem,
50 +)
51 +from vllm.multimodal.parse import (
52 + AudioProcessorItems,
53 + ImageProcessorItems,
54 + MultiModalDataItems,
55 + MultiModalDataParser,
56 +)
57 +from vllm.multimodal.processing import BaseDummyInputsBuilder
58 +from vllm.multimodal.processing.processor import (
59 + BaseMultiModalProcessor,
60 + BaseProcessingInfo,
61 + PromptReplacement,
62 + PromptUpdate,
63 + PromptUpdateDetails,
64 +)
65 +from vllm.sequence import IntermediateTensors
66 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
67 +
68 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
69 +from .utils import (
70 + AutoWeightsLoader,
71 + WeightsMapper,
72 + init_vllm_registered_model,
73 + maybe_prefix,
74 +)
75 +
76 +logger = init_logger(__name__)
77 +
78 +# Video constants — match transformers Gemma4VideoProcessor defaults.
79 +_VIDEO_MAX_SOFT_TOKENS = 70 # soft tokens per video frame (vs 280 for images)
80 +_VIDEO_MAX_FRAMES = 32 # max sampled frames per video
81 +
82 +
83 +# ---------------------------------------------------------------------------
84 +# Input schema
85 +# ---------------------------------------------------------------------------
86 +
87 +
88 +class Gemma4ImagePixelInputs(TensorSchema):
89 + """
90 + Pre-patchified image inputs from the Gemma4 image processor.
91 +
92 + Dimensions:
93 + - bn: Batch size * number of images
94 + - np: Number of patches (max_patches = max_soft_tokens * pooling_kernel_size²)
95 + - pp: Patch pixels (patch_size² * 3)
96 +
97 + The HF Gemma4ImageProcessor outputs pixel_values as
98 + (batch, max_patches, patch_pixels) — already patchified with
99 + zero-padding for patches beyond the real image content.
100 + pixel_position_ids provides (x, y) coordinates per patch,
101 + with (-1, -1) for padding patches.
102 + """
103 +
104 + type: Literal["pixel_values"] = "pixel_values"
105 + pixel_values: Annotated[
106 + torch.Tensor,
107 + TensorShape("bn", "np", "pp"),
108 + ]
109 + pixel_position_ids: Annotated[
110 + torch.Tensor,
111 + TensorShape("bn", "np", 2),
112 + ]
113 +
114 +
115 +class Gemma4AudioInputs(TensorSchema):
116 + """
117 + Dimensions:
118 + - bn: Batch size * number of audios
119 + - s: Sequence length (MEL spectrogram frames)
120 + - f: Number of features (MEL bins)
121 + """
122 +
123 + type: Literal["audio"] = "audio"
124 + input_features_padded: Annotated[torch.Tensor, TensorShape("bn", "s", "f")]
125 + input_features_mask: Annotated[torch.Tensor, TensorShape("bn", "s")]
126 +
127 +
128 +Gemma4ImageInputs = Gemma4ImagePixelInputs
129 +
130 +
131 +class Gemma4VideoInputs(TensorSchema):
132 + """Video frame inputs — same tensor format as image inputs.
133 +
134 + Gemma4 has no separate video tower; video frames are processed
135 + through the vision tower at lower resolution (max_soft_tokens=70).
136 + """
137 +
138 + type: Literal["pixel_values_videos"] = "pixel_values_videos"
139 + pixel_values_videos: Annotated[
140 + torch.Tensor,
141 + TensorShape("bn", "np", "pp"),
142 + ]
143 + pixel_position_ids_videos: Annotated[
144 + torch.Tensor,
145 + TensorShape("bn", "np", 2),
146 + ]
147 +
148 +
149 +# ---------------------------------------------------------------------------
150 +# Processing info
151 +# ---------------------------------------------------------------------------
152 +
153 +
154 +class Gemma4ProcessingInfo(BaseProcessingInfo):
155 + def get_hf_config(self):
156 + return self.ctx.get_hf_config(Gemma4Config)
157 +
158 + def get_default_tok_params(self):
159 + """Gemma4's chat template already embeds a literal ``<bos>`` token in
160 + the rendered text. If ``add_special_tokens=True`` (the base-class
161 + default), the tokenizer prepends *another* BOS, producing a
162 + ``[2, 2, ...]`` double-BOS sequence that the model was not trained on.
163 +
164 + Setting ``add_special_tokens=False`` here prevents the duplicate and
165 + ensures both ``llm.generate()`` and the chat/completions API behave
166 + correctly.
167 + """
168 + params = super().get_default_tok_params()
169 + params = params.with_kwargs(add_special_tokens=False)
170 + return params
171 +
172 + def get_hf_processor(self, **kwargs: object) -> Gemma4Processor:
173 + return self.ctx.get_hf_processor(
174 + Gemma4Processor,
175 + **kwargs,
176 + )
177 +
178 + def validate_num_items(self, modality: str, num_items: int) -> None:
179 + if (
180 + modality == "audio"
181 + and num_items > 0
182 + and self.get_hf_config().audio_config is None
183 + ):
184 + model = self.ctx.model_config.model
185 + raise ValueError(
186 + f"Audio input was provided but the model "
187 + f"'{model}' does not have an audio tower. "
188 + f"Audio inference is only supported for Gemma4 "
189 + f"models that include an audio_config "
190 + f"(i.e., models that include an audio_config)."
191 + )
192 + super().validate_num_items(modality, num_items)
193 +
194 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
195 + limits: dict[str, int | None] = {"image": None}
196 + if self.get_hf_config().audio_config is not None:
197 + limits["audio"] = None
198 + limits["video"] = None
199 + return limits
200 +
201 + def get_mm_max_tokens_per_item(
202 + self, seq_len: int, mm_counts: Mapping[str, int]
203 + ) -> Mapping[str, int] | None:
204 + config = self.get_hf_config()
205 + # Upper bound: the pooler outputs default_output_length slots
206 + # per image (280). After padding is stripped the actual count
207 + # is ≤ this value, but vLLM needs the max for memory planning.
208 + tokens_per_image = config.vision_config.default_output_length
209 + tokens: dict[str, int] = {"image": tokens_per_image}
210 + if config.audio_config is not None:
211 + # Audio max tokens from the processor's audio_seq_length.
212 + processor = self.get_hf_processor()
213 + tokens["audio"] = processor.audio_seq_length
214 + # Video: each frame ≤ 70 soft tokens + boi + eoi + ~6 ts tokens.
215 + tokens["video"] = _VIDEO_MAX_FRAMES * (_VIDEO_MAX_SOFT_TOKENS + 2 + 6)
216 + return tokens
217 +
218 + def get_data_parser(self) -> MultiModalDataParser:
219 + config = self.get_hf_config()
220 + kwargs: dict[str, Any] = {"video_needs_metadata": True}
221 + if getattr(config, "audio_config", None) is not None:
222 + processor = self.get_hf_processor()
223 + kwargs["target_sr"] = processor.feature_extractor.sampling_rate
224 + return MultiModalDataParser(**kwargs)
225 +
226 + def _compute_num_soft_tokens(
227 + self,
228 + image_width: int,
229 + image_height: int,
230 + max_soft_tokens: int | None = None,
231 + ) -> int:
232 + """Compute the number of soft tokens the vision tower produces
233 + for an image of the given dimensions, after padding is stripped.
234 +
235 + Args:
236 + max_soft_tokens: Override for the vision config's
237 + ``default_output_length``. When *None*, the value from
238 + the model config is used.
239 + """
240 + vision_cfg = self.get_hf_config().vision_config
241 + patch_size = vision_cfg.patch_size
242 + pooling_kernel_size = vision_cfg.pooling_kernel_size
243 +
244 + if max_soft_tokens is None:
245 + max_soft_tokens = vision_cfg.default_output_length
246 +
247 + unit = patch_size * pooling_kernel_size
248 + max_patches = max_soft_tokens * pooling_kernel_size**2
249 + num_patches_orig = (image_height / patch_size) * (image_width / patch_size)
250 + scale = math.sqrt(max_patches / num_patches_orig)
251 + target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit)
252 + target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit)
253 + num_patches = (target_h // patch_size) * (target_w // patch_size)
254 + return num_patches // (pooling_kernel_size**2)
255 +
256 + def get_image_repl(
257 + self,
258 + *,
259 + image_width: int,
260 + image_height: int,
261 + processor: Gemma4Processor | None,
262 + max_soft_tokens: int | None = None,
263 + ) -> PromptUpdateDetails[list[int]]:
264 + """Return the dynamic image token sequence for this image.
265 +
266 + Computes the exact number of soft tokens the vision tower will
267 + produce after stripping padding.
268 +
269 + Args:
270 + max_soft_tokens: Override for the default token budget.
271 + When *None*, falls back to the model config value.
272 + """
273 + if processor is None:
274 + processor = self.get_hf_processor()
275 +
276 + num_soft = self._compute_num_soft_tokens(
277 + image_width,
278 + image_height,
279 + max_soft_tokens=max_soft_tokens,
280 + )
281 + config = self.get_hf_config()
282 + token_ids = (
283 + [config.boi_token_id]
284 + + [processor.image_token_id] * num_soft
285 + + [config.eoi_token_id]
286 + )
287 + return PromptUpdateDetails.select_token_id(token_ids, processor.image_token_id)
288 +
289 + def get_audio_repl(
290 + self,
291 + *,
292 + audio_len: int,
293 + processor: Gemma4Processor | None,
294 + ) -> PromptUpdateDetails[list[int]]:
295 + """Return the dynamic audio token sequence for this audio.
296 +
297 + Computes the number of soft tokens from the audio waveform
298 + length using ``ceil(duration_ms / audio_ms_per_token)``.
299 + """
300 + if processor is None:
301 + processor = self.get_hf_processor()
302 +
303 + sampling_rate = processor.feature_extractor.sampling_rate
304 + num_tokens = processor._compute_audio_num_tokens(
305 + torch.zeros(audio_len), sampling_rate
306 + )
307 + config = self.get_hf_config()
308 + token_ids = (
309 + [config.boa_token_id]
310 + + [processor.audio_token_id] * num_tokens
311 + + [config.eoa_token_id]
312 + )
313 + return PromptUpdateDetails.select_token_id(token_ids, processor.audio_token_id)
314 +
315 + def get_video_repl(
316 + self,
317 + *,
318 + timestamps: list[float],
319 + num_soft_tokens_per_frame: list[int],
320 + processor: Gemma4Processor,
321 + ) -> PromptUpdateDetails[list[int]]:
322 + """Build the full token replacement for one video.
323 +
324 + Produces the same interleaved sequence as the HF Gemma4Processor:
325 + mm:ss <boi><|video|>*N<eoi> mm:ss <boi><|video|>*N<eoi> ...
326 + """
327 + tokenizer = self.ctx.get_tokenizer()
328 + config = self.get_hf_config()
329 +
330 + boi_token_id = config.boi_token_id
331 + eoi_token_id = config.eoi_token_id
332 + video_token_id = processor.video_token_id
333 +
334 + all_token_ids: list[int] = []
335 + for i, (ts, n_tokens) in enumerate(zip(timestamps, num_soft_tokens_per_frame)):
336 + # mm:ss timestamp — matches transformers: int-truncated,
337 + # zero-padded.
338 + minutes = int(ts // 60)
339 + seconds = int(ts % 60)
340 + ts_str = f"{minutes:02d}:{seconds:02d}"
341 +
342 + prefix = f" {ts_str} " if i > 0 else f"{ts_str} "
343 + ts_token_ids = tokenizer.encode(prefix, add_special_tokens=False)
344 + all_token_ids.extend(ts_token_ids)
345 +
346 + all_token_ids.append(boi_token_id)
347 + all_token_ids.extend([video_token_id] * n_tokens)
348 + all_token_ids.append(eoi_token_id)
349 +
350 + return PromptUpdateDetails.select_token_id(all_token_ids, video_token_id)
351 +
352 +
353 +# ---------------------------------------------------------------------------
354 +# Dummy inputs builder
355 +# ---------------------------------------------------------------------------
356 +
357 +
358 +class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]):
359 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
360 + num_images = mm_counts.get("image", 0)
361 + num_audios = mm_counts.get("audio", 0)
362 + num_videos = mm_counts.get("video", 0)
363 + processor = self.info.get_hf_processor()
364 + # Use image_token (<|image|>) with tab prefix — this is what the
365 + # Gemma4 chat template inserts per image (\t<|image|>).
366 + # _get_prompt_updates targets image_token and expands it to the
367 + # full_image_sequence.
368 + text = ("\t" + processor.image_token) * num_images
369 + if num_audios > 0 and processor.audio_token:
370 + text += processor.audio_token * num_audios
371 + if num_videos > 0:
372 + text += processor.video_token * num_videos
373 + return text
374 +
375 + def get_dummy_mm_data(
376 + self,
377 + seq_len: int,
378 + mm_counts: Mapping[str, int],
379 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
380 + ) -> MultiModalDataDict:
381 + num_images = mm_counts.get("image", 0)
382 + num_audios = mm_counts.get("audio", 0)
383 + num_videos = mm_counts.get("video", 0)
384 + processor = self.info.get_hf_processor()
385 + image_processor = processor.image_processor
386 + # Use processor's configured image size for dummies.
387 + # Gemma4ImageProcessor sets size=None (it uses patch_size /
388 + # max_soft_tokens instead of the standard size dict), so we
389 + # guard against None with `or {}`.
390 + size = getattr(image_processor, "size", None) or {}
391 + img_width = size.get("width", 224)
392 + img_height = size.get("height", 224)
393 +
394 + image_overrides = mm_options.get("image") if mm_options else None
395 + audio_overrides = mm_options.get("audio") if mm_options else None
396 + video_overrides = mm_options.get("video") if mm_options else None
397 +
398 + data: MultiModalDataDict = {
399 + "image": self._get_dummy_images(
400 + width=img_width,
401 + height=img_height,
402 + num_images=num_images,
403 + overrides=image_overrides,
404 + ),
405 + }
406 +
407 + if num_audios > 0:
408 + audio_len = processor.feature_extractor.fft_length
409 + data["audio"] = self._get_dummy_audios(
410 + length=audio_len,
411 + num_audios=num_audios,
412 + overrides=audio_overrides,
413 + )
414 +
415 + if num_videos > 0:
416 + data["video"] = self._get_dummy_videos(
417 + width=img_width,
418 + height=img_height,
419 + num_frames=_VIDEO_MAX_FRAMES,
420 + num_videos=num_videos,
421 + overrides=video_overrides,
422 + )
423 +
424 + return data
425 +
426 + def _get_dummy_videos(
427 + self,
428 + *,
429 + width: int,
430 + height: int,
431 + num_frames: int,
432 + num_videos: int,
433 + overrides: VideoDummyOptions | None = None,
434 + ) -> list[VideoItem]:
435 + num_frames = max(num_frames, 2)
436 + videos = super()._get_dummy_videos(
437 + width=width,
438 + height=height,
439 + num_frames=num_frames,
440 + num_videos=num_videos,
441 + overrides=overrides,
442 + )
443 + videos = [v.copy() for v in videos]
444 +
445 + video_items: list[VideoItem] = []
446 + for video in videos:
447 + video_num_frames = video.shape[0]
448 + video_metadata = {
449 + "fps": 2.0,
450 + "duration": video_num_frames / 2.0,
451 + "total_num_frames": video_num_frames,
452 + "frames_indices": list(range(video_num_frames)),
453 + "video_backend": "opencv",
454 + "do_sample_frames": False,
455 + }
456 + video_items.append((video, video_metadata))
457 +
458 + return video_items
459 +
460 +
461 +# ---------------------------------------------------------------------------
462 +# Multimodal processor
463 +# ---------------------------------------------------------------------------
464 +
465 +
466 +class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
467 + def _call_hf_processor(
468 + self,
469 + prompt: str,
470 + mm_data: Mapping[str, object],
471 + mm_kwargs: Mapping[str, object],
472 + tok_kwargs: Mapping[str, object],
473 + ) -> BatchFeature:
474 + # Validate max_soft_tokens early and exit cleanly on bad values.
475 + _SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120)
476 +
477 + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(mm_kwargs)
478 + val = merged_kwargs.get("max_soft_tokens")
479 + if val is None:
480 + val = merged_kwargs.get("images_kwargs", {}).get("max_soft_tokens")
481 +
482 + if val is not None and val not in _SUPPORTED_SOFT_TOKENS:
483 + logger.error(
484 + "Unsupported max_soft_tokens value: %d. Valid values are %s. Exiting.",
485 + val,
486 + _SUPPORTED_SOFT_TOKENS,
487 + )
488 + sys.exit(1)
489 +
490 + mm_data = dict(mm_data)
491 +
492 + # ---- VIDEO HANDLING ----
493 + # Gemma4 decomposes video into timestamped image frames.
494 + # Each frame is processed with max_soft_tokens=70 through the
495 + # same vision tower, matching transformers processing_gemma4.py.
496 + video_outputs: dict[str, Any] = {}
497 + if videos := mm_data.pop("videos", []):
498 + processor = self.info.get_hf_processor()
499 +
500 + all_video_pixel_values: list[torch.Tensor] = []
501 + all_video_position_ids: list[torch.Tensor] = []
502 + video_num_soft_tokens_per_video: list[list[int]] = []
503 + video_timestamps_per_video: list[list[float]] = []
504 + video_frame_counts: list[int] = []
505 +
506 + for item in videos:
507 + video_array, metadata = item
508 +
509 + # Convert frames to PIL images
510 + if isinstance(video_array, np.ndarray):
511 + frames = [
512 + PILImage.fromarray(video_array[i])
513 + for i in range(video_array.shape[0])
514 + ]
515 + else:
516 + frames = list(video_array)
517 +
518 + # Compute timestamps from metadata (same as transformers)
519 + fps = metadata.get("fps") or 24
520 + frame_indices = metadata.get("frames_indices", list(range(len(frames))))
521 + timestamps = [idx / fps for idx in frame_indices]
522 +
523 + # Process frames as images with max_soft_tokens=70
524 + video_mm_kwargs = dict(mm_kwargs)
525 + video_mm_kwargs["max_soft_tokens"] = _VIDEO_MAX_SOFT_TOKENS
526 +
527 + dummy_prompt = ("\t" + processor.image_token) * len(frames)
528 +
529 + frame_outputs = super()._call_hf_processor(
530 + prompt=dummy_prompt,
531 + mm_data={"images": frames},
532 + mm_kwargs=video_mm_kwargs,
533 + tok_kwargs=tok_kwargs,
534 + )
535 +
536 + # Remap HF key name
537 + if "image_position_ids" in frame_outputs:
538 + frame_outputs["pixel_position_ids"] = frame_outputs.pop(
539 + "image_position_ids"
540 + )
541 +
542 + all_video_pixel_values.append(frame_outputs["pixel_values"])
543 + all_video_position_ids.append(frame_outputs["pixel_position_ids"])
544 +
545 + # Compute soft tokens per frame
546 + num_soft_per_frame = []
547 + for img in frames:
548 + w, h = img.size
549 + n = self.info._compute_num_soft_tokens(
550 + w, h, max_soft_tokens=_VIDEO_MAX_SOFT_TOKENS
551 + )
552 + num_soft_per_frame.append(n)
553 +
554 + video_num_soft_tokens_per_video.append(num_soft_per_frame)
555 + video_timestamps_per_video.append(timestamps)
556 + video_frame_counts.append(len(frames))
557 +
558 + # Build expanded replacement text and replace the
559 + # <|video|> placeholder in the prompt.
560 + # Use split(token, 1) to avoid collision — the
561 + # replacement text itself contains <|video|> tokens.
562 + ts_strs = [f"{int(s // 60):02d}:{int(s % 60):02d}" for s in timestamps]
563 + replacement = " ".join(
564 + f"{t} {processor.boi_token}"
565 + f"{processor.video_token * n}"
566 + f"{processor.eoi_token}"
567 + for t, n in zip(ts_strs, num_soft_per_frame)
568 + )
569 + parts = prompt.split(processor.video_token, 1)
570 + if len(parts) == 2:
571 + prompt = parts[0] + replacement + parts[1]
572 +
573 + video_outputs = {
574 + "pixel_values_videos": torch.cat(all_video_pixel_values, dim=0),
575 + "pixel_position_ids_videos": torch.cat(all_video_position_ids, dim=0),
576 + "video_frame_counts": torch.tensor(video_frame_counts),
577 + "video_num_soft_tokens": video_num_soft_tokens_per_video,
578 + "video_timestamps": video_timestamps_per_video,
579 + }
580 +
581 + # The processor accepts 'audio' not 'audios'.
582 + if "audios" in mm_data:
583 + mm_data["audio"] = mm_data.pop("audios")
584 +
585 + # Warn if any audio waveform exceeds the model's max duration.
586 + if "audio" in mm_data:
587 + processor = self.info.get_hf_processor()
588 + sr = processor.feature_extractor.sampling_rate
589 + max_tokens = processor.audio_seq_length
590 + ms_per_tok = processor.audio_ms_per_token
591 + max_duration_s = max_tokens * ms_per_tok / 1000.0
592 + audios = mm_data["audio"]
593 + if not isinstance(audios, (list, tuple)):
594 + audios = [audios]
595 + for i, waveform in enumerate(audios):
596 + duration_s = len(waveform) / sr
597 + if duration_s > max_duration_s:
598 + logger.warning(
599 + "Audio duration exceeds max: %f > %f seconds",
600 + duration_s,
601 + max_duration_s,
602 + )
603 + # vLLM's call_hf_processor (context.py) re-merges
604 + # mm_processor_kwargs from the model config on every call via:
605 + # config_kwargs | incoming_kwargs (right side wins)
606 + #
607 + # If we strip max_soft_tokens from incoming, the re-merge puts
608 + # back the config's global default (e.g. 280), ignoring any
609 + # per-prompt override. Instead, we keep it in the kwargs with
610 + # the validated per-prompt value so it wins during the merge.
611 + #
612 + # NOTE: This requires a corresponding type annotation on the
613 + # HF side (Gemma4ProcessorKwargs.images_kwargs) so that
614 + # _merge_kwargs routes max_soft_tokens into images_kwargs.
615 + patched_mm_kwargs = dict(mm_kwargs)
616 + if val is not None:
617 + patched_mm_kwargs["max_soft_tokens"] = val
618 +
619 + processed_outputs = super()._call_hf_processor(
620 + prompt,
621 + mm_data,
622 + patched_mm_kwargs,
623 + tok_kwargs,
624 + )
625 +
626 + # HF uses 'image_position_ids'; vLLM uses 'pixel_position_ids'.
627 + # Remap here to keep a single translation point.
628 + if "image_position_ids" in processed_outputs:
629 + processed_outputs["pixel_position_ids"] = processed_outputs.pop(
630 + "image_position_ids"
631 + )
632 +
633 + if "input_features" in processed_outputs:
634 + # Keep padded features for batched audio tower execution.
635 + processed_outputs["input_features_padded"] = processed_outputs[
636 + "input_features"
637 + ]
638 + # Unpad per-item so each item's cache entry is self-contained.
639 + unpadded_features = [
640 + f[mask]
641 + for f, mask in zip(
642 + processed_outputs["input_features"],
643 + processed_outputs["input_features_mask"],
644 + )
645 + ]
646 + processed_outputs["input_features"] = unpadded_features
647 +
648 + # Merge video outputs into the final result
649 + combined_outputs = dict(processed_outputs, **video_outputs)
650 + return BatchFeature(combined_outputs)
651 +
652 + def _get_mm_fields_config(
653 + self,
654 + hf_inputs: BatchFeature,
655 + hf_processor_mm_kwargs: Mapping[str, object],
656 + ) -> Mapping[str, MultiModalFieldConfig]:
657 + fields = dict(
658 + pixel_values=MultiModalFieldConfig.batched("image"),
659 + pixel_position_ids=MultiModalFieldConfig.batched("image"),
660 + input_features_padded=MultiModalFieldConfig.batched("audio"),
661 + input_features_mask=MultiModalFieldConfig.batched("audio"),
662 + )
663 +
664 + # Video fields: frames stored flat, split per video by
665 + # video_frame_counts.
666 + video_frame_counts = hf_inputs.get("video_frame_counts")
667 + if video_frame_counts is not None:
668 + vfc = video_frame_counts
669 + if not isinstance(vfc, torch.Tensor):
670 + vfc = torch.tensor(vfc)
671 + fields.update(
672 + pixel_values_videos=(
673 + MultiModalFieldConfig.flat_from_sizes("video", vfc)
674 + ),
675 + pixel_position_ids_videos=(
676 + MultiModalFieldConfig.flat_from_sizes("video", vfc)
677 + ),
678 + video_frame_counts=MultiModalFieldConfig.batched(
679 + "video",
680 + ),
681 + video_num_soft_tokens=MultiModalFieldConfig.batched(
682 + "video", keep_on_cpu=True
683 + ),
684 + video_timestamps=MultiModalFieldConfig.batched(
685 + "video", keep_on_cpu=True
686 + ),
687 + )
688 +
689 + return fields
690 +
691 + def _get_prompt_updates(
692 + self,
693 + mm_items: MultiModalDataItems,
694 + hf_processor_mm_kwargs: Mapping[str, Any],
695 + out_mm_kwargs: MultiModalKwargsItems,
696 + ) -> Sequence[PromptUpdate]:
697 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
698 +
699 + prompt_updates = []
700 +
701 + if "image" in mm_items:
702 + # Target image_token (<|image|>) — the single placeholder the
703 + # Gemma4 chat template inserts once per image in the prompt.
704 + # vLLM tokenizes the prompt without token expansion, so only
705 + # one image_token exists per image in the token stream.
706 + # The replacement expands it to the full image sequence
707 + # (boi + N×image_token + eoi, where N = max_soft_tokens).
708 + image_token = hf_processor.image_token
709 +
710 + def get_replacement_image(item_idx: int):
711 + images = mm_items.get_items("image", ImageProcessorItems)
712 + image_size = images.get_image_size(item_idx)
713 + # Resolve the effective max_soft_tokens by merging
714 + # per-prompt kwargs with the config-level defaults,
715 + # consistent with how _call_hf_processor resolves it.
716 + # Without this merge, a missing per-prompt override
717 + # would fall back to vision_cfg.default_output_length
718 + # instead of the config's mm_processor_kwargs default.
719 + merged_kwargs = self.info.ctx.get_merged_mm_kwargs(
720 + hf_processor_mm_kwargs,
721 + )
722 + max_soft_tokens = merged_kwargs.get("max_soft_tokens")
723 + return self.info.get_image_repl(
724 + image_width=image_size.width,
725 + image_height=image_size.height,
726 + processor=hf_processor,
727 + max_soft_tokens=max_soft_tokens,
728 + )
729 +
730 + prompt_updates.append(
731 + PromptReplacement(
732 + modality="image",
733 + target=image_token,
734 + replacement=get_replacement_image,
735 + )
736 + )
737 +
738 + if "video" in mm_items:
739 + video_token = hf_processor.video_token
740 +
741 + def get_replacement_video(item_idx: int):
742 + out_item = out_mm_kwargs["video"][item_idx]
743 + timestamps = out_item["video_timestamps"].data
744 + num_soft = out_item["video_num_soft_tokens"].data
745 + return self.info.get_video_repl(
746 + timestamps=timestamps,
747 + num_soft_tokens_per_frame=num_soft,
748 + processor=hf_processor,
749 + )
750 +
751 + prompt_updates.append(
752 + PromptReplacement(
753 + modality="video",
754 + target=video_token,
755 + replacement=get_replacement_video,
756 + )
757 + )
758 +
759 + if "audio" in mm_items:
760 + audio_token = hf_processor.audio_token
761 +
762 + def get_replacement_audio(item_idx: int):
763 + audios = mm_items.get_items("audio", AudioProcessorItems)
764 + audio_len = audios.get_audio_length(item_idx)
765 + return self.info.get_audio_repl(
766 + audio_len=audio_len,
767 + processor=hf_processor,
768 + )
769 +
770 + prompt_updates.append(
771 + PromptReplacement(
772 + modality="audio",
773 + target=audio_token,
774 + replacement=get_replacement_audio,
775 + )
776 + )
777 +
778 + return prompt_updates
779 +
780 + # NOTE: Gemma3/Gemma3n override _apply_token_matches and
781 + # _find_mm_placeholders to merge adjacent newline tokens that arise
782 + # when full_image_sequence contains "\n\n" wrappers. Gemma4's
783 + # full_image_sequence has NO newlines (just BOI + 280×image_token +
784 + # EOI), so the base class implementations work correctly as-is.
785 +
786 +
787 +# ---------------------------------------------------------------------------
788 +# Multimodal embedder
789 +# ---------------------------------------------------------------------------
790 +
791 +
792 +class Gemma4MultimodalEmbedder(nn.Module):
793 + """Projects vision/audio soft tokens into LM embedding space.
794 +
795 + Architecture:
796 + inputs_embeds → embedding_projection → embedding_post_projection_norm
797 +
798 + Unlike Gemma3n which has separate hard/soft embedding paths with
799 + per-path normalization and a learned embedding table, Gemma4 uses a
800 + simplified 2-layer design: a linear projection followed by RMSNorm
801 + (without learnable scale). The checkpoint confirms this — only
802 + ``embedding_projection.weight`` exists; there is no embedding table
803 + or pre-projection norm weights.
804 + """
805 +
806 + def __init__(
807 + self,
808 + multimodal_config: Gemma4VisionConfig | Gemma4AudioConfig,
809 + text_config: Gemma4TextConfig,
810 + ):
811 + super().__init__()
812 +
813 + self.eps = multimodal_config.rms_norm_eps
814 + self.text_hidden_size = text_config.hidden_size
815 +
816 + # Audio tower uses output_proj_dims (1536) rather than hidden_size
817 + # (1024); vision uses hidden_size (768) directly.
818 + embedding_dim = (
819 + getattr(multimodal_config, "output_proj_dims", None)
820 + or multimodal_config.hidden_size
821 + )
822 +
823 + self.embedding_projection = ReplicatedLinear(
824 + embedding_dim,
825 + self.text_hidden_size,
826 + bias=False,
827 + )
828 +
829 + self.embedding_post_projection_norm = RMSNorm(
830 + self.text_hidden_size,
831 + eps=self.eps,
832 + has_weight=False,
833 + )
834 +
835 + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
836 + """Project soft tokens from a multimodal tower into LM space."""
837 + embs_proj, _ = self.embedding_projection(inputs_embeds)
838 + return self.embedding_post_projection_norm(embs_proj)
839 +
840 +
841 +# ---------------------------------------------------------------------------
842 +# Main model
843 +# ---------------------------------------------------------------------------
844 +
845 +
846 +@MULTIMODAL_REGISTRY.register_processor(
847 + Gemma4MultiModalProcessor,
848 + info=Gemma4ProcessingInfo,
849 + dummy_inputs=Gemma4DummyInputsBuilder,
850 +)
851 +class Gemma4ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
852 + packed_modules_mapping = {
853 + "qkv_proj": [
854 + "q_proj",
855 + "k_proj",
856 + "v_proj",
857 + ],
858 + "gate_up_proj": [
859 + "gate_proj",
860 + "up_proj",
861 + ],
862 + }
863 +
864 + # Maps checkpoint prefixes to vLLM module paths.
865 + hf_to_vllm_mapper = WeightsMapper(
866 + orig_to_new_prefix={
867 + "model.embed_audio.": "embed_audio.",
868 + "model.embed_vision.": "embed_vision.",
869 + "model.language_model.": "language_model.model.",
870 + "model.vision_tower.": "vision_tower.",
871 + "model.audio_tower.": "audio_tower.",
872 + "lm_head.": "language_model.lm_head.",
873 + "model": "language_model.model",
874 + }
875 + )
876 +
877 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
878 + super().__init__()
879 + config = vllm_config.model_config.hf_config
880 + quant_config = vllm_config.quant_config
881 + multimodal_config = vllm_config.model_config.multimodal_config
882 + self.config = config
883 + self.quant_config = quant_config
884 + self.multimodal_config = multimodal_config
885 +
886 + # ---- Vision tower (shared by image and video) ----
887 + with self._mark_tower_model(vllm_config, {"image", "video"}):
888 + self.vision_tower = AutoModel.from_config(config=config.vision_config)
889 + self.embed_vision = Gemma4MultimodalEmbedder(
890 + config.vision_config, config.text_config
891 + )
892 +
893 + # ---- Audio tower (variants with audio_config) ----
894 + if config.audio_config is not None:
895 + with self._mark_tower_model(vllm_config, "audio"):
896 + self.audio_tower = AutoModel.from_config(config=config.audio_config)
897 + # AutoModel.from_config does NOT call post_init(),
898 + # which is needed to initialize buffers that are absent
899 + # from the checkpoint (e.g. inv_timescales for relative
900 + # position embeddings, softcap, gradient_clipping).
901 + self.audio_tower.post_init()
902 + self.embed_audio = Gemma4MultimodalEmbedder(
903 + config.audio_config, config.text_config
904 + )
905 + else:
906 + self.audio_tower = None
907 + self.embed_audio = None
908 +
909 + # ---- Language model (vLLM optimised) ----
910 + with self._mark_language_model(vllm_config):
911 + self.language_model: Gemma4ForCausalLM = init_vllm_registered_model(
912 + vllm_config=vllm_config,
913 + hf_config=config.text_config,
914 + prefix=maybe_prefix(prefix, "language_model"),
915 + architectures=["Gemma4ForCausalLM"],
916 + )
917 +
918 + # Pre-allocate PLE buffer for CUDA graph compatibility.
919 + # Some variants have hidden_size_per_layer_input=None (no PLE).
920 + ple_dim = config.text_config.hidden_size_per_layer_input
921 + if ple_dim is not None:
922 + self.per_layer_embeddings = torch.zeros(
923 + vllm_config.scheduler_config.max_num_batched_tokens,
924 + config.text_config.num_hidden_layers,
925 + ple_dim,
926 + device=(self.language_model.model.embed_tokens.weight.device),
927 + dtype=(self.language_model.model.embed_tokens.weight.dtype),
928 + )
929 + else:
930 + self.per_layer_embeddings = None
931 +
932 + self.make_empty_intermediate_tensors = (
933 + self.language_model.make_empty_intermediate_tensors
934 + )
935 +
936 + # --- MixtureOfExperts delegation to language_model ---
937 + self.expert_weights = self.language_model.expert_weights
938 + self.moe_layers = self.language_model.moe_layers
939 + self.num_moe_layers = self.language_model.num_moe_layers
940 + self.num_logical_experts = self.language_model.num_logical_experts
941 + self.num_physical_experts = self.language_model.num_physical_experts
942 + self.num_local_physical_experts = self.language_model.num_local_physical_experts
943 + self.num_routed_experts = self.language_model.num_routed_experts
944 + self.num_expert_groups = self.language_model.num_expert_groups
945 + self.num_shared_experts = self.language_model.num_shared_experts
946 + self.num_redundant_experts = self.language_model.num_redundant_experts
947 +
948 + # ------------------------------------------------------------------ #
949 + # Input parsing
950 + # ------------------------------------------------------------------ #
951 +
952 + def _parse_and_validate_image_input(
953 + self, **kwargs: object
954 + ) -> Gemma4ImageInputs | None:
955 + pixel_values = kwargs.pop("pixel_values", None)
956 + pixel_position_ids = kwargs.pop("pixel_position_ids", None)
957 + image_embeds = kwargs.pop("image_embeds", None)
958 + assert image_embeds is None, "Gemma4 does not support image_embeds."
959 + if pixel_values is None:
960 + return None
961 + return Gemma4ImagePixelInputs(
962 + pixel_values=pixel_values,
963 + pixel_position_ids=pixel_position_ids,
964 + )
965 +
966 + def _parse_and_validate_audio_input(
967 + self, **kwargs: object
968 + ) -> Gemma4AudioInputs | None:
969 + input_features_padded = kwargs.pop("input_features_padded", None)
970 + if input_features_padded is None:
971 + return None
972 + input_features_mask = kwargs.pop("input_features_mask", None)
973 + if input_features_mask is None:
974 + return None
975 + return Gemma4AudioInputs(
976 + input_features_padded=input_features_padded,
977 + input_features_mask=input_features_mask,
978 + )
979 +
980 + def _parse_and_validate_video_input(
981 + self, **kwargs: object
982 + ) -> dict[str, torch.Tensor] | None:
983 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)
984 + pixel_position_ids_videos = kwargs.pop("pixel_position_ids_videos", None)
985 + video_frame_counts = kwargs.pop("video_frame_counts", None)
986 + if pixel_values_videos is None:
987 + return None
988 + return {
989 + "pixel_values_videos": pixel_values_videos,
990 + "pixel_position_ids_videos": pixel_position_ids_videos,
991 + "video_frame_counts": video_frame_counts,
992 + }
993 +
994 + def _parse_and_validate_multimodal_inputs(
995 + self, **kwargs: object
996 + ) -> dict[str, Gemma4ImageInputs | Gemma4AudioInputs | Gemma4VideoInputs | None]:
997 + mm_input_by_modality = {}
998 + for input_key in list(kwargs):
999 + if (
1000 + input_key in ("pixel_values", "image_embeds")
1001 + and "image" not in mm_input_by_modality
1002 + ):
1003 + mm_input_by_modality["image"] = self._parse_and_validate_image_input(
1004 + **kwargs
1005 + )
1006 + if (
1007 + input_key == "pixel_values_videos"
1008 + and "video" not in mm_input_by_modality
1009 + ):
1010 + mm_input_by_modality["video"] = self._parse_and_validate_video_input(
1011 + **kwargs
1012 + )
1013 + if (
1014 + input_key == "input_features_padded"
1015 + and "audio" not in mm_input_by_modality
1016 + ):
1017 + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(
1018 + **kwargs
1019 + )
1020 + return mm_input_by_modality
1021 +
1022 + # ------------------------------------------------------------------ #
1023 + # Image processing
1024 + # ------------------------------------------------------------------ #
1025 +
1026 + def _process_image_input(
1027 + self,
1028 + image_input: Gemma4ImageInputs,
1029 + ) -> list[torch.Tensor]:
1030 + pixel_values = image_input["pixel_values"]
1031 + pixel_position_ids = image_input["pixel_position_ids"]
1032 +
1033 + # The HF image processor now outputs pre-patchified data:
1034 + # pixel_values: (num_images, max_patches, patch_pixels)
1035 + # pixel_position_ids: (num_images, max_patches, 2)
1036 + # We call the vision tower's forward() directly, which handles
1037 + # patch embedding, encoding, pooling, padding removal, and
1038 + # optional standardization internally.
1039 + vt = self.vision_tower
1040 + pooling_k2 = self.config.vision_config.pooling_kernel_size**2
1041 +
1042 + # TODO: Move this per-image loop into the input processor to
1043 + # reduce dynamism at the model runner / engine core. This
1044 + # requires spatially padding all images to uniform (H_max,
1045 + # W_max) in _call_hf_processor() so they arrive as a single
1046 + # stacked tensor, tracking padded regions via image_sizes
1047 + # metadata, and validating numerical equivalence with the
1048 + # current per-image path.
1049 + #
1050 + # Process each image individually through the vision tower.
1051 + # The vision tower's forward() strips padding and returns a
1052 + # flat tensor of valid tokens. We process per-image to get
1053 + # variable-length outputs matching the dynamic token count
1054 + # from get_image_repl.
1055 + per_image_features = []
1056 + for i in range(pixel_values.shape[0]):
1057 + pv = pixel_values[i].unsqueeze(0) # (1, max_patches, patch_pixels)
1058 + pp = pixel_position_ids[i].unsqueeze(0) # (1, max_patches, 2)
1059 +
1060 + # Derive the pooler's output_length from the total patch
1061 + # count (including padding). The vision tower encoder
1062 + # processes ALL patches — padding patches get zero hidden
1063 + # states but still occupy sequence positions. The pooler's
1064 + # _avg_pool_by_positions requires:
1065 + # input_seq_len / output_length == k²
1066 + # where k == pooling_kernel_size. The image processor
1067 + # allocates max_patches = max_soft_tokens * k² total slots,
1068 + # so output_length = max_patches / k² == max_soft_tokens.
1069 + # Without this, the pooler falls back to
1070 + # config.image_seq_length (e.g. 280), which fails when a
1071 + # different max_soft_tokens was used at preprocessing time.
1072 + max_patches = pv.shape[1]
1073 + output_length = max_patches // pooling_k2
1074 +
1075 + vt_output = vt(pv, pp, output_length=output_length)
1076 + # last_hidden_state: (num_valid_tokens, hidden_size)
1077 + # — already flat with padding stripped by the vision tower
1078 + per_image_features.append(vt_output.last_hidden_state)
1079 +
1080 + # Project each image's features into LM embedding space.
1081 + # Per-image loop is required because images have variable
1082 + # token counts after padding removal.
1083 + # Cast to match the projection layer's dtype (model may be
1084 + # bf16 while the vision tower outputs fp32).
1085 + target_dtype = self.embed_vision.embedding_projection.weight.dtype
1086 + return [
1087 + self.embed_vision(inputs_embeds=img.unsqueeze(0).to(target_dtype)).squeeze(
1088 + 0
1089 + )
1090 + for img in per_image_features
1091 + ]
1092 +
1093 + # ------------------------------------------------------------------ #
1094 + # Video processing (frames through vision tower)
1095 + # ------------------------------------------------------------------ #
1096 +
1097 + def _process_video_input(
1098 + self,
1099 + video_input: dict[str, torch.Tensor],
1100 + ) -> list[torch.Tensor]:
1101 + """Process video frames through the vision tower.
1102 +
1103 + Reuses the image processing pipeline — Gemma4 has no separate
1104 + video tower; video frames are just images at lower resolution
1105 + (max_soft_tokens=70).
1106 +
1107 + Returns one concatenated embedding tensor per video (not per
1108 + frame), because vLLM treats one video as one multimodal item.
1109 + The flat_from_sizes field config groups all frames of a video
1110 + together, so embed_multimodal must return one tensor per video.
1111 + """
1112 + pixel_values = video_input["pixel_values_videos"]
1113 + pixel_position_ids = video_input["pixel_position_ids_videos"]
1114 + frame_counts = video_input["video_frame_counts"]
1115 +
1116 + vt = self.vision_tower
1117 + pooling_k2 = self.config.vision_config.pooling_kernel_size**2
1118 + target_dtype = self.embed_vision.embedding_projection.weight.dtype
1119 +
1120 + # Split flat tensors into per-video chunks
1121 + if isinstance(frame_counts, torch.Tensor):
1122 + fc_list = frame_counts.tolist()
1123 + else:
1124 + fc_list = list(frame_counts)
1125 +
1126 + pv_per_video = torch.split(pixel_values, fc_list, dim=0)
1127 + pp_per_video = torch.split(pixel_position_ids, fc_list, dim=0)
1128 +
1129 + per_video_embeddings = []
1130 + for pv_chunk, pp_chunk in zip(pv_per_video, pp_per_video):
1131 + frame_embs = []
1132 + for i in range(pv_chunk.shape[0]):
1133 + pv = pv_chunk[i].unsqueeze(0)
1134 + pp = pp_chunk[i].unsqueeze(0)
1135 +
1136 + max_patches = pv.shape[1]
1137 + output_length = max_patches // pooling_k2
1138 +
1139 + vt_output = vt(pv, pp, output_length=output_length)
1140 + frame_emb = self.embed_vision(
1141 + inputs_embeds=(
1142 + vt_output.last_hidden_state.unsqueeze(0).to(target_dtype)
1143 + )
1144 + ).squeeze(0)
1145 + frame_embs.append(frame_emb)
1146 +
1147 + # Concatenate all frames of this video into one tensor.
1148 + per_video_embeddings.append(torch.cat(frame_embs, dim=0))
1149 +
1150 + return per_video_embeddings
1151 +
1152 + # ------------------------------------------------------------------ #
1153 + # Audio processing
1154 + # ------------------------------------------------------------------ #
1155 +
1156 + def _process_audio_input(
1157 + self,
1158 + audio_input: Gemma4AudioInputs,
1159 + ) -> list[torch.Tensor]:
1160 + input_features = audio_input["input_features_padded"].squeeze(1)
1161 + input_features_mask = audio_input["input_features_mask"].squeeze(1)
1162 +
1163 + # Run audio tower — mask uses standard HF convention
1164 + # (True=valid, False=padding).
1165 + audio_outputs = self.audio_tower(input_features, input_features_mask)
1166 + if isinstance(audio_outputs, tuple):
1167 + audio_encodings, audio_mask = audio_outputs
1168 + else:
1169 + audio_encodings = audio_outputs.last_hidden_state
1170 + audio_mask = audio_outputs.attention_mask
1171 +
1172 + # Project into LM embedding space.
1173 + audio_features = self.embed_audio(inputs_embeds=audio_encodings)
1174 +
1175 + # Strip padding per-batch element: only keep real (non-padding)
1176 + # tokens. audio_mask is True for valid positions (HF convention).
1177 + per_audio = []
1178 + for enc, mask in zip(audio_features, audio_mask, strict=True):
1179 + per_audio.append(enc[mask]) # [num_real, hidden_size]
1180 +
1181 + return per_audio
1182 +
1183 + # ------------------------------------------------------------------ #
1184 + # MultiModalEmbeddings interface
1185 + # ------------------------------------------------------------------ #
1186 +
1187 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1188 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
1189 + multimodal_embeddings: list[torch.Tensor] = []
1190 +
1191 + for modality, multimodal_input in mm_input_by_modality.items():
1192 + if multimodal_input is None:
1193 + continue
1194 + if modality == "image":
1195 + multimodal_embeddings.extend(
1196 + self._process_image_input(multimodal_input)
1197 + )
1198 + elif modality == "video":
1199 + multimodal_embeddings.extend(
1200 + self._process_video_input(multimodal_input)
1201 + )
1202 + elif modality == "audio":
1203 + multimodal_embeddings.extend(
1204 + self._process_audio_input(multimodal_input)
1205 + )
1206 +
1207 + return multimodal_embeddings
1208 +
1209 + def embed_input_ids(
1210 + self,
1211 + input_ids: torch.Tensor,
1212 + multimodal_embeddings: MultiModalEmbeddings | None = None,
1213 + *,
1214 + is_multimodal: torch.Tensor | None = None,
1215 + ) -> torch.Tensor:
1216 + # Cache per-layer embeddings (PLE) for the language model's
1217 + # forward pass. During profiling embed_input_ids is not called,
1218 + # so the pre-allocated zeros are used instead.
1219 + if self.per_layer_embeddings is not None:
1220 + # Mask multimodal tokens (image/audio) to 0 for PLE
1221 + # computation (using token_type_ids == 0 as text_mask).
1222 + # Replicate this: map image token positions to token 0.
1223 + if is_multimodal is not None:
1224 + is_multimodal = is_multimodal.to(input_ids.device)
1225 + ple_input_ids = torch.where(
1226 + is_multimodal, torch.zeros_like(input_ids), input_ids
1227 + )
1228 + else:
1229 + ple_input_ids = input_ids
1230 +
1231 + per_layer_inputs = self.language_model.model.get_per_layer_inputs(
1232 + ple_input_ids
1233 + )
1234 + if per_layer_inputs is not None:
1235 + per_layer_inputs = per_layer_inputs.reshape(
1236 + -1,
1237 + self.config.text_config.num_hidden_layers,
1238 + self.config.text_config.hidden_size_per_layer_input,
1239 + )
1240 + self.per_layer_embeddings[: per_layer_inputs.shape[0]].copy_(
1241 + per_layer_inputs
1242 + )
1243 +
1244 + if multimodal_embeddings is None or is_multimodal is None:
1245 + return super().embed_input_ids(input_ids)
1246 +
1247 + return super().embed_input_ids(
1248 + input_ids,
1249 + multimodal_embeddings=multimodal_embeddings,
1250 + is_multimodal=is_multimodal,
1251 + )
1252 +
1253 + # ------------------------------------------------------------------ #
1254 + # Forward
1255 + # ------------------------------------------------------------------ #
1256 +
1257 + def forward(
1258 + self,
1259 + input_ids: torch.Tensor,
1260 + positions: torch.Tensor,
1261 + intermediate_tensors: IntermediateTensors | None = None,
1262 + inputs_embeds: torch.Tensor | None = None,
1263 + **kwargs: object,
1264 + ) -> IntermediateTensors:
1265 + if intermediate_tensors is not None:
1266 + inputs_embeds = None
1267 +
1268 + # Select the pre-cached PLEs for this batch (None when PLE
1269 + # is disabled for variants without PLE).
1270 + per_layer_inputs = (
1271 + self.per_layer_embeddings[: inputs_embeds.shape[0]]
1272 + if self.per_layer_embeddings is not None and inputs_embeds is not None
1273 + else None
1274 + )
1275 +
1276 + hidden_states = self.language_model.model(
1277 + input_ids,
1278 + positions,
1279 + per_layer_inputs=per_layer_inputs,
1280 + intermediate_tensors=intermediate_tensors,
1281 + inputs_embeds=inputs_embeds,
1282 + **kwargs,
1283 + )
1284 +
1285 + return hidden_states
1286 +
1287 + def compute_logits(
1288 + self,
1289 + hidden_states: torch.Tensor,
1290 + ) -> torch.Tensor | None:
1291 + return self.language_model.compute_logits(hidden_states)
1292 +
1293 + # ------------------------------------------------------------------ #
1294 + # Weight loading
1295 + # ------------------------------------------------------------------ #
1296 +
1297 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1298 + # Some checkpoints have vestigial embed_vision.embedding and
1299 + # embed_audio.embedding weights from the Gemma3n architecture
1300 + # that are not used by Gemma4's MultimodalEmbedder (which only
1301 + # has embedding_projection + embedding_post_projection_norm).
1302 + ignore_prefixes = [
1303 + "embed_vision.embedding.",
1304 + "embed_audio.embedding.",
1305 + ]
1306 + # Models without audio tower should skip
1307 + # audio weights entirely.
1308 + if self.audio_tower is None:
1309 + ignore_prefixes.extend(
1310 + [
1311 + "audio_tower.",
1312 + "embed_audio.",
1313 + ]
1314 + )
1315 + loader = AutoWeightsLoader(
1316 + self,
1317 + ignore_unexpected_prefixes=ignore_prefixes,
1318 + )
1319 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1320 +
1321 + # ------------------------------------------------------------------ #
1322 + # LoRA / multimodal mapping
1323 + # ------------------------------------------------------------------ #
1324 +
1325 + def get_mm_mapping(self) -> MultiModelKeys:
1326 + """Get the module prefix mapping for multimodal models."""
1327 + return MultiModelKeys.from_string_field(
1328 + language_model="language_model",
1329 + connector=["embed_vision", "embed_audio"],
1330 + tower_model=["vision_tower", "audio_tower"],
1331 + )
1332 +
1333 + @classmethod
1334 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1335 + if modality == "image":
1336 + return "<image_soft_token>"
1337 + if modality == "audio":
1338 + return "<audio_soft_token>"
1339 + if modality == "video":
1340 + return "<|video|>"
1341 + raise ValueError(f"Unsupported modality: {modality}")
vllm/model_executor/models/gemma4_utils.py new
+292
@@ -0,0 +1,292 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
4 +
5 +"""Gemma4 output parsing utilities for offline inference.
6 +
7 +Standalone functions that parse decoded model text to extract structured
8 +thinking content and tool calls from Gemma4 models. These are pure-Python
9 +utilities with zero heavy dependencies — they work on raw decoded strings
10 +from any inference backend (vLLM, HuggingFace, TGI, etc.).
11 +
12 +Usage with vLLM offline inference::
13 +
14 + from vllm import LLM, SamplingParams
15 + from vllm.model_executor.models.gemma4_utils import (
16 + parse_output,
17 + parse_tool_calls,
18 + )
19 +
20 + llm = LLM(model="google/gemma-4-it")
21 + outputs = llm.generate(prompt, SamplingParams(...))
22 + text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False)
23 +
24 + # Extract thinking / answer (works with or without enable_thinking)
25 + result = parse_output(text)
26 + print(result["thinking"]) # chain-of-thought or None
27 + print(result["answer"]) # final answer
28 +
29 + # Extract tool calls
30 + tool_calls = parse_tool_calls(text)
31 + for tc in tool_calls:
32 + print(f"{tc['name']}({tc['arguments']})")
33 +
34 +Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users
35 +do not need a transformers dependency for output parsing.
36 +"""
37 +
38 +import json
39 +
40 +import regex as re
41 +
42 +# ---- Thinking Mode Utility ----
43 +
44 +# Thinking delimiter tokens as they appear in decoded text.
45 +# Gemma4 uses <|channel> (start) and <channel|> (end) as thinking delimiters.
46 +_THINKING_START_TAG = "<|channel>"
47 +_THINKING_END_TAG = "<channel|>"
48 +
49 +# Sentinel tokens that may appear in decoded output.
50 +_TURN_END_TAG = "<turn|>"
51 +
52 +
53 +def parse_thinking_output(text: str) -> dict[str, str | None]:
54 + """Parse decoded Gemma4 model output.
55 +
56 + Use this on **all** Gemma4 output regardless of whether thinking mode
57 + was enabled. It handles three cases:
58 +
59 + 1. **Thinking enabled, tags present** — splits on ``<|channel>``/
60 + ``<channel|>`` to separate chain-of-thought from the answer and
61 + strips the ``thought\\n`` role label.
62 + 2. **Thinking disabled, spurious label** — strips the bare
63 + ``thought\\n`` prefix that some Gemma4 models emit even
64 + without thinking mode.
65 + 3. **Clean output** — returns the text unchanged.
66 +
67 + The answer text is always cleaned of trailing sentinel tokens
68 + (``<turn|>``, ``<eos>``, etc.).
69 +
70 + Args:
71 + text: Decoded model output text (from ``tokenizer.decode(...)``).
72 +
73 + Returns:
74 + A dict with keys:
75 + - ``"thinking"``: The chain-of-thought text, or ``None`` if no
76 + thinking delimiters were found.
77 + - ``"answer"``: The final answer text.
78 +
79 + Example::
80 +
81 + >>> from vllm.model_executor.models.gemma4_utils import parse_thinking_output
82 + >>> output_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
83 + >>> result = parse_thinking_output(output_text)
84 + >>> print(result["thinking"]) # chain-of-thought reasoning or None
85 + >>> print(result["answer"]) # final answer
86 + """
87 + if _THINKING_END_TAG in text:
88 + parts = text.split(_THINKING_END_TAG, 1)
89 + thinking_block = parts[0]
90 + answer = _clean_answer(parts[1])
91 +
92 + # Extract thinking content: strip the start tag if present
93 + if _THINKING_START_TAG in thinking_block:
94 + thinking = thinking_block.split(_THINKING_START_TAG, 1)[1]
95 + else:
96 + thinking = thinking_block
97 +
98 + # Strip the "thought\n" channel role label the model emits inside
99 + # <|channel>thought\n...<channel|> (analogous to "user\n" in
100 + # <|turn>user\n...<turn|>).
101 + thinking = _strip_thought_label(thinking.strip())
102 + thinking = thinking.strip()
103 +
104 + return {"thinking": thinking, "answer": answer}
105 +
106 + # No thinking delimiters found.
107 + # Strip spurious "thought\n" role label that some Gemma4 models sometimes
108 + # emit even without thinking mode enabled, then clean trailing tokens.
109 + answer = _strip_thought_label(text)
110 + answer = _clean_answer(answer)
111 + return {"thinking": None, "answer": answer}
112 +
113 +
114 +def _strip_thought_label(text: str) -> str:
115 + """Strip the spurious ``thought\\n`` label from the start of text.
116 +
117 + Only strips when ``thought`` appears as the very first word followed by
118 + a newline — preserving the word ``thought`` in any other context.
119 + """
120 + if text.startswith("thought\n"):
121 + return text[len("thought\n") :]
122 + return text
123 +
124 +
125 +def _clean_answer(text: str) -> str:
126 + """Clean trailing sentinel tokens from the answer text.
127 +
128 + Strips ``<turn|>``, ``<eos>``, and surrounding whitespace that the
129 + model appends at the end of its response.
130 + """
131 + text = text.strip()
132 + # Strip trailing <turn|> (Gemma4 turn-end marker)
133 + if text.endswith(_TURN_END_TAG):
134 + text = text[: -len(_TURN_END_TAG)].rstrip()
135 + # Strip trailing <eos> if present
136 + if text.endswith("<eos>"):
137 + text = text[:-5].rstrip()
138 + return text
139 +
140 +
141 +# ---- Tool Call Parsing Utility ----
142 +#
143 +# NOTE: For the OpenAI-compatible API server tool parser (streaming +
144 +# non-streaming), see vllm/tool_parsers/gemma4_tool_parser.py.
145 +# This module provides offline inference utilities for direct user import.
146 +
147 +# Tool call delimiter tokens as they appear in decoded text.
148 +# Standard format: <|tool_call>call:name{args}<tool_call|>
149 +_TOOL_CALL_START_TAG = "<|tool_call>"
150 +_TOOL_CALL_END_TAG = "<tool_call|>"
151 +_TOOL_RESPONSE_START_TAG = "<|tool_response>"
152 +
153 +# Gemma4 escape token as it appears in decoded text.
154 +_ESCAPE_TOKEN = '<|"|>'
155 +
156 +
157 +def _parse_tool_arguments(args_str: str) -> dict[str, str]:
158 + """Parse tool call arguments from the Gemma4 compact format.
159 +
160 + Handles the ``key:<|"|>value<|"|>`` format used by Gemma4, with fallback
161 + to heuristic key-value extraction. Also tolerates the slightly different
162 + ``key: "value"`` format (space + plain quotes) that some chat templates
163 + produce.
164 +
165 + Args:
166 + args_str: Raw argument string from inside ``call:name{...}``.
167 +
168 + Returns:
169 + Dictionary of argument name → value.
170 + """
171 + if not args_str or not args_str.strip():
172 + return {}
173 +
174 + # Replace Gemma4 escape tokens with standard quotes.
175 + cleaned = args_str.replace(_ESCAPE_TOKEN, '"')
176 +
177 + # Try JSON parsing first (handles nested values, arrays, etc.).
178 + try:
179 + parsed = json.loads("{" + cleaned + "}")
180 + # Ensure all values are strings for consistency.
181 + return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}
182 + except (json.JSONDecodeError, ValueError):
183 + pass
184 +
185 + # Fallback: extract key:"value" pairs (allow optional space after colon).
186 + arguments = {}
187 + for key, value in re.findall(r'(\w+):\s*"([^"]*)"', cleaned):
188 + arguments[key] = value
189 +
190 + if not arguments:
191 + # Last resort: extract key:value pairs (unquoted).
192 + for key, value in re.findall(r"(\w+):\s*([^,}]+)", args_str):
193 + arguments[key] = value.strip().strip('"').replace(_ESCAPE_TOKEN, "")
194 +
195 + return arguments
196 +
197 +
198 +def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]:
199 + """Parse tool calls from decoded Gemma4 model output.
200 +
201 + Uses a tiered parsing strategy to handle known output variations in
202 + Gemma4 models, which may emit
203 + non-standard tool call formats.
204 +
205 + Parsing tiers:
206 + 1. **Standard**: ``<|tool_call>call:name{args}<tool_call|>``
207 + (special token IDs 48/49 in decoded text)
208 + 2. **Fallback** (when ``strict=False``): bare ``call:name{args}``
209 + patterns, including ``<call>name{args}`` (fragmented tokens from
210 + multimodal inputs)
211 +
212 + Args:
213 + text: Decoded model output text (from ``tokenizer.decode(...,
214 + skip_special_tokens=False)``).
215 + strict: If ``True``, only match the standard ``<|tool_call>`` format.
216 + If ``False`` (default), also try fallback patterns for
217 + known Gemma4 output variations.
218 +
219 + Returns:
220 + A list of dicts, each with keys:
221 + - ``"name"``: The tool function name (e.g. ``"get_weather"``).
222 + - ``"arguments"``: A dict of argument name → value.
223 +
224 + Example::
225 +
226 + >>> from vllm.model_executor.models.gemma4_utils import (
227 + ... parse_tool_calls
228 + ... )
229 + >>> output = tokenizer.decode(outputs[0], skip_special_tokens=False)
230 + >>> tool_calls = parse_tool_calls(output)
231 + >>> for tc in tool_calls:
232 + ... print(f"Call: {tc['name']}({tc['arguments']})")
233 + """
234 + results = []
235 +
236 + # Tier 1: Standard format with special tokens.
237 + # <|tool_call>call:name{args}<tool_call|>
238 + # Note: Some Gemma4 models emit <turn|> instead of <tool_call|>.
239 + standard_pattern = r"<\|tool_call\>call:(\w+)\{(.*?)\}(?:<tool_call\|>|<turn\|>)"
240 + for match in re.finditer(standard_pattern, text, re.DOTALL):
241 + name, args_str = match.group(1), match.group(2)
242 + results.append(
243 + {
244 + "name": name,
245 + "arguments": _parse_tool_arguments(args_str),
246 + }
247 + )
248 +
249 + if results or strict:
250 + return results
251 +
252 + # Tier 2: Fallback for known Gemma4 output variations.
253 + # Matches: <call>name{args}, call:name{args}, or bare call:name{args}<eos>
254 + fallback_pattern = r"(?:<call>|(?:^|\s)call:)(\w+)\{(.*?)\}"
255 + for match in re.finditer(fallback_pattern, text, re.DOTALL):
256 + name, args_str = match.group(1), match.group(2)
257 + results.append(
258 + {
259 + "name": name,
260 + "arguments": _parse_tool_arguments(args_str),
261 + }
262 + )
263 +
264 + return results
265 +
266 +
267 +def has_tool_response_tag(text: str) -> bool:
268 + """Check if model output properly ends with a tool response tag.
269 +
270 + Some Gemma4 models sometimes emit ``<eos>`` instead of
271 + ``<|tool_response>`` after a tool call. This helper detects
272 + whether the model used the proper termination, so callers can
273 + decide whether to inject ``<|tool_response>`` into the next prompt.
274 +
275 + Args:
276 + text: Decoded model output text.
277 +
278 + Returns:
279 + ``True`` if the output ends with ``<|tool_response>``
280 + (proper behavior), ``False`` otherwise.
281 +
282 + Example::
283 +
284 + >>> from vllm.model_executor.models.gemma4_utils import (
285 + ... has_tool_response_tag
286 + ... )
287 + >>> if not has_tool_response_tag(model_output):
288 + ... # Model used <eos> instead — inject <|tool_response> manually
289 + ... next_prompt = "<|tool_response>" + tool_result
290 + """
291 + stripped = text.rstrip()
292 + return stripped.endswith(_TOOL_RESPONSE_START_TAG)
vllm/model_executor/models/registry.py
+2
@@ -111,6 +111,7 @@ _TEXT_GENERATION_MODELS = {
111 "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),
112 "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),
113 "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
114 + "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"),
115 "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
116 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
117 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
@@ -383,6 +384,7 @@ _MULTIMODAL_MODELS = {
384 "gemma3n_mm",
385 "Gemma3nForConditionalGeneration",
386 ),
387 + "Gemma4ForConditionalGeneration": ("gemma4_mm", "Gemma4ForConditionalGeneration"),
388 "GlmAsrForConditionalGeneration": ("glmasr", "GlmAsrForConditionalGeneration"),
389 "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),
390 "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),
vllm/model_executor/models/utils.py
+8 -1
@@ -233,8 +233,15 @@ class AutoWeightsLoader:
233 ):
234 """
235 Add tensor names that are not in the model params that may be in the
236 - safetensors, e.g., batch normalization stats.
236 + safetensors, e.g., batch normalization stats and registered buffers.
237 """
238 + # Add persistent registered buffers.
239 + # Non-persistent buffers are excluded, matching PyTorch state_dict().
240 + non_persistent = getattr(module, "_non_persistent_buffers_set", set())
241 + for buf_name, buf in module.named_buffers(recurse=False):
242 + if buf_name not in child_params and buf_name not in non_persistent:
243 + child_params[buf_name] = buf
244 +
245 if isinstance(
246 module,
247 (
vllm/reasoning/__init__.py
+4
@@ -32,6 +32,10 @@ _REASONING_PARSERS_TO_REGISTER = {
32 "ernie45_reasoning_parser",
33 "Ernie45ReasoningParser",
34 ),
35 + "gemma4": (
36 + "gemma4_reasoning_parser",
37 + "Gemma4ReasoningParser",
38 + ),
39 "glm45": (
40 "deepseek_v3_reasoning_parser",
41 "DeepSeekV3ReasoningWithThinkingParser",
vllm/reasoning/gemma4_reasoning_parser.py new
+193
@@ -0,0 +1,193 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Sequence
5 +from typing import TYPE_CHECKING
6 +
7 +from vllm.entrypoints.openai.engine.protocol import DeltaMessage
8 +from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
9 +from vllm.tokenizers import TokenizerLike
10 +
11 +if TYPE_CHECKING:
12 + from vllm.entrypoints.openai.chat_completion.protocol import (
13 + ChatCompletionRequest,
14 + )
15 + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
16 +
17 +# Role label that Gemma4 emits at the start of the thinking channel.
18 +# The model generates: <|channel>thought\n...reasoning...<channel|>
19 +# This prefix must be stripped to expose only the actual reasoning content.
20 +_THOUGHT_PREFIX = "thought\n"
21 +
22 +
23 +class Gemma4ReasoningParser(BaseThinkingReasoningParser):
24 + """
25 + Reasoning parser for Google Gemma4 thinking models.
26 +
27 + Gemma4 uses <|channel>...<channel|> tokens to delimit reasoning/thinking
28 + content within its output. Thinking mode is activated by passing
29 + ``enable_thinking=True`` in the chat template kwargs, which injects a
30 + system turn containing <|think|> (token 98) to trigger chain-of-thought
31 + reasoning.
32 +
33 + Output pattern when thinking is enabled::
34 +
35 + <|channel>thought
36 + ...chain of thought reasoning...<channel|>
37 + Final answer text here.
38 +
39 + The ``thought\\n`` role label inside the channel delimiters is a
40 + structural artefact (analogous to ``user\\n`` in ``<|turn>user\\n...``).
41 + This parser strips it so that downstream consumers see only the
42 + actual reasoning text, consistent with the offline parser
43 + (``vllm.reasoning.gemma4_utils._strip_thought_label``).
44 + """
45 +
46 + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
47 + super().__init__(tokenizer, *args, **kwargs)
48 + # Instance state for streaming prefix stripping.
49 + # Tracks only the reasoning text received from the base parser,
50 + # independent of current_text (which may contain pre-reasoning
51 + # content and lacks special token text due to
52 + # skip_special_tokens=True).
53 + self._reasoning_text: str = ""
54 + self._prefix_stripped: bool = False
55 +
56 + @property
57 + def start_token(self) -> str:
58 + """The token that starts reasoning content."""
59 + return "<|channel>"
60 +
61 + @property
62 + def end_token(self) -> str:
63 + """The token that ends reasoning content."""
64 + return "<channel|>"
65 +
66 + # ------------------------------------------------------------------
67 + # Non-streaming path
68 + # ------------------------------------------------------------------
69 +
70 + def extract_reasoning(
71 + self,
72 + model_output: str,
73 + request: "ChatCompletionRequest | ResponsesRequest",
74 + ) -> tuple[str | None, str | None]:
75 + """Extract reasoning, stripping the ``thought\\n`` role label."""
76 + if self.start_token not in model_output and self.end_token not in model_output:
77 + # Default to content history if no tags are present
78 + # (or if they were stripped)
79 + return None, model_output
80 +
81 + reasoning, content = super().extract_reasoning(model_output, request)
82 + if reasoning is not None:
83 + reasoning = _strip_thought_label(reasoning)
84 + return reasoning, content
85 +
86 + # ------------------------------------------------------------------
87 + # Streaming path
88 + # ------------------------------------------------------------------
89 +
90 + def extract_reasoning_streaming(
91 + self,
92 + previous_text: str,
93 + current_text: str,
94 + delta_text: str,
95 + previous_token_ids: Sequence[int],
96 + current_token_ids: Sequence[int],
97 + delta_token_ids: Sequence[int],
98 + ) -> DeltaMessage | None:
99 + """Extract streaming reasoning, stripping ``thought\\n`` from the
100 + first reasoning delta(s).
101 +
102 + The ``thought\\n`` prefix may arrive as a single delta or split
103 + across multiple deltas (e.g. ``"thought"`` then ``"\\n"``). We
104 + buffer early reasoning tokens until we can determine whether the
105 + prefix is present, then emit the buffered content minus the
106 + prefix.
107 +
108 + Unlike the previous implementation which reconstructed accumulated
109 + reasoning from ``current_text``, this uses instance state
110 + (``_reasoning_text``) to track only the reasoning content returned
111 + by the base parser. This is necessary because
112 + ``skip_special_tokens=True`` (the vLLM default) causes the
113 + ``<|channel>`` delimiter to be invisible in ``current_text``,
114 + making it impossible to separate pre-reasoning content from
115 + reasoning content via string matching.
116 + """
117 + result = super().extract_reasoning_streaming(
118 + previous_text,
119 + current_text,
120 + delta_text,
121 + previous_token_ids,
122 + current_token_ids,
123 + delta_token_ids,
124 + )
125 + if result is None:
126 + return None
127 +
128 + if result.reasoning is None:
129 + return result
130 +
131 + # Accumulate ONLY the reasoning text from base parser results.
132 + # This is immune to pre-reasoning content pollution.
133 + self._reasoning_text += result.reasoning
134 +
135 + # Once the prefix has been handled, all subsequent reasoning
136 + # deltas pass through unchanged.
137 + if self._prefix_stripped:
138 + return result
139 +
140 + # ---- Prefix stripping logic ----
141 +
142 + # Case 1: We've accumulated enough to confirm the prefix is
143 + # present. Strip it and pass through the remainder.
144 + if self._reasoning_text.startswith(_THOUGHT_PREFIX):
145 + prefix_len = len(_THOUGHT_PREFIX)
146 + # How much reasoning was accumulated before this delta?
147 + prev_reasoning_len = len(self._reasoning_text) - len(result.reasoning)
148 + if prev_reasoning_len >= prefix_len:
149 + # Prefix was already consumed by prior deltas; this
150 + # delta is entirely real content — pass through.
151 + self._prefix_stripped = True
152 + return result
153 + else:
154 + # Part or all of the prefix is in this delta.
155 + chars_of_prefix_in_delta = prefix_len - prev_reasoning_len
156 + stripped = result.reasoning[chars_of_prefix_in_delta:]
157 + if stripped:
158 + self._prefix_stripped = True
159 + result.reasoning = stripped
160 + return result
161 + else:
162 + # This entire delta was prefix — suppress it.
163 + # Don't set _prefix_stripped yet; there may be more
164 + # prefix chars to consume in the next delta.
165 + if len(self._reasoning_text) >= prefix_len:
166 + self._prefix_stripped = True
167 + return None
168 +
169 + # Case 2: Accumulated text is a strict prefix of
170 + # _THOUGHT_PREFIX (e.g. we've only seen "thou" so far).
171 + # Buffer by suppressing — we can't yet tell if this will
172 + # become the full prefix or diverge.
173 + if _THOUGHT_PREFIX.startswith(self._reasoning_text):
174 + return None
175 +
176 + # Case 3: Accumulated text doesn't match the thought prefix
177 + # at all. This means prior deltas were buffered (suppressed
178 + # by Case 2) but the text diverged. Re-emit the full
179 + # accumulated text to avoid data loss.
180 + self._prefix_stripped = True
181 + result.reasoning = self._reasoning_text
182 + return result
183 +
184 +
185 +def _strip_thought_label(text: str) -> str:
186 + """Remove the ``thought\\n`` role label from the beginning of text.
187 +
188 + Mirrors ``vllm.reasoning.gemma4_utils._strip_thought_label`` from the
189 + offline parser.
190 + """
191 + if text.startswith(_THOUGHT_PREFIX):
192 + return text[len(_THOUGHT_PREFIX) :]
193 + return text
vllm/reasoning/gemma4_utils.py new
+130
@@ -0,0 +1,130 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
4 +
5 +"""Gemma4 thinking/reasoning output parsing utilities for offline inference.
6 +
7 +Standalone functions that parse decoded model text to extract structured
8 +thinking content from Gemma4 models. These are pure-Python utilities with
9 +zero heavy dependencies — they work on raw decoded strings from any
10 +inference backend (vLLM, HuggingFace, TGI, etc.).
11 +
12 +For the OpenAI-compatible API reasoning parser (streaming +
13 +non-streaming), see ``vllm.reasoning.gemma4_reasoning_parser``.
14 +For tool call parsing, see ``vllm.tool_parsers.gemma4_utils``.
15 +
16 +Usage with vLLM offline inference::
17 +
18 + from vllm import LLM, SamplingParams
19 + from vllm.reasoning.gemma4_utils import parse_thinking_output
20 +
21 + llm = LLM(model="google/gemma-4-it")
22 + outputs = llm.generate(prompt, SamplingParams(...))
23 + text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False)
24 +
25 + # Extract thinking / answer (works with or without enable_thinking)
26 + result = parse_thinking_output(text)
27 + print(result["thinking"]) # chain-of-thought or None
28 + print(result["answer"]) # final answer
29 +
30 +Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users
31 +do not need a transformers dependency for output parsing.
32 +"""
33 +
34 +# ---- Thinking Mode Utility ----
35 +
36 +# Thinking delimiter tokens as they appear in decoded text.
37 +# Gemma4 uses <|channel> (start) and <channel|> (end) as thinking delimiters.
38 +_THINKING_START_TAG = "<|channel>"
39 +_THINKING_END_TAG = "<channel|>"
40 +
41 +# Sentinel tokens that may appear in decoded output.
42 +_TURN_END_TAG = "<turn|>"
43 +
44 +
45 +def parse_thinking_output(text: str) -> dict[str, str | None]:
46 + """Parse decoded Gemma4 model output.
47 +
48 + Use this on **all** Gemma4 output regardless of whether thinking mode
49 + was enabled. It handles three cases:
50 +
51 + 1. **Thinking enabled, tags present** — splits on ``<|channel>``/
52 + ``<channel|>`` to separate chain-of-thought from the answer and
53 + strips the ``thought\\n`` role label.
54 + 2. **Thinking disabled, spurious label** — strips the bare
55 + ``thought\\n`` prefix that some Gemma4 models emit even
56 + without thinking mode.
57 + 3. **Clean output** — returns the text unchanged.
58 +
59 + The answer text is always cleaned of trailing sentinel tokens
60 + (``<turn|>``, ``<eos>``, etc.).
61 +
62 + Args:
63 + text: Decoded model output text (from ``tokenizer.decode(...)``).
64 +
65 + Returns:
66 + A dict with keys:
67 + - ``"thinking"``: The chain-of-thought text, or ``None`` if no
68 + thinking delimiters were found.
69 + - ``"answer"``: The final answer text.
70 +
71 + Example::
72 +
73 + >>> from vllm.reasoning.gemma4_utils import parse_thinking_output
74 + >>> output_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
75 + >>> result = parse_thinking_output(output_text)
76 + >>> print(result["thinking"]) # chain-of-thought reasoning or None
77 + >>> print(result["answer"]) # final answer
78 + """
79 + if _THINKING_END_TAG in text:
80 + parts = text.split(_THINKING_END_TAG, 1)
81 + thinking_block = parts[0]
82 + answer = _clean_answer(parts[1])
83 +
84 + # Extract thinking content: strip the start tag if present
85 + if _THINKING_START_TAG in thinking_block:
86 + thinking = thinking_block.split(_THINKING_START_TAG, 1)[1]
87 + else:
88 + thinking = thinking_block
89 +
90 + # Strip the "thought\n" channel role label the model emits inside
91 + # <|channel>thought\n...<channel|> (analogous to "user\n" in
92 + # <|turn>user\n...<turn|>).
93 + thinking = _strip_thought_label(thinking.strip())
94 + thinking = thinking.strip()
95 +
96 + return {"thinking": thinking, "answer": answer}
97 +
98 + # No thinking delimiters found.
99 + # Strip spurious "thought\n" role label that some Gemma4 models sometimes
100 + # emit even without thinking mode enabled, then clean trailing tokens.
101 + answer = _strip_thought_label(text)
102 + answer = _clean_answer(answer)
103 + return {"thinking": None, "answer": answer}
104 +
105 +
106 +def _strip_thought_label(text: str) -> str:
107 + """Strip the spurious ``thought\\n`` label from the start of text.
108 +
109 + Only strips when ``thought`` appears as the very first word followed by
110 + a newline — preserving the word ``thought`` in any other context.
111 + """
112 + if text.startswith("thought\n"):
113 + return text[len("thought\n") :]
114 + return text
115 +
116 +
117 +def _clean_answer(text: str) -> str:
118 + """Clean trailing sentinel tokens from the answer text.
119 +
120 + Strips ``<turn|>``, ``<eos>``, and surrounding whitespace that the
121 + model appends at the end of its response.
122 + """
123 + text = text.strip()
124 + # Strip trailing <turn|> (Gemma4 turn-end marker)
125 + if text.endswith(_TURN_END_TAG):
126 + text = text[: -len(_TURN_END_TAG)].rstrip()
127 + # Strip trailing <eos> if present
128 + if text.endswith("<eos>"):
129 + text = text[:-5].rstrip()
130 + return text
vllm/tool_parsers/__init__.py
+4
@@ -154,6 +154,10 @@ _TOOL_PARSERS_TO_REGISTER = {
154 "functiongemma_tool_parser",
155 "FunctionGemmaToolParser",
156 ),
157 + "gemma4": (
158 + "gemma4_tool_parser",
159 + "Gemma4ToolParser",
160 + ),
161 }
162
163
vllm/tool_parsers/gemma4_tool_parser.py new
+724
@@ -0,0 +1,724 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tool call parser for Google Gemma4 models.
5 +
6 +Gemma4 uses a custom serialization format (not JSON) for tool calls::
7 +
8 + <|tool_call>call:func_name{key:<|"|>value<|"|>,num:42}<tool_call|>
9 +
10 +Strings are delimited by ``<|"|>`` (token 52), keys are unquoted, and
11 +multiple tool calls are concatenated without separators.
12 +
13 +Used when ``--enable-auto-tool-choice --tool-call-parser gemma4`` are set.
14 +
15 +For offline inference tool call parsing (direct ``tokenizer.decode()`` output),
16 +see ``vllm.tool_parsers.gemma4_utils.parse_tool_calls``.
17 +"""
18 +
19 +import json
20 +from collections.abc import Sequence
21 +
22 +import regex as re
23 +
24 +from vllm.entrypoints.chat_utils import make_tool_call_id
25 +from vllm.entrypoints.openai.chat_completion.protocol import (
26 + ChatCompletionRequest,
27 +)
28 +from vllm.entrypoints.openai.engine.protocol import (
29 + DeltaFunctionCall,
30 + DeltaMessage,
31 + DeltaToolCall,
32 + ExtractedToolCallInformation,
33 + FunctionCall,
34 + ToolCall,
35 +)
36 +from vllm.entrypoints.openai.responses.protocol import (
37 + ResponsesRequest,
38 +)
39 +from vllm.logger import init_logger
40 +from vllm.tokenizers import TokenizerLike
41 +from vllm.tool_parsers.abstract_tool_parser import ToolParser
42 +from vllm.tool_parsers.utils import find_common_prefix
43 +
44 +logger = init_logger(__name__)
45 +
46 +# Gemma4 special tokens for tool calls
47 +TOOL_CALL_START = "<|tool_call>"
48 +TOOL_CALL_END = "<tool_call|>"
49 +STRING_DELIM = '<|"|>'
50 +
51 +
52 +# ---------------------------------------------------------------------------
53 +# Gemma4 argument parser (used by both streaming and non-streaming paths)
54 +# ---------------------------------------------------------------------------
55 +
56 +
57 +def _parse_gemma4_value(value_str: str) -> object:
58 + """Parse a single Gemma4 value (after key:) into a Python object."""
59 + value_str = value_str.strip()
60 + if not value_str:
61 + return value_str
62 +
63 + # Boolean
64 + if value_str == "true":
65 + return True
66 + if value_str == "false":
67 + return False
68 +
69 + # Number (int or float)
70 + try:
71 + if "." in value_str:
72 + return float(value_str)
73 + return int(value_str)
74 + except ValueError:
75 + pass
76 +
77 + # Bare string (no <|"|> delimiters — shouldn't happen but be safe)
78 + return value_str
79 +
80 +
81 +def _parse_gemma4_args(args_str: str) -> dict:
82 + """Parse Gemma4's custom key:value format into a Python dict.
83 +
84 + Format examples::
85 +
86 + location:<|"|>Tokyo<|"|>
87 + location:<|"|>San Francisco<|"|>,unit:<|"|>celsius<|"|>
88 + count:42,flag:true
89 + nested:{inner_key:<|"|>val<|"|>}
90 + items:[<|"|>a<|"|>,<|"|>b<|"|>]
91 +
92 + Returns a dict ready for ``json.dumps()``.
93 + """
94 + if not args_str or not args_str.strip():
95 + return {}
96 +
97 + result: dict = {}
98 + i = 0
99 + n = len(args_str)
100 +
101 + while i < n:
102 + # Skip whitespace and commas
103 + while i < n and args_str[i] in (" ", ",", "\n", "\t"):
104 + i += 1
105 + if i >= n:
106 + break
107 +
108 + # Parse key (unquoted, ends at ':')
109 + key_start = i
110 + while i < n and args_str[i] != ":":
111 + i += 1
112 + if i >= n:
113 + break
114 + key = args_str[key_start:i].strip()
115 + i += 1 # skip ':'
116 +
117 + # Parse value
118 + if i >= n:
119 + result[key] = ""
120 + break
121 +
122 + # Skip whitespace after ':'
123 + while i < n and args_str[i] in (" ", "\n", "\t"):
124 + i += 1
125 + if i >= n:
126 + result[key] = ""
127 + break
128 +
129 + # String value: <|"|>...<|"|>
130 + if args_str[i:].startswith(STRING_DELIM):
131 + i += len(STRING_DELIM)
132 + val_start = i
133 + end_pos = args_str.find(STRING_DELIM, i)
134 + if end_pos == -1:
135 + # Unterminated string — take rest
136 + result[key] = args_str[val_start:]
137 + break
138 + result[key] = args_str[val_start:end_pos]
139 + i = end_pos + len(STRING_DELIM)
140 +
141 + # Nested object: {...}
142 + elif args_str[i] == "{":
143 + depth = 1
144 + obj_start = i + 1
145 + i += 1
146 + while i < n and depth > 0:
147 + if args_str[i:].startswith(STRING_DELIM):
148 + # Skip over string contents to avoid counting { inside strings
149 + i += len(STRING_DELIM)
150 + next_delim = args_str.find(STRING_DELIM, i)
151 + i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
152 + continue
153 + if args_str[i] == "{":
154 + depth += 1
155 + elif args_str[i] == "}":
156 + depth -= 1
157 + i += 1
158 + result[key] = _parse_gemma4_args(args_str[obj_start : i - 1])
159 +
160 + # Array: [...]
161 + elif args_str[i] == "[":
162 + depth = 1
163 + arr_start = i + 1
164 + i += 1
165 + while i < n and depth > 0:
166 + if args_str[i:].startswith(STRING_DELIM):
167 + i += len(STRING_DELIM)
168 + next_delim = args_str.find(STRING_DELIM, i)
169 + i = n if next_delim == -1 else next_delim + len(STRING_DELIM)
170 + continue
171 + if args_str[i] == "[":
172 + depth += 1
173 + elif args_str[i] == "]":
174 + depth -= 1
175 + i += 1
176 + arr_content = args_str[arr_start : i - 1]
177 + result[key] = _parse_gemma4_array(arr_content)
178 +
179 + # Bare value (number, boolean, etc.)
180 + else:
181 + val_start = i
182 + while i < n and args_str[i] not in (",", "}", "]"):
183 + i += 1
184 + result[key] = _parse_gemma4_value(args_str[val_start:i])
185 +
186 + return result
187 +
188 +
189 +def _parse_gemma4_array(arr_str: str) -> list:
190 + """Parse a Gemma4 array content string into a Python list."""
191 + items: list = []
192 + i = 0
193 + n = len(arr_str)
194 +
195 + while i < n:
196 + while i < n and arr_str[i] in (" ", ",", "\n", "\t"):
197 + i += 1
198 + if i >= n:
199 + break
200 +
201 + # String element
202 + if arr_str[i:].startswith(STRING_DELIM):
203 + i += len(STRING_DELIM)
204 + end_pos = arr_str.find(STRING_DELIM, i)
205 + if end_pos == -1:
206 + items.append(arr_str[i:])
207 + break
208 + items.append(arr_str[i:end_pos])
209 + i = end_pos + len(STRING_DELIM)
210 +
211 + # Nested object
212 + elif arr_str[i] == "{":
213 + depth = 1
214 + obj_start = i + 1
215 + i += 1
216 + while i < n and depth > 0:
217 + if arr_str[i:].startswith(STRING_DELIM):
218 + i += len(STRING_DELIM)
219 + nd = arr_str.find(STRING_DELIM, i)
220 + i = nd + len(STRING_DELIM) if nd != -1 else n
221 + continue
222 + if arr_str[i] == "{":
223 + depth += 1
224 + elif arr_str[i] == "}":
225 + depth -= 1
226 + i += 1
227 + items.append(_parse_gemma4_args(arr_str[obj_start : i - 1]))
228 +
229 + # Nested array
230 + elif arr_str[i] == "[":
231 + depth = 1
232 + sub_start = i + 1
233 + i += 1
234 + while i < n and depth > 0:
235 + if arr_str[i] == "[":
236 + depth += 1
237 + elif arr_str[i] == "]":
238 + depth -= 1
239 + i += 1
240 + items.append(_parse_gemma4_array(arr_str[sub_start : i - 1]))
241 +
242 + # Bare value
243 + else:
244 + val_start = i
245 + while i < n and arr_str[i] not in (",", "]"):
246 + i += 1
247 + items.append(_parse_gemma4_value(arr_str[val_start:i]))
248 +
249 + return items
250 +
251 +
252 +# ---------------------------------------------------------------------------
253 +# Parser
254 +# ---------------------------------------------------------------------------
255 +
256 +
257 +class Gemma4ToolParser(ToolParser):
258 + """
259 + Tool call parser for Google Gemma4 models.
260 +
261 + Handles the Gemma4 function call format::
262 +
263 + <|tool_call>call:func_name{key:<|"|>value<|"|>}<tool_call|>
264 +
265 + Used when ``--enable-auto-tool-choice --tool-call-parser gemma4``
266 + are set.
267 +
268 + Streaming strategy: **accumulate-then-parse-then-diff**
269 +
270 + Instead of trying to convert Gemma4's custom format to JSON
271 + token-by-token (which fails because Gemma4 uses bare keys, custom
272 + delimiters, and structural braces that differ from JSON), this parser:
273 +
274 + 1. Accumulates the raw Gemma4 argument string during streaming
275 + 2. Parses it with ``_parse_gemma4_args()`` into a Python dict
276 + 3. Converts to JSON with ``json.dumps()``
277 + 4. Diffs against the previously-streamed JSON string
278 + 5. Emits only the new JSON fragment as the delta
279 +
280 + This follows the same pattern used by FunctionGemma, Hermes, and Llama
281 + tool parsers.
282 + """
283 +
284 + def __init__(self, tokenizer: TokenizerLike):
285 + super().__init__(tokenizer)
286 +
287 + if not self.model_tokenizer:
288 + raise ValueError(
289 + "The model tokenizer must be passed to the ToolParser "
290 + "constructor during construction."
291 + )
292 +
293 + # Token strings
294 + self.tool_call_start_token = TOOL_CALL_START
295 + self.tool_call_end_token = TOOL_CALL_END
296 +
297 + # Token IDs
298 + self.tool_call_start_token_id = self.vocab.get(TOOL_CALL_START)
299 + self.tool_call_end_token_id = self.vocab.get(TOOL_CALL_END)
300 +
301 + if self.tool_call_start_token_id is None:
302 + raise RuntimeError(
303 + "Gemma4 ToolParser could not locate the tool call start "
304 + f"token '{TOOL_CALL_START}' in the tokenizer!"
305 + )
306 +
307 + # Regex for non-streaming: extract complete tool calls.
308 + # Supports function names with letters, digits, underscores,
309 + # hyphens, and dots (e.g. "get-weather", "module.func").
310 + self.tool_call_regex = re.compile(
311 + r"<\|tool_call>call:([\w\-\.]+)\{(.*?)\}<tool_call\|>",
312 + re.DOTALL,
313 + )
314 +
315 + # Streaming state — reset per-request via _reset_streaming_state()
316 + self._reset_streaming_state()
317 +
318 + # Delta buffer for handling multi-token special sequences
319 + self.buffered_delta_text = ""
320 +
321 + def _reset_streaming_state(self) -> None:
322 + """Reset all streaming state for a new request."""
323 + self.current_tool_id = -1
324 + self.current_tool_name_sent = False
325 + self.prev_tool_call_arr: list[dict] = []
326 + self.streamed_args_for_tool: list[str] = []
327 +
328 + def adjust_request(
329 + self, request: ChatCompletionRequest | ResponsesRequest
330 + ) -> ChatCompletionRequest | ResponsesRequest:
331 + request = super().adjust_request(request)
332 + if (
333 + isinstance(request, ChatCompletionRequest)
334 + and request.tools
335 + and request.tool_choice != "none"
336 + ):
337 + # Don't skip special tokens — <|tool_call> etc. are needed
338 + request.skip_special_tokens = False
339 + return request
340 +
341 + # ------------------------------------------------------------------
342 + # Delta buffering for multi-token special sequences
343 + # ------------------------------------------------------------------
344 +
345 + def _buffer_delta_text(self, delta_text: str) -> str:
346 + """Buffer incoming delta text to handle multi-token special sequences.
347 +
348 + Accumulates partial tokens that could be the start of
349 + ``<|tool_call>`` or ``<tool_call|>`` and only flushes them
350 + when the complete sequence is recognized or the sequence breaks.
351 +
352 + This prevents partial special tokens (e.g., ``<|tool``) from being
353 + emitted prematurely as content text.
354 + """
355 + combined = self.buffered_delta_text + delta_text
356 +
357 + # Check if combined ends with a complete special token
358 + if combined.endswith(TOOL_CALL_START) or combined.endswith(TOOL_CALL_END):
359 + self.buffered_delta_text = ""
360 + return combined
361 +
362 + # Check if combined ends with a partial prefix of a special token
363 + for tag in [TOOL_CALL_START, TOOL_CALL_END]:
364 + for i in range(1, len(tag)):
365 + if combined.endswith(tag[:i]):
366 + self.buffered_delta_text = combined[-i:]
367 + return combined[:-i]
368 +
369 + # No partial match — flush everything
370 + self.buffered_delta_text = ""
371 + return combined
372 +
373 + # ------------------------------------------------------------------
374 + # Non-streaming extraction
375 + # ------------------------------------------------------------------
376 +
377 + def extract_tool_calls(
378 + self,
379 + model_output: str,
380 + request: ChatCompletionRequest,
381 + ) -> ExtractedToolCallInformation:
382 + if self.tool_call_start_token not in model_output:
383 + return ExtractedToolCallInformation(
384 + tools_called=False, tool_calls=[], content=model_output
385 + )
386 +
387 + try:
388 + matches = self.tool_call_regex.findall(model_output)
389 + if not matches:
390 + return ExtractedToolCallInformation(
391 + tools_called=False, tool_calls=[], content=model_output
392 + )
393 +
394 + tool_calls: list[ToolCall] = []
395 + for func_name, args_str in matches:
396 + arguments = _parse_gemma4_args(args_str)
397 + tool_calls.append(
398 + ToolCall(
399 + type="function",
400 + function=FunctionCall(
401 + name=func_name,
402 + arguments=json.dumps(arguments, ensure_ascii=False),
403 + ),
404 + )
405 + )
406 +
407 + # Content = text before first tool call (if any)
408 + content_end = model_output.find(self.tool_call_start_token)
409 + content = model_output[:content_end].strip() if content_end > 0 else None
410 +
411 + return ExtractedToolCallInformation(
412 + tools_called=True,
413 + tool_calls=tool_calls,
414 + content=content if content else None,
415 + )
416 +
417 + except Exception:
418 + logger.exception("Error extracting tool calls from Gemma4 response")
419 + return ExtractedToolCallInformation(
420 + tools_called=False, tool_calls=[], content=model_output
421 + )
422 +
423 + # ------------------------------------------------------------------
424 + # Streaming extraction — accumulate-then-parse-then-diff
425 + # ------------------------------------------------------------------
426 +
427 + def extract_tool_calls_streaming(
428 + self,
429 + previous_text: str,
430 + current_text: str,
431 + delta_text: str,
432 + previous_token_ids: Sequence[int],
433 + current_token_ids: Sequence[int],
434 + delta_token_ids: Sequence[int],
435 + request: ChatCompletionRequest,
436 + ) -> DeltaMessage | None:
437 + # Buffer delta text to handle multi-token special sequences
438 + delta_text = self._buffer_delta_text(delta_text)
439 + # Reconstruct current_text after buffering to stay in sync
440 + current_text = previous_text + delta_text
441 +
442 + # If no tool call token seen yet, emit as content
443 + if self.tool_call_start_token not in current_text:
444 + if delta_text:
445 + return DeltaMessage(content=delta_text)
446 + return None
447 +
448 + try:
449 + return self._extract_streaming(
450 + previous_text=previous_text,
451 + current_text=current_text,
452 + delta_text=delta_text,
453 + )
454 + except Exception:
455 + logger.exception("Error in Gemma4 streaming tool call extraction")
456 + return None
457 +
458 + def _extract_streaming(
459 + self,
460 + previous_text: str,
461 + current_text: str,
462 + delta_text: str,
463 + ) -> DeltaMessage | None:
464 + """Tag-counting streaming parser.
465 +
466 + Uses the proven approach from FunctionGemma/Hermes: count start/end
467 + tags in previous vs current text to determine phase, then
468 + accumulate-parse-diff for arguments.
469 +
470 + Format: ``<|tool_call>call:name{args}<tool_call|>``
471 + """
472 + start_count = current_text.count(self.tool_call_start_token)
473 + end_count = current_text.count(self.tool_call_end_token)
474 + prev_start_count = previous_text.count(self.tool_call_start_token)
475 + prev_end_count = previous_text.count(self.tool_call_end_token)
476 +
477 + # Case 1: Not inside any tool call — emit as content
478 + if (
479 + start_count == end_count
480 + and prev_end_count == end_count
481 + and self.tool_call_end_token not in delta_text
482 + ):
483 + if delta_text:
484 + return DeltaMessage(content=delta_text)
485 + return None
486 +
487 + # Case 2: Starting a new tool call
488 + if start_count > prev_start_count and start_count > end_count:
489 + self.current_tool_id += 1
490 + self.current_tool_name_sent = False
491 + self.streamed_args_for_tool.append("")
492 + self.prev_tool_call_arr.append({})
493 + logger.debug("Starting new tool call %d", self.current_tool_id)
494 + # Don't return yet — fall through to try parsing if there's
495 + # content after <|tool_call> in this same delta
496 + # (but usually it's just the token itself, so return None)
497 + if len(delta_text) <= len(self.tool_call_start_token):
498 + return None
499 +
500 + # Case 3: Tool call just ended
501 + if end_count > prev_end_count:
502 + return self._handle_tool_call_end(current_text)
503 +
504 + # Case 4: In the middle of a tool call — parse partial content
505 + if start_count > end_count:
506 + return self._handle_tool_call_middle(current_text)
507 +
508 + # Default: generate text outside tool calls
509 + if delta_text:
510 + text = delta_text.replace(self.tool_call_start_token, "")
511 + text = text.replace(self.tool_call_end_token, "")
512 + if text:
513 + return DeltaMessage(content=text)
514 + return None
515 +
516 + def _extract_partial_call(self, current_text: str) -> tuple[str | None, str]:
517 + """Extract function name and raw argument string from partial text.
518 +
519 + Returns (func_name, raw_args_str) or (None, "") if not parseable yet.
520 + """
521 + # Get the text after the last <|tool_call> token
522 + last_start = current_text.rfind(self.tool_call_start_token)
523 + if last_start == -1:
524 + return None, ""
525 +
526 + partial_call = current_text[last_start + len(self.tool_call_start_token) :]
527 +
528 + # Strip end token if present
529 + if self.tool_call_end_token in partial_call:
530 + partial_call = partial_call.split(self.tool_call_end_token)[0]
531 +
532 + # Expect "call:name{args...}" or "call:name{args...}"
533 + if not partial_call.startswith("call:"):
534 + return None, ""
535 +
536 + func_part = partial_call[5:] # skip "call:"
537 +
538 + if "{" not in func_part:
539 + # Still accumulating function name, not ready yet
540 + return None, ""
541 +
542 + func_name, _, args_part = func_part.partition("{")
543 + func_name = func_name.strip()
544 +
545 + # Strip trailing '}' if present (Gemma4 structural brace)
546 + if args_part.endswith("}"):
547 + args_part = args_part[:-1]
548 +
549 + return func_name, args_part
550 +
551 + def _handle_tool_call_middle(self, current_text: str) -> DeltaMessage | None:
552 + """Handle streaming when we're inside an active tool call.
553 +
554 + Accumulates the raw Gemma4 arguments, parses them into JSON, and
555 + diffs against the previously-streamed JSON to emit only the new
556 + fragment.
557 + """
558 + func_name, args_part = self._extract_partial_call(current_text)
559 +
560 + if func_name is None:
561 + return None
562 +
563 + # Step 1: Send function name (once)
564 + if not self.current_tool_name_sent and func_name:
565 + self.current_tool_name_sent = True
566 + self.prev_tool_call_arr[self.current_tool_id] = {
567 + "name": func_name,
568 + "arguments": {},
569 + }
570 + return DeltaMessage(
571 + tool_calls=[
572 + DeltaToolCall(
573 + index=self.current_tool_id,
574 + type="function",
575 + id=make_tool_call_id(),
576 + function=DeltaFunctionCall(
577 + name=func_name,
578 + arguments="",
579 + ).model_dump(exclude_none=True),
580 + )
581 + ]
582 + )
583 +
584 + # Step 2: Parse and diff arguments
585 + if self.current_tool_name_sent and args_part:
586 + return self._emit_argument_diff(args_part)
587 +
588 + return None
589 +
590 + def _handle_tool_call_end(self, current_text: str) -> DeltaMessage | None:
591 + """Handle streaming when a tool call has just completed.
592 +
593 + Performs a final parse of the complete tool call and flushes
594 + any remaining un-streamed argument fragments.
595 + """
596 + if self.current_tool_id < 0 or self.current_tool_id >= len(
597 + self.prev_tool_call_arr
598 + ):
599 + logger.debug(
600 + "Tool call end detected but no active tool call (current_tool_id=%d)",
601 + self.current_tool_id,
602 + )
603 + return None
604 +
605 + # Parse the complete tool call using regex for accuracy
606 + all_matches = self.tool_call_regex.findall(current_text)
607 + if self.current_tool_id < len(all_matches):
608 + _, args_str = all_matches[self.current_tool_id]
609 + final_args = _parse_gemma4_args(args_str)
610 + final_args_json = json.dumps(final_args, ensure_ascii=False)
611 +
612 + prev_streamed = self.streamed_args_for_tool[self.current_tool_id]
613 + if len(final_args_json) > len(prev_streamed):
614 + diff = final_args_json[len(prev_streamed) :]
615 + self.streamed_args_for_tool[self.current_tool_id] = final_args_json
616 + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = final_args
617 +
618 + return DeltaMessage(
619 + tool_calls=[
620 + DeltaToolCall(
621 + index=self.current_tool_id,
622 + function=DeltaFunctionCall(arguments=diff).model_dump(
623 + exclude_none=True
624 + ),
625 + )
626 + ]
627 + )
628 +
629 + return None
630 +
631 + def _emit_argument_diff(self, raw_args_str: str) -> DeltaMessage | None:
632 + """Parse raw Gemma4 arguments, convert to JSON, diff, and emit.
633 +
634 + This is the core of the accumulate-then-parse-then-diff strategy:
635 + 1. Parse ``raw_args_str`` with ``_parse_gemma4_args()``
636 + 2. Convert to JSON string with ``json.dumps()``
637 + 3. Withhold trailing closing characters (``"}``) that may move
638 + as more tokens arrive
639 + 4. Diff against previously streamed JSON and emit only new chars
640 +
641 + **Why withholding is necessary:**
642 +
643 + Gemma4's custom format produces *structurally incomplete* JSON
644 + during streaming. For example, when ``<|"|>Paris`` arrives
645 + without a closing delimiter, ``_parse_gemma4_args`` treats it
646 + as a complete value and produces ``{"location": "Paris"}``. But
647 + when ``, France<|"|>`` arrives next, the JSON becomes
648 + ``{"location": "Paris, France"}``. If we had sent the closing
649 + ``"}`` from the first parse, the concatenated client output
650 + would be ``{"location": "Paris"}France"}``, which is garbage.
651 +
652 + The solution: **never send trailing closing chars during
653 + streaming**. They get flushed by ``_handle_tool_call_end()``
654 + when the ``<tool_call|>`` end marker arrives.
655 +
656 + Args:
657 + raw_args_str: The raw Gemma4 argument text accumulated so far
658 + (without the surrounding ``{`` ``}``).
659 +
660 + Returns:
661 + DeltaMessage with the argument diff, or None if no new content.
662 + """
663 + try:
664 + current_args = _parse_gemma4_args(raw_args_str)
665 + except Exception:
666 + logger.debug(
667 + "Could not parse partial Gemma4 args yet: %s",
668 + raw_args_str[:100],
669 + )
670 + return None
671 +
672 + if not current_args:
673 + return None
674 +
675 + current_args_json = json.dumps(current_args, ensure_ascii=False)
676 +
677 + # Withhold trailing closing characters that may shift as more
678 + # tokens arrive. Strip trailing '}', '"', and ']' sequences
679 + # to get the "safe prefix".
680 + safe_json = current_args_json
681 + while safe_json and safe_json[-1] in ("}", '"', "]"):
682 + safe_json = safe_json[:-1]
683 +
684 + prev_streamed = self.streamed_args_for_tool[self.current_tool_id]
685 +
686 + if not safe_json or safe_json == prev_streamed:
687 + return None
688 +
689 + # Use find_common_prefix to handle cases where the value changed
690 + # structurally (e.g., a string grew).
691 + if prev_streamed:
692 + prefix = find_common_prefix(prev_streamed, safe_json)
693 + sent_len = len(prev_streamed)
694 + prefix_len = len(prefix)
695 +
696 + if prefix_len < sent_len:
697 + # Structure changed — we sent too much. Truncate our
698 + # tracking to the common prefix and wait for the final
699 + # flush in _handle_tool_call_end.
700 + self.streamed_args_for_tool[self.current_tool_id] = prefix
701 + return None
702 +
703 + # Stream the new stable portion
704 + diff = safe_json[sent_len:]
705 + else:
706 + # First emission
707 + diff = safe_json
708 +
709 + if diff:
710 + self.streamed_args_for_tool[self.current_tool_id] = safe_json
711 + self.prev_tool_call_arr[self.current_tool_id]["arguments"] = current_args
712 +
713 + return DeltaMessage(
714 + tool_calls=[
715 + DeltaToolCall(
716 + index=self.current_tool_id,
717 + function=DeltaFunctionCall(arguments=diff).model_dump(
718 + exclude_none=True
719 + ),
720 + )
721 + ]
722 + )
723 +
724 + return None
vllm/tool_parsers/gemma4_utils.py new
+183
@@ -0,0 +1,183 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.
4 +
5 +"""Gemma4 tool call parsing utilities for offline inference.
6 +
7 +Standalone functions that parse decoded model text to extract tool calls
8 +from Gemma4 models. These are pure-Python utilities with zero heavy
9 +dependencies — they work on raw decoded strings from any inference
10 +backend (vLLM, HuggingFace, TGI, etc.).
11 +
12 +For the OpenAI-compatible API server tool parser (streaming +
13 +non-streaming), see ``vllm.tool_parsers.gemma4_tool_parser``.
14 +For thinking/reasoning output parsing, see
15 +``vllm.reasoning.gemma4_utils``.
16 +
17 +Usage with vLLM offline inference::
18 +
19 + from vllm import LLM, SamplingParams
20 + from vllm.tool_parsers.gemma4_utils import (
21 + parse_tool_calls,
22 + has_tool_response_tag,
23 + )
24 +
25 + llm = LLM(model="google/gemma-4-it")
26 + outputs = llm.generate(prompt, SamplingParams(...))
27 + text = tokenizer.decode(outputs[0].outputs[0].token_ids, skip_special_tokens=False)
28 +
29 + # Extract tool calls
30 + tool_calls = parse_tool_calls(text)
31 + for tc in tool_calls:
32 + print(f"{tc['name']}({tc['arguments']})")
33 +
34 +Ported from ``transformers.models.gemma4.utils_gemma4`` so that vLLM users
35 +do not need a transformers dependency for output parsing.
36 +"""
37 +
38 +import json
39 +
40 +import regex as re
41 +
42 +# Tool call delimiter tokens as they appear in decoded text.
43 +# Standard format: <|tool_call>call:name{args}<tool_call|>
44 +_TOOL_CALL_START_TAG = "<|tool_call>"
45 +_TOOL_CALL_END_TAG = "<tool_call|>"
46 +_TOOL_RESPONSE_START_TAG = "<|tool_response>"
47 +
48 +# Gemma4 escape token as it appears in decoded text.
49 +_ESCAPE_TOKEN = '<|"|>'
50 +
51 +
52 +def _parse_tool_arguments(args_str: str) -> dict[str, str]:
53 + """Parse tool call arguments from the Gemma4 compact format.
54 +
55 + Handles the ``key:<|"|>value<|"|>`` format used by Gemma4, with fallback
56 + to heuristic key-value extraction. Also tolerates the slightly different
57 + ``key: "value"`` format (space + plain quotes) that some chat templates
58 + produce.
59 +
60 + Args:
61 + args_str: Raw argument string from inside ``call:name{...}``.
62 +
63 + Returns:
64 + Dictionary of argument name → value.
65 + """
66 + if not args_str or not args_str.strip():
67 + return {}
68 +
69 + # Replace Gemma4 escape tokens with standard quotes.
70 + cleaned = args_str.replace(_ESCAPE_TOKEN, '"')
71 +
72 + # Try JSON parsing first (handles nested values, arrays, etc.).
73 + try:
74 + parsed = json.loads("{" + cleaned + "}")
75 + # Ensure all values are strings for consistency.
76 + return {k: str(v) if not isinstance(v, str) else v for k, v in parsed.items()}
77 + except (json.JSONDecodeError, ValueError):
78 + pass
79 +
80 + # Fallback: extract key:"value" pairs (allow optional space after colon).
81 + arguments = {}
82 + for key, value in re.findall(r'(\w+):\s*"([^"]*)"', cleaned):
83 + arguments[key] = value
84 +
85 + if not arguments:
86 + # Last resort: extract key:value pairs (unquoted).
87 + for key, value in re.findall(r"(\w+):\s*([^,}]+)", args_str):
88 + arguments[key] = value.strip().strip('"').replace(_ESCAPE_TOKEN, "")
89 +
90 + return arguments
91 +
92 +
93 +def parse_tool_calls(text: str, *, strict: bool = False) -> list[dict]:
94 + """Parse tool calls from decoded Gemma4 model output.
95 +
96 + Uses a tiered parsing strategy to handle known output variations in
97 + Gemma4 models, which may emit
98 + non-standard tool call formats.
99 +
100 + Parsing tiers:
101 + 1. **Standard**: ``<|tool_call>call:name{args}<tool_call|>``
102 + (special token IDs 48/49 in decoded text)
103 + 2. **Fallback** (when ``strict=False``): bare ``call:name{args}``
104 + patterns, including ``<call>name{args}`` (fragmented tokens from
105 + multimodal inputs)
106 +
107 + Args:
108 + text: Decoded model output text (from ``tokenizer.decode(...,
109 + skip_special_tokens=False)``).
110 + strict: If ``True``, only match the standard ``<|tool_call>`` format.
111 + If ``False`` (default), also try fallback patterns for
112 + known Gemma4 output variations.
113 +
114 + Returns:
115 + A list of dicts, each with keys:
116 + - ``"name"``: The tool function name (e.g. ``"get_weather"``).
117 + - ``"arguments"``: A dict of argument name → value.
118 +
119 + Example::
120 +
121 + >>> from vllm.tool_parsers.gemma4_utils import parse_tool_calls
122 + >>> output = tokenizer.decode(outputs[0], skip_special_tokens=False)
123 + >>> tool_calls = parse_tool_calls(output)
124 + >>> for tc in tool_calls:
125 + ... print(f"Call: {tc['name']}({tc['arguments']})")
126 + """
127 + results = []
128 +
129 + # Tier 1: Standard format with special tokens.
130 + # <|tool_call>call:name{args}<tool_call|>
131 + # Note: Some Gemma4 models emit <turn|> instead of <tool_call|>.
132 + standard_pattern = r"<\|tool_call\>call:(\w+)\{(.*?)\}(?:<tool_call\|>|<turn\|>)"
133 + for match in re.finditer(standard_pattern, text, re.DOTALL):
134 + name, args_str = match.group(1), match.group(2)
135 + results.append(
136 + {
137 + "name": name,
138 + "arguments": _parse_tool_arguments(args_str),
139 + }
140 + )
141 +
142 + if results or strict:
143 + return results
144 +
145 + # Tier 2: Fallback for known Gemma4 output variations.
146 + # Matches: <call>name{args}, call:name{args}, or bare call:name{args}<eos>
147 + fallback_pattern = r"(?:<call>|(?:^|\s)call:)(\w+)\{(.*?)\}"
148 + for match in re.finditer(fallback_pattern, text, re.DOTALL):
149 + name, args_str = match.group(1), match.group(2)
150 + results.append(
151 + {
152 + "name": name,
153 + "arguments": _parse_tool_arguments(args_str),
154 + }
155 + )
156 +
157 + return results
158 +
159 +
160 +def has_tool_response_tag(text: str) -> bool:
161 + """Check if model output properly ends with a tool response tag.
162 +
163 + Some Gemma4 models sometimes emit ``<eos>`` instead of
164 + ``<|tool_response>`` after a tool call. This helper detects
165 + whether the model used the proper termination, so callers can
166 + decide whether to inject ``<|tool_response>`` into the next prompt.
167 +
168 + Args:
169 + text: Decoded model output text.
170 +
171 + Returns:
172 + ``True`` if the output ends with ``<|tool_response>``
173 + (proper behavior), ``False`` otherwise.
174 +
175 + Example::
176 +
177 + >>> from vllm.tool_parsers.gemma4_utils import has_tool_response_tag
178 + >>> if not has_tool_response_tag(model_output):
179 + ... # Model used <eos> instead — inject <|tool_response> manually
180 + ... next_prompt = "<|tool_response>" + tool_result
181 + """
182 + stripped = text.rstrip()
183 + return stripped.endswith(_TOOL_RESPONSE_START_TAG)
vllm/transformers_utils/model_arch_config_convertor.py
+12
@@ -448,6 +448,16 @@ class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
448 return getattr(self.hf_text_config, "num_nextn_predict_layers", 1)
449
450
451 +class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
452 + def get_head_size(self) -> int:
453 + # Gemma4 uses dual head dimensions: head_dim (sliding attention)
454 + # and global_head_dim (full attention). Return the largest so
455 + # that attention backends allocate buffers large enough for both.
456 + head_dim = getattr(self.hf_text_config, "head_dim", 0)
457 + global_head_dim = getattr(self.hf_text_config, "global_head_dim", 0)
458 + return max(head_dim, global_head_dim) or super().get_head_size()
459 +
460 +
461 # hf_config.model_type -> convertor class
462 MODEL_ARCH_CONFIG_CONVERTORS = {
463 "cohere_asr": CohereAsrModelArchConfigConvertor,
@@ -471,4 +481,6 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
481 "ernie_mtp": ErnieMTPModelArchConfigConvertor,
482 "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,
483 "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,
484 + "gemma4": Gemma4ModelArchConfigConvertor,
485 + "gemma4_text": Gemma4ModelArchConfigConvertor,
486 }