@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1f76efaa21

[Model] Add K2-Horizon model support (#55063)

Bowen Tan committed Sep 3, 2026 at 01:17 UTC 1f76efaa2195485b92cb04215aba6fb8f5fe523d
10 files changed +2613
docs/models/supported_models.md
+1
@@ -414,6 +414,7 @@ th {
414 | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
415 | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
416 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
417 +| `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ |
418 | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
419 | `Lfm2ForCausalLM` | LFM2 | `LiquidAI/LFM2-1.2B`, `LiquidAI/LFM2-700M`, `LiquidAI/LFM2-350M`, etc. | ✅︎ | ✅︎ |
420 | `Lfm2MoeForCausalLM` | LFM2MoE | `LiquidAI/LFM2-8B-A1B-preview`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+5
@@ -364,6 +364,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
364 "random": "ai21labs/Jamba-tiny-random",
365 },
366 ),
367 + "K2HorizonForCausalLM": _HfExamplesInfo(
368 + "IFM/K2-Horizon-36B",
369 + trust_remote_code=True,
370 + is_available_online=False,
371 + ),
372 "KimiLinearForCausalLM": _HfExamplesInfo(
373 "moonshotai/Kimi-Linear-48B-A3B-Instruct", trust_remote_code=True
374 ),
tests/reasoning/test_k2_horizon_reasoning_parser.py new
+214
@@ -0,0 +1,214 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import string
5 +from collections.abc import Sequence
6 +
7 +import pytest
8 +
9 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
10 +from vllm.parser.parser_manager import ParserManager
11 +from vllm.reasoning import ReasoningParserManager
12 +from vllm.reasoning.k2_horizon_reasoning_parser import K2HorizonReasoningParser
13 +
14 +pytestmark = pytest.mark.skip_global_cleanup
15 +
16 +EFFORT_TOKENS = {
17 + "high": ("<ifm|think>", "</ifm|think>"),
18 + "medium": ("<ifm|think_fast>", "</ifm|think_fast>"),
19 + "low": ("<ifm|think_faster>", "</ifm|think_faster>"),
20 +}
21 +TOOL_CALLS_START = "<ifm|tool_calls>"
22 +
23 +
24 +class K2Tokenizer:
25 + def __init__(self) -> None:
26 + special_tokens = [
27 + token for token_pair in EFFORT_TOKENS.values() for token in token_pair
28 + ] + [TOOL_CALLS_START]
29 + tokens = special_tokens + list(string.printable)
30 + self._vocab = {token: index for index, token in enumerate(tokens, start=1)}
31 +
32 + def get_vocab(self) -> dict[str, int]:
33 + return dict(self._vocab)
34 +
35 + def encode(self, text: str) -> list[int]:
36 + tokens: list[int] = []
37 + position = 0
38 + specials = sorted(EFFORT_TOKENS["high"] + (TOOL_CALLS_START,), key=len)
39 + while position < len(text):
40 + for special in specials:
41 + if text.startswith(special, position):
42 + tokens.append(self._vocab[special])
43 + position += len(special)
44 + break
45 + else:
46 + tokens.append(self._vocab[text[position]])
47 + position += 1
48 + return tokens
49 +
50 + def decode(
51 + self,
52 + ids: Sequence[int] | int,
53 + skip_special_tokens: bool = False,
54 + ) -> str:
55 + del skip_special_tokens
56 + if isinstance(ids, int):
57 + ids = [ids]
58 + inverse_vocab = {token_id: token for token, token_id in self._vocab.items()}
59 + return "".join(inverse_vocab[token_id] for token_id in ids)
60 +
61 +
62 +def _request(**kwargs) -> ChatCompletionRequest:
63 + return ChatCompletionRequest(
64 + model="k2-horizon",
65 + messages=[{"role": "user", "content": "test"}],
66 + **kwargs,
67 + )
68 +
69 +
70 +@pytest.mark.parametrize("effort", ["high", "medium", "low"])
71 +def test_reasoning_effort_selects_tokens(effort: str):
72 + parser = K2HorizonReasoningParser(
73 + K2Tokenizer(),
74 + chat_template_kwargs={"reasoning_effort": effort},
75 + )
76 +
77 + assert (parser.start_token, parser.end_token) == EFFORT_TOKENS[effort]
78 +
79 +
80 +def test_reasoning_effort_precedence_and_default():
81 + tokenizer = K2Tokenizer()
82 + request = _request(
83 + reasoning_effort="medium",
84 + chat_template_kwargs={"reasoning_effort": "low"},
85 + )
86 + resolved_kwargs = (
87 + request.build_chat_params(None, "auto")
88 + .with_defaults({"reasoning_effort": "high"})
89 + .chat_template_kwargs
90 + )
91 +
92 + parser = K2HorizonReasoningParser(
93 + tokenizer,
94 + chat_template_kwargs=resolved_kwargs,
95 + )
96 + default_parser = K2HorizonReasoningParser(tokenizer)
97 +
98 + assert parser.start_token == EFFORT_TOKENS["medium"][0]
99 + assert default_parser.start_token == EFFORT_TOKENS["high"][0]
100 +
101 +
102 +@pytest.mark.parametrize("effort", ["none", "unknown", None])
103 +def test_invalid_reasoning_effort_raises(effort: str | None):
104 + with pytest.raises(ValueError, match="Unsupported reasoning_effort"):
105 + K2HorizonReasoningParser(
106 + K2Tokenizer(),
107 + chat_template_kwargs={"reasoning_effort": effort},
108 + )
109 +
110 +
111 +def test_non_streaming_explicit_and_implicit_boundaries():
112 + parser = K2HorizonReasoningParser(K2Tokenizer())
113 + request = _request()
114 +
115 + assert parser.extract_reasoning("<ifm|think>plan</ifm|think>answer", request) == (
116 + "plan",
117 + "answer",
118 + )
119 + assert parser.extract_reasoning(
120 + f"plan{TOOL_CALLS_START}<ifm|tool_call>ping</ifm|tool_call>", request
121 + ) == (
122 + "plan",
123 + f"{TOOL_CALLS_START}<ifm|tool_call>ping</ifm|tool_call>",
124 + )
125 +
126 +
127 +@pytest.mark.parametrize("reasoning", ["plan", " ", "\n", " \n"])
128 +def test_character_split_reasoning_stream(reasoning: str):
129 + parser = K2HorizonReasoningParser(K2Tokenizer())
130 + reasoning_parts: list[str] = []
131 + content_parts: list[str] = []
132 +
133 + for char in f"{reasoning}</ifm|think>answer":
134 + delta = parser.extract_reasoning_streaming(
135 + previous_text="",
136 + current_text="",
137 + delta_text=char,
138 + previous_token_ids=[],
139 + current_token_ids=[],
140 + delta_token_ids=[],
141 + )
142 + if delta is not None:
143 + if delta.reasoning is not None:
144 + reasoning_parts.append(delta.reasoning)
145 + if delta.content is not None:
146 + content_parts.append(delta.content)
147 +
148 + assert "".join(reasoning_parts) == reasoning
149 + assert "".join(content_parts) == "answer"
150 +
151 +
152 +def test_latest_boundary_controls_reasoning_state():
153 + tokenizer = K2Tokenizer()
154 + parser = K2HorizonReasoningParser(tokenizer)
155 + vocab = tokenizer.get_vocab()
156 +
157 + assert parser.is_reasoning_end([vocab["<ifm|think>"], vocab[TOOL_CALLS_START]])
158 + assert not parser.is_reasoning_end([vocab[TOOL_CALLS_START], vocab["<ifm|think>"]])
159 +
160 +
161 +def test_composed_streaming_reasoning_to_tool_handoff():
162 + tokenizer = K2Tokenizer()
163 + tools = [
164 + {
165 + "type": "function",
166 + "function": {
167 + "name": "ping",
168 + "parameters": {"type": "object", "properties": {}},
169 + },
170 + }
171 + ]
172 + request = _request(tools=tools, tool_choice="auto")
173 + parser_cls = ParserManager.get_parser(
174 + tool_parser_name="k2_horizon",
175 + reasoning_parser_name="k2_horizon",
176 + enable_auto_tools=True,
177 + )
178 + assert parser_cls is not None
179 + parser = parser_cls(
180 + tokenizer,
181 + request.tools,
182 + chat_template_kwargs={"reasoning_effort": "high"},
183 + )
184 + output = "plan<ifm|tool_calls><ifm|tool_call>ping</ifm|tool_call></ifm|tool_calls>"
185 + reasoning_parts: list[str] = []
186 + tool_names: list[str] = []
187 +
188 + for index, char in enumerate(output):
189 + delta = parser.parse_delta(
190 + delta_text=char,
191 + delta_token_ids=tokenizer.encode(char),
192 + request=request,
193 + prompt_token_ids=[] if index == 0 else None,
194 + finished=index == len(output) - 1,
195 + )
196 + if delta is not None:
197 + if delta.reasoning is not None:
198 + reasoning_parts.append(delta.reasoning)
199 + tool_names.extend(
200 + tool_call.function.name
201 + for tool_call in delta.tool_calls
202 + if tool_call.function is not None
203 + and tool_call.function.name is not None
204 + )
205 +
206 + assert "".join(reasoning_parts) == "plan"
207 + assert tool_names == ["ping"]
208 +
209 +
210 +def test_reasoning_parser_registered():
211 + assert (
212 + ReasoningParserManager.get_reasoning_parser("k2_horizon")
213 + is K2HorizonReasoningParser
214 + )
tests/tool_parsers/test_k2_horizon_tool_parser.py new
+318
@@ -0,0 +1,318 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from unittest.mock import MagicMock
6 +
7 +import pytest
8 +
9 +from vllm.entrypoints.generate.base.protocol import DeltaMessage
10 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
11 +from vllm.parser.parser_manager import ParserManager
12 +from vllm.tool_parsers import ToolParserManager
13 +from vllm.tool_parsers.k2_horizon_tool_parser import K2HorizonToolParser
14 +
15 +pytestmark = pytest.mark.skip_global_cleanup
16 +
17 +TOOLS = [
18 + {
19 + "type": "function",
20 + "function": {
21 + "name": "lookup",
22 + "parameters": {
23 + "type": "object",
24 + "properties": {
25 + "query": {"type": "string"},
26 + "limit": {"type": "integer"},
27 + "enabled": {"type": "boolean"},
28 + },
29 + },
30 + },
31 + },
32 + {
33 + "type": "function",
34 + "function": {
35 + "name": "ping",
36 + "parameters": {"type": "object", "properties": {}},
37 + },
38 + },
39 +]
40 +
41 +
42 +@pytest.fixture
43 +def tokenizer():
44 + tokenizer = MagicMock()
45 + tokenizer.get_vocab.return_value = {}
46 + return tokenizer
47 +
48 +
49 +def _request(
50 + tool_format: str | None = None,
51 + tool_choice="auto",
52 +) -> ChatCompletionRequest:
53 + chat_template_kwargs = (
54 + {"tool_call_format": tool_format} if tool_format is not None else None
55 + )
56 + return ChatCompletionRequest(
57 + model="k2-horizon",
58 + messages=[{"role": "user", "content": "test"}],
59 + tools=TOOLS,
60 + tool_choice=tool_choice,
61 + chat_template_kwargs=chat_template_kwargs,
62 + )
63 +
64 +
65 +def _group(*calls: str, prefix: str = "", suffix: str = "") -> str:
66 + return prefix + "<ifm|tool_calls>" + "\n".join(calls) + "</ifm|tool_calls>" + suffix
67 +
68 +
69 +def _json_call(name: str, arguments: dict) -> str:
70 + payload = json.dumps({"name": name, "arguments": arguments})
71 + return f"<ifm|tool_call>{payload}</ifm|tool_call>"
72 +
73 +
74 +def _xml_call(
75 + name: str,
76 + arguments: list[tuple[str, str]],
77 + *,
78 + typed: bool = False,
79 +) -> str:
80 + parts = [f"<ifm|tool_call>{name}"]
81 + for key, value in arguments:
82 + parts.append(f"<ifm|arg_key>{key}</ifm|arg_key>")
83 + if typed:
84 + arg_type = "integer" if key == "limit" else "string"
85 + parts.append(f"<ifm|arg_type>{arg_type}</ifm|arg_type>")
86 + parts.append(f"<ifm|arg_value>{value}</ifm|arg_value>")
87 + parts.append("</ifm|tool_call>")
88 + return "".join(parts)
89 +
90 +
91 +def _collect_stream(
92 + parser: K2HorizonToolParser,
93 + request: ChatCompletionRequest,
94 + output: str,
95 +) -> tuple[str, list[tuple[str, dict]]]:
96 + messages: list[DeltaMessage | None] = []
97 + for char in output:
98 + messages.append(
99 + parser.extract_tool_calls_streaming(
100 + previous_text="",
101 + current_text="",
102 + delta_text=char,
103 + previous_token_ids=[],
104 + current_token_ids=[],
105 + delta_token_ids=[],
106 + request=request,
107 + )
108 + )
109 + messages.append(parser.finish_streaming())
110 +
111 + content: list[str] = []
112 + calls: list[tuple[str, dict]] = []
113 + for message in messages:
114 + if message is None:
115 + continue
116 + if message.content is not None:
117 + content.append(message.content)
118 + for tool_call in message.tool_calls:
119 + assert tool_call.function is not None
120 + assert tool_call.function.name is not None
121 + assert tool_call.function.arguments is not None
122 + calls.append(
123 + (
124 + tool_call.function.name,
125 + json.loads(tool_call.function.arguments),
126 + )
127 + )
128 + return "".join(content), calls
129 +
130 +
131 +def test_xml_default_parses_multiple_calls_and_coerces_schema(tokenizer):
132 + parser = K2HorizonToolParser(tokenizer, _request().tools)
133 + output = _group(
134 + _xml_call("lookup", [("query", "123"), ("limit", "3")]),
135 + _xml_call("ping", []),
136 + prefix="Before ",
137 + suffix=" after.",
138 + )
139 +
140 + result = parser.extract_tool_calls(output, _request())
141 +
142 + assert parser.tool_format == "xml"
143 + assert result.tools_called
144 + assert result.content == "Before after."
145 + assert [call.function.name for call in result.tool_calls] == ["lookup", "ping"]
146 + assert json.loads(result.tool_calls[0].function.arguments) == {
147 + "query": "123",
148 + "limit": 3,
149 + }
150 +
151 +
152 +@pytest.mark.parametrize(
153 + ("tool_format", "call", "expected"),
154 + [
155 + ("json", _json_call("lookup", {"limit": "2"}), {"limit": 2}),
156 + (
157 + "xml_typed",
158 + _xml_call("lookup", [("query", "weather"), ("limit", "2")], typed=True),
159 + {"query": "weather", "limit": 2},
160 + ),
161 + ],
162 +)
163 +def test_request_level_format_overrides(tokenizer, tool_format, call, expected):
164 + request = _request(tool_format)
165 + parser = K2HorizonToolParser(tokenizer, request.tools)
166 +
167 + result = parser.extract_tool_calls(_group(call), request)
168 +
169 + assert result.tools_called
170 + assert json.loads(result.tool_calls[0].function.arguments) == expected
171 +
172 +
173 +@pytest.mark.parametrize(
174 + "tool_format",
175 + ["yaml", "python", "xml_untyped", "", None, 1],
176 +)
177 +def test_invalid_request_format_rejected_before_generation(tokenizer, tool_format):
178 + parser = K2HorizonToolParser(tokenizer, _request().tools)
179 + request = _request().model_copy(
180 + update={"chat_template_kwargs": {"tool_call_format": tool_format}}
181 + )
182 +
183 + with pytest.raises(ValueError, match="Unsupported tool_call_format"):
184 + parser.adjust_request(request)
185 +
186 +
187 +@pytest.mark.parametrize(
188 + ("tool_request", "call"),
189 + [
190 + (_request(), _json_call("ping", {})),
191 + (_request("json"), _xml_call("ping", [])),
192 + (_request("xml"), _xml_call("lookup", [("query", "x")], typed=True)),
193 + (_request("xml_typed"), _xml_call("lookup", [("query", "x")])),
194 + ],
195 +)
196 +def test_mismatched_formats_fail_closed(tokenizer, tool_request, call):
197 + parser = K2HorizonToolParser(tokenizer, tool_request.tools)
198 + output = _group(call)
199 +
200 + result = parser.extract_tool_calls(output, tool_request)
201 +
202 + assert not result.tools_called
203 + assert result.content == output
204 +
205 +
206 +@pytest.mark.parametrize(
207 + "tool_choice",
208 + [
209 + "auto",
210 + "required",
211 + {"type": "function", "function": {"name": "lookup"}},
212 + ],
213 +)
214 +def test_auto_required_and_named_use_ifm_parser(tokenizer, tool_choice):
215 + request = _request("json", tool_choice)
216 + parser_cls = ParserManager.get_parser(
217 + tool_parser_name="k2_horizon",
218 + enable_auto_tools=True,
219 + )
220 + assert parser_cls is not None
221 + parser = parser_cls(tokenizer, request.tools)
222 +
223 + assert parser.adjust_request(request) is request
224 + assert request.structured_outputs is None
225 + _, content, tool_calls = parser.parse(
226 + _group(_json_call("lookup", {"enabled": "true"})),
227 + request,
228 + enable_auto_tools=True,
229 + )
230 +
231 + assert content is None
232 + assert tool_calls is not None
233 + assert json.loads(tool_calls[0].arguments) == {"enabled": True}
234 +
235 +
236 +def test_unknown_or_wrong_named_tool_fails_closed(tokenizer):
237 + named_request = _request(
238 + "json", {"type": "function", "function": {"name": "lookup"}}
239 + )
240 + parser = K2HorizonToolParser(tokenizer, named_request.tools)
241 +
242 + unknown = parser.extract_tool_calls(
243 + _group(_json_call("missing", {})), _request("json")
244 + )
245 + wrong_named = parser.extract_tool_calls(
246 + _group(_json_call("ping", {})), named_request
247 + )
248 +
249 + assert not unknown.tools_called
250 + assert not wrong_named.tools_called
251 +
252 +
253 +@pytest.mark.parametrize(
254 + "output",
255 + [
256 + "<ifm|tool_calls><ifm|tool_call>ping",
257 + "<ifm|tool_calls><ifm|tool_call>ping</ifm|tool_call>",
258 + _group("junk"),
259 + ],
260 +)
261 +def test_malformed_and_incomplete_input_is_content(tokenizer, output):
262 + request = _request()
263 + parser = K2HorizonToolParser(tokenizer, request.tools)
264 +
265 + result = parser.extract_tool_calls(output, request)
266 +
267 + assert not result.tools_called
268 + assert result.content == output
269 +
270 +
271 +@pytest.mark.parametrize(
272 + ("tool_request", "call", "expected"),
273 + [
274 + (_request(), _xml_call("lookup", [("limit", "2")]), {"limit": 2}),
275 + (_request("json"), _json_call("lookup", {"limit": "2"}), {"limit": 2}),
276 + (
277 + _request("xml_typed"),
278 + _xml_call("lookup", [("limit", "2")], typed=True),
279 + {"limit": 2},
280 + ),
281 + ],
282 +)
283 +def test_character_split_streaming(tokenizer, tool_request, call, expected):
284 + parser = K2HorizonToolParser(tokenizer, tool_request.tools)
285 + output = _group(call, prefix="Before", suffix="after")
286 +
287 + content, calls = _collect_stream(parser, tool_request, output)
288 +
289 + assert content == "Beforeafter"
290 + assert calls == [("lookup", expected)]
291 +
292 +
293 +def test_streaming_flushes_incomplete_markup_as_content(tokenizer):
294 + request = _request()
295 + parser = K2HorizonToolParser(tokenizer, request.tools)
296 + output = "Before<ifm|tool_calls><ifm|tool_call>ping"
297 +
298 + content, calls = _collect_stream(parser, request, output)
299 +
300 + assert content == output
301 + assert calls == []
302 +
303 +
304 +def test_whitespace_only_surrounding_content_is_not_preserved(tokenizer):
305 + request = _request()
306 + parser = K2HorizonToolParser(tokenizer, request.tools)
307 +
308 + result = parser.extract_tool_calls(
309 + _group(_xml_call("ping", []), prefix=" \n", suffix="\t "), request
310 + )
311 +
312 + assert result.tools_called
313 + assert result.content is None
314 +
315 +
316 +def test_tool_parser_registered():
317 + assert ToolParserManager.get_tool_parser("k2_horizon") is K2HorizonToolParser
318 + assert K2HorizonToolParser.supports_required_and_named is False
vllm/model_executor/models/k2_horizon.py new
+1445
@@ -0,0 +1,1445 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2026 MBZUAI-IFM.
5 +# Copyright 2024 The Qwen team.
6 +# Copyright 2023 The vLLM team.
7 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8 +#
9 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 +# and OPT implementations in this library. It has been modified from its
11 +# original forms to accommodate minor architectural differences compared
12 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 +#
14 +# Licensed under the Apache License, Version 2.0 (the "License");
15 +# you may not use this file except in compliance with the License.
16 +# You may obtain a copy of the License at
17 +#
18 +# http://www.apache.org/licenses/LICENSE-2.0
19 +#
20 +# Unless required by applicable law or agreed to in writing, software
21 +# distributed under the License is distributed on an "AS IS" BASIS,
22 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 +# See the License for the specific language governing permissions and
24 +# limitations under the License.
25 +"""Inference-only K2Horizon model compatible with HuggingFace weights."""
26 +
27 +import math
28 +import typing
29 +from collections.abc import Callable, Iterable
30 +from itertools import islice
31 +from typing import Any
32 +
33 +import torch
34 +import torch.nn.functional as F
35 +from torch import nn
36 +
37 +from vllm import _custom_ops as ops
38 +from vllm.compilation.decorators import support_torch_compile
39 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
40 +from vllm.distributed import (
41 + get_ep_group,
42 + get_pp_group,
43 + get_tensor_model_parallel_rank,
44 + get_tensor_model_parallel_world_size,
45 + tensor_model_parallel_all_gather,
46 +)
47 +from vllm.logger import init_logger
48 +from vllm.model_executor.layers.activation import SiluAndMul
49 +from vllm.model_executor.layers.attention import Attention
50 +from vllm.model_executor.layers.fused_moe import (
51 + FusedMoEFactory,
52 + fused_moe_make_expert_params_mapping,
53 +)
54 +from vllm.model_executor.layers.fused_moe.config import (
55 + _get_config_dtype_str,
56 +)
57 +from vllm.model_executor.layers.fused_moe.fused_moe import (
58 + _get_config_quant_dtype,
59 + _prepare_expert_assignment,
60 + dispatch_fused_moe_kernel,
61 + try_get_optimal_moe_config,
62 +)
63 +from vllm.model_executor.layers.fused_moe.utils import (
64 + moe_kernel_quantize_input,
65 +)
66 +from vllm.model_executor.layers.layernorm import RMSNorm
67 +from vllm.model_executor.layers.linear import (
68 + ColumnParallelLinear,
69 + MergedColumnParallelLinear,
70 + QKVParallelLinear,
71 + ReplicatedLinear,
72 + RowParallelLinear,
73 +)
74 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
75 +from vllm.model_executor.layers.quantization import QuantizationConfig
76 +from vllm.model_executor.layers.rotary_embedding import get_rope
77 +from vllm.model_executor.layers.vocab_parallel_embedding import (
78 + ParallelLMHead,
79 + VocabParallelEmbedding,
80 +)
81 +from vllm.model_executor.model_loader.weight_utils import (
82 + default_weight_loader,
83 + maybe_remap_kv_scale_name,
84 +)
85 +from vllm.model_executor.models.interfaces import (
86 + EagleModelMixin,
87 + MixtureOfExperts,
88 + SupportsEagle,
89 + SupportsEagle3,
90 + SupportsLoRA,
91 + SupportsPP,
92 +)
93 +from vllm.model_executor.models.utils import (
94 + AutoWeightsLoader,
95 + PPMissingLayer,
96 + extract_layer_index,
97 + is_pp_missing_parameter,
98 + make_empty_intermediate_tensors_factory,
99 + make_layers,
100 + maybe_prefix,
101 + sequence_parallel_chunk,
102 +)
103 +from vllm.sequence import IntermediateTensors
104 +from vllm.triton_utils import tl
105 +
106 +logger = init_logger(__name__)
107 +
108 +
109 +def split_to_interleaved(x):
110 + # Split halves: x0 x1 x2 x3 ... y0 y1 y2 y3 ...
111 + # Interleaved: x0 y0 x1 y1 x2 y2 x3 y3 ...
112 + return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1)
113 +
114 +
115 +def interleaved_to_split(x):
116 + # Interleaved: x0 y0 x1 y1 x2 y2 x3 y3 ...
117 + # Split halves: x0 x1 x2 x3 ... y0 y1 y2 y3 ...
118 + return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1)
119 +
120 +
121 +def calc_router_weights(
122 + router_logits: torch.Tensor,
123 + e_score_correction_bias: torch.Tensor | None,
124 + score_func: str,
125 + top_k: int,
126 + scaling_factor: float | None,
127 + renormalize: bool | None = None,
128 +) -> tuple[torch.Tensor, torch.Tensor]:
129 + if score_func == "softmax":
130 + routing_scores = F.softmax(router_logits, dim=-1, dtype=torch.float32)
131 + else:
132 + assert score_func == "sigmoid", (
133 + f"Unsupported router score function: {score_func}"
134 + )
135 + routing_scores = torch.sigmoid(router_logits.to(torch.float32))
136 +
137 + selection_scores = routing_scores
138 + if e_score_correction_bias is not None:
139 + selection_scores = selection_scores + e_score_correction_bias.to(
140 + selection_scores
141 + )
142 +
143 + selected_indices = torch.topk(selection_scores, top_k, dim=-1).indices
144 + routing_weights = torch.gather(routing_scores, dim=-1, index=selected_indices)
145 + if renormalize is None:
146 + renormalize = top_k > 1
147 + if renormalize:
148 + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
149 + if scaling_factor is not None:
150 + routing_weights = routing_weights * scaling_factor
151 + return routing_weights.to(router_logits.dtype), selected_indices
152 +
153 +
154 +def fused_mova_impl(
155 + config,
156 + hidden_states: torch.Tensor,
157 + w1: torch.Tensor,
158 + topk_weights: torch.Tensor,
159 + topk_ids: torch.Tensor,
160 + use_fp8_w8a8: bool = False,
161 + use_int8_w8a8: bool = False,
162 + use_int8_w8a16: bool = False,
163 + use_int4_w4a16: bool = False,
164 + ocp_mx_scheme: str | None = None,
165 + per_channel_quant: bool = False,
166 + global_num_experts: int = -1,
167 + expert_map: torch.Tensor | None = None,
168 + w1_scale: torch.Tensor | None = None,
169 + w1_zp: torch.Tensor | None = None,
170 + a1_scale: torch.Tensor | None = None,
171 + block_shape: list[int] | None = None,
172 + w1_bias: torch.Tensor | None = None,
173 +) -> torch.Tensor:
174 + if ocp_mx_scheme is not None:
175 + raise NotImplementedError(
176 + f"Using ocp_mx_scheme={ocp_mx_scheme} in functional fused_experts call is "
177 + "deprecated. Please use OCP_MXQuantizationEmulationTritonExperts."
178 + )
179 +
180 + # Check constraints.
181 + if use_int4_w4a16:
182 + assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
183 + else:
184 + assert hidden_states.size(1) == w1.size(2), (
185 + f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}"
186 + )
187 +
188 + assert topk_weights.size() == topk_ids.size(), "topk shape mismatch"
189 + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
190 + assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
191 + assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16]
192 +
193 + num_tokens = hidden_states.size(0)
194 + E, N, K = w1.size()
195 +
196 + if global_num_experts == -1:
197 + global_num_experts = E
198 + top_k_num = topk_ids.size(1)
199 + M = num_tokens
200 +
201 + # Note: for use_int8_w8a16 or use_int4_w4a16, the activations are
202 + # quantized prior to calling fused_experts.
203 + quant_dtype = _get_config_quant_dtype(
204 + use_fp8_w8a8=use_fp8_w8a8,
205 + use_int8_w8a8=use_int8_w8a8,
206 + )
207 +
208 + intermediate_cache1 = torch.empty(
209 + (M, top_k_num, N), device=hidden_states.device, dtype=hidden_states.dtype
210 + )
211 +
212 + if hidden_states.dtype == torch.bfloat16:
213 + compute_type = tl.bfloat16
214 + elif hidden_states.dtype == torch.float16:
215 + compute_type = tl.float16
216 + elif hidden_states.dtype == torch.float32:
217 + compute_type = tl.float32
218 + else:
219 + raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
220 +
221 + out_hidden_states = torch.empty(
222 + (M, N),
223 + device=hidden_states.device,
224 + dtype=hidden_states.dtype,
225 + )
226 +
227 + qhidden_states, a1q_scale = moe_kernel_quantize_input(
228 + A=hidden_states,
229 + A_scale=a1_scale,
230 + quant_dtype=quant_dtype,
231 + per_act_token_quant=per_channel_quant,
232 + block_shape=block_shape,
233 + )
234 +
235 + sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
236 + topk_ids,
237 + config,
238 + num_tokens,
239 + top_k_num,
240 + global_num_experts,
241 + expert_map,
242 + use_int8_w8a16=use_int8_w8a16,
243 + use_int4_w4a16=use_int4_w4a16,
244 + block_shape=block_shape,
245 + ignore_invalid_experts=True,
246 + )
247 +
248 + dispatch_fused_moe_kernel(
249 + qhidden_states,
250 + w1,
251 + intermediate_cache1,
252 + a1q_scale,
253 + w1_scale,
254 + w1_zp,
255 + topk_weights,
256 + sorted_token_ids,
257 + expert_ids,
258 + num_tokens_post_padded,
259 + False, # apply_router_weight_on_input,
260 + top_k_num,
261 + config,
262 + compute_type=compute_type,
263 + use_fp8_w8a8=use_fp8_w8a8,
264 + use_int8_w8a8=use_int8_w8a8,
265 + use_int8_w8a16=use_int8_w8a16,
266 + use_int4_w4a16=use_int4_w4a16,
267 + per_channel_quant=per_channel_quant,
268 + block_shape=block_shape,
269 + B_bias=w1_bias,
270 + )
271 +
272 + intermediate_cache1 = F.silu(intermediate_cache1)
273 + intermediate_cache1 *= topk_weights[:, :, None].to(intermediate_cache1.dtype)
274 +
275 + ops.moe_sum(intermediate_cache1, out_hidden_states)
276 +
277 + return out_hidden_states
278 +
279 +
280 +class K2HorizonRMSNorm(RMSNorm):
281 + def __init__(self, hidden_size: int, n_groups: int, eps: float = 1e-6):
282 + assert hidden_size % n_groups == 0
283 + super().__init__(hidden_size, eps=eps)
284 + self.hidden_size = hidden_size
285 + self.n_groups = n_groups
286 + self.group_size = hidden_size // n_groups
287 +
288 + def forward(self, x, residual=None):
289 + if self.n_groups == 1:
290 + return super().forward(x, residual)
291 +
292 + if residual is not None:
293 + x = x + residual
294 + residual = x
295 +
296 + orig_shape = x.shape
297 + x_grouped = x.reshape(-1, self.n_groups, self.group_size).reshape(
298 + -1, self.group_size
299 + )
300 + y_grouped = torch.empty_like(x_grouped)
301 +
302 + # Reuse vLLM's RMSNorm CUDA kernel on each group.
303 + ops.rms_norm(
304 + y_grouped,
305 + x_grouped,
306 + torch.ones(self.group_size, dtype=x.dtype, device=x.device),
307 + self.variance_epsilon,
308 + )
309 +
310 + y = y_grouped.reshape(orig_shape) * self.weight
311 + return (y, residual) if residual is not None else y
312 +
313 +
314 +class K2HorizonMLP(nn.Module):
315 + def __init__(
316 + self,
317 + hidden_size: int,
318 + intermediate_size: int,
319 + hidden_act: str,
320 + quant_config: QuantizationConfig | None = None,
321 + reduce_results: bool = True,
322 + expert_gate: torch.nn.Linear | None = None,
323 + prefix: str = "",
324 + ) -> None:
325 + super().__init__()
326 + self.gate_up_proj = MergedColumnParallelLinear(
327 + hidden_size,
328 + [intermediate_size] * 2,
329 + bias=False,
330 + quant_config=quant_config,
331 + prefix=f"{prefix}.gate_up_proj",
332 + )
333 + self.down_proj = RowParallelLinear(
334 + intermediate_size,
335 + hidden_size,
336 + bias=False,
337 + quant_config=quant_config,
338 + reduce_results=reduce_results,
339 + prefix=f"{prefix}.down_proj",
340 + )
341 + if hidden_act != "silu":
342 + raise ValueError(
343 + f"Unsupported activation: {hidden_act}. Only silu is supported for now."
344 + )
345 + self.act_fn = SiluAndMul()
346 + self.expert_gate = expert_gate
347 +
348 + def forward(self, x):
349 + gate_up, _ = self.gate_up_proj(x)
350 + out = self.act_fn(gate_up)
351 + out, _ = self.down_proj(out)
352 +
353 + if self.expert_gate is not None:
354 + out = F.sigmoid(self.expert_gate(x)[0]) * out
355 +
356 + return out
357 +
358 +
359 +class K2HorizonSparseMoeBlock(nn.Module):
360 + def __init__(
361 + self,
362 + vllm_config: VllmConfig,
363 + prefix: str = "",
364 + ):
365 + super().__init__()
366 +
367 + config = vllm_config.model_config.hf_text_config
368 + parallel_config = vllm_config.parallel_config
369 + quant_config = vllm_config.quant_config
370 +
371 + self.tp_size = get_tensor_model_parallel_world_size()
372 +
373 + self.ep_group = get_ep_group().device_group
374 + self.ep_rank = get_ep_group().rank_in_group
375 + self.ep_size = self.ep_group.size()
376 + self.n_routed_experts = config.num_experts
377 +
378 + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
379 +
380 + if self.tp_size > config.num_experts:
381 + raise ValueError(
382 + f"Tensor parallel size {self.tp_size} is greater than "
383 + f"the number of experts {config.num_experts}."
384 + )
385 +
386 + # Load balancing settings.
387 + vllm_config = get_current_vllm_config()
388 + eplb_config = vllm_config.parallel_config.eplb_config
389 + self.enable_eplb = parallel_config.enable_eplb
390 +
391 + self.n_logical_experts = self.n_routed_experts
392 + self.n_redundant_experts = eplb_config.num_redundant_experts
393 + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
394 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
395 +
396 + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
397 + self.physical_expert_end = (
398 + self.physical_expert_start + self.n_local_physical_experts
399 + )
400 +
401 + self.gate = ReplicatedLinear(
402 + config.hidden_size,
403 + config.num_experts,
404 + bias=config.moe_gate_bias,
405 + skip_bias_add=True,
406 + quant_config=quant_config,
407 + prefix=f"{prefix}.gate",
408 + )
409 + if self.gate.bias is not None:
410 + self.gate.bias = nn.Parameter(self.gate.bias.float(), requires_grad=False)
411 +
412 + self.num_shared_experts = config.num_shared_experts
413 + if config.num_shared_experts > 0:
414 + self.shared_experts = K2HorizonMLP(
415 + hidden_size=config.hidden_size,
416 + intermediate_size=config.moe_intermediate_size
417 + * config.num_shared_experts,
418 + hidden_act=config.hidden_act,
419 + quant_config=quant_config,
420 + reduce_results=False,
421 + expert_gate=None,
422 + prefix=f"{prefix}.shared_experts",
423 + )
424 + else:
425 + self.shared_experts = None
426 +
427 + self.experts = FusedMoEFactory(
428 + shared_experts=self.shared_experts,
429 + gate=self.gate,
430 + num_experts=self.n_routed_experts,
431 + top_k=config.num_experts_per_tok,
432 + hidden_size=config.hidden_size,
433 + intermediate_size=config.moe_intermediate_size,
434 + renormalize=config.norm_topk_prob,
435 + quant_config=quant_config,
436 + prefix=f"{prefix}.experts",
437 + enable_eplb=self.enable_eplb,
438 + num_redundant_experts=self.n_redundant_experts,
439 + is_sequence_parallel=self.is_sequence_parallel,
440 + scoring_func=config.router_score_func,
441 + e_score_correction_bias=self.gate.bias,
442 + routed_scaling_factor=config.router_scaling_factor,
443 + )
444 +
445 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
446 + assert hidden_states.dim() <= 2, (
447 + "K2HorizonSparseMoeBlock only supports 1D or 2D inputs"
448 + )
449 + is_input_1d = hidden_states.dim() == 1
450 + num_tokens, hidden_dim = hidden_states.shape
451 + hidden_states = hidden_states.view(-1, hidden_dim)
452 +
453 + if self.is_sequence_parallel:
454 + hidden_states = sequence_parallel_chunk(hidden_states)
455 +
456 + final_hidden_states = self.experts(
457 + hidden_states=hidden_states, router_logits=hidden_states
458 + )
459 +
460 + if self.is_sequence_parallel:
461 + final_hidden_states = tensor_model_parallel_all_gather(
462 + final_hidden_states, 0
463 + )
464 + final_hidden_states = final_hidden_states[:num_tokens]
465 +
466 + # return to 1d if input is 1d
467 + return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
468 +
469 +
470 +class K2HorizonAttention(nn.Module):
471 + def __init__(
472 + self,
473 + hidden_size: int,
474 + num_heads: int,
475 + num_kv_heads: int,
476 + rope_parameters: dict[str, Any],
477 + rope_head_dim: int,
478 + query_key_norm: bool,
479 + max_position_embeddings: int = 8192,
480 + head_dim: int | None = None,
481 + rms_norm_eps: float = 1e-06,
482 + qkv_bias: bool = False,
483 + cache_config: CacheConfig | None = None,
484 + quant_config: QuantizationConfig | None = None,
485 + prefix: str = "",
486 + dual_chunk_attention_config: dict[str, Any] | None = None,
487 + gate_func: str | None = None,
488 + ) -> None:
489 + super().__init__()
490 + self.hidden_size = hidden_size
491 + tp_size = get_tensor_model_parallel_world_size()
492 + self.total_num_heads = num_heads
493 + assert self.total_num_heads % tp_size == 0
494 + self.num_heads = self.total_num_heads // tp_size
495 + self.total_num_kv_heads = num_kv_heads
496 + if self.total_num_kv_heads >= tp_size:
497 + # Number of KV heads is greater than TP size, so we partition
498 + # the KV heads across multiple tensor parallel GPUs.
499 + assert self.total_num_kv_heads % tp_size == 0
500 + else:
501 + # Number of KV heads is less than TP size, so we replicate
502 + # the KV heads across multiple tensor parallel GPUs.
503 + assert tp_size % self.total_num_kv_heads == 0
504 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
505 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)
506 + self.q_size = self.num_heads * self.head_dim
507 + self.kv_size = self.num_kv_heads * self.head_dim
508 + self.scaling = self.head_dim**-0.5
509 + self.max_position_embeddings = max_position_embeddings
510 + self.dual_chunk_attention_config = dual_chunk_attention_config
511 +
512 + self.qkv_proj = QKVParallelLinear(
513 + hidden_size,
514 + self.head_dim,
515 + self.total_num_heads,
516 + self.total_num_kv_heads,
517 + bias=qkv_bias,
518 + quant_config=quant_config,
519 + prefix=f"{prefix}.qkv_proj",
520 + )
521 +
522 + self.o_proj = RowParallelLinear(
523 + self.total_num_heads * self.head_dim,
524 + hidden_size,
525 + bias=qkv_bias,
526 + quant_config=quant_config,
527 + prefix=f"{prefix}.o_proj",
528 + )
529 +
530 + self.rope_head_dim = rope_head_dim or self.head_dim
531 + self.rotary_emb = get_rope(
532 + self.rope_head_dim,
533 + max_position=max_position_embeddings,
534 + rope_parameters=rope_parameters,
535 + dual_chunk_attention_config=dual_chunk_attention_config,
536 + )
537 + self.attn = Attention(
538 + self.num_heads,
539 + self.head_dim,
540 + self.scaling,
541 + num_kv_heads=self.num_kv_heads,
542 + cache_config=cache_config,
543 + quant_config=quant_config,
544 + prefix=f"{prefix}.attn",
545 + **{
546 + "layer_idx": extract_layer_index(prefix),
547 + "dual_chunk_attention_config": dual_chunk_attention_config,
548 + }
549 + if dual_chunk_attention_config
550 + else {},
551 + )
552 +
553 + self.query_key_norm = query_key_norm
554 + if self.query_key_norm:
555 + self.q_norm = K2HorizonRMSNorm(
556 + hidden_size=self.q_size, n_groups=self.num_heads, eps=rms_norm_eps
557 + )
558 + self.k_norm = K2HorizonRMSNorm(
559 + hidden_size=self.kv_size, n_groups=self.num_kv_heads, eps=rms_norm_eps
560 + )
561 +
562 + self.gate_func = gate_func
563 + if self.gate_func is not None:
564 + self.gate_proj = ColumnParallelLinear(
565 + hidden_size,
566 + self.total_num_heads * self.head_dim,
567 + bias=False,
568 + quant_config=quant_config,
569 + prefix=f"{prefix}.gate_proj",
570 + )
571 +
572 + self.tp_size = get_tensor_model_parallel_world_size()
573 + self.tp_rank = get_tensor_model_parallel_rank()
574 +
575 + def forward(
576 + self,
577 + positions: torch.Tensor,
578 + hidden_states: torch.Tensor,
579 + ) -> torch.Tensor:
580 + qkv, _ = self.qkv_proj(hidden_states)
581 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
582 + if self.query_key_norm:
583 + q = self.q_norm(q)
584 + k = self.k_norm(k)
585 +
586 + if self.rope_head_dim == self.head_dim:
587 + q, k = self.rotary_emb(positions, q, k)
588 + else:
589 + q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
590 + k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
591 +
592 + q_rope, q_nope = torch.split(
593 + split_to_interleaved(q),
594 + split_size_or_sections=[
595 + self.rope_head_dim,
596 + self.head_dim - self.rope_head_dim,
597 + ],
598 + dim=-1,
599 + )
600 + k_rope, k_nope = torch.split(
601 + split_to_interleaved(k),
602 + split_size_or_sections=[
603 + self.rope_head_dim,
604 + self.head_dim - self.rope_head_dim,
605 + ],
606 + dim=-1,
607 + )
608 +
609 + q_rope, k_rope = self.rotary_emb(
610 + positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
611 + )
612 +
613 + q = interleaved_to_split(
614 + torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
615 + ).reshape(*q.shape[:-2], -1)
616 +
617 + k = interleaved_to_split(
618 + torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
619 + ).reshape(*k.shape[:-2], -1)
620 +
621 + attn_output = self.attn(q, k, v)
622 +
623 + if self.gate_func is not None:
624 + gate, _ = self.gate_proj(hidden_states)
625 + if self.gate_func == "silu":
626 + gate = F.silu(gate)
627 + else:
628 + assert self.gate_func == "softplus"
629 + gate = F.softplus(gate, beta=math.log(2))
630 +
631 + attn_output = attn_output * gate
632 +
633 + output, _ = self.o_proj(attn_output)
634 + return output
635 +
636 +
637 +class K2HorizonMoVAAttention(nn.Module):
638 + def __init__(
639 + self,
640 + vllm_config: VllmConfig,
641 + hidden_size: int,
642 + num_heads: int,
643 + num_kv_heads: int,
644 + rope_parameters: dict[str, Any],
645 + rope_head_dim: int,
646 + query_key_norm: bool,
647 + num_experts: int,
648 + num_experts_per_tok: int,
649 + moe_gate_bias: bool,
650 + router_score_func: str,
651 + router_scaling_factor: float,
652 + max_position_embeddings: int = 8192,
653 + head_dim: int | None = None,
654 + rms_norm_eps: float = 1e-06,
655 + qkv_bias: bool = False,
656 + cache_config: CacheConfig | None = None,
657 + quant_config: QuantizationConfig | None = None,
658 + prefix: str = "",
659 + dual_chunk_attention_config: dict[str, Any] | None = None,
660 + gate_func: str | None = None,
661 + ) -> None:
662 + super().__init__()
663 + self.hidden_size = hidden_size
664 + tp_size = get_tensor_model_parallel_world_size()
665 + self.total_num_heads = num_heads
666 + assert self.total_num_heads % tp_size == 0
667 + self.num_heads = self.total_num_heads // tp_size
668 + self.total_num_kv_heads = num_kv_heads
669 + if self.total_num_kv_heads >= tp_size:
670 + # Number of KV heads is greater than TP size, so we partition
671 + # the KV heads across multiple tensor parallel GPUs.
672 + assert self.total_num_kv_heads % tp_size == 0
673 + else:
674 + # Number of KV heads is less than TP size, so we replicate
675 + # the KV heads across multiple tensor parallel GPUs.
676 + assert tp_size % self.total_num_kv_heads == 0
677 +
678 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
679 + self.head_dim = head_dim or (hidden_size // self.total_num_heads)
680 + self.q_size = self.num_heads * self.head_dim
681 + self.kv_size = self.num_kv_heads * self.head_dim
682 + self.scaling = self.head_dim**-0.5
683 + self.max_position_embeddings = max_position_embeddings
684 + self.dual_chunk_attention_config = dual_chunk_attention_config
685 +
686 + self.num_experts = num_experts
687 + self.num_experts_per_tok = num_experts_per_tok
688 + self.router_score_func = router_score_func
689 + self.router_scaling_factor = router_scaling_factor
690 +
691 + self.qk_proj = MergedColumnParallelLinear(
692 + hidden_size,
693 + [self.total_num_heads * self.head_dim, self.kv_size * tp_size],
694 + bias=qkv_bias,
695 + quant_config=quant_config,
696 + prefix=f"{prefix}.qk_proj",
697 + )
698 +
699 + self.v_router = ReplicatedLinear(
700 + hidden_size,
701 + num_experts,
702 + bias=moe_gate_bias,
703 + skip_bias_add=True,
704 + quant_config=quant_config,
705 + prefix=f"{prefix}.v_router",
706 + )
707 + if self.v_router.bias is not None:
708 + self.v_router.bias = nn.Parameter(
709 + self.v_router.bias.float(), requires_grad=False
710 + )
711 +
712 + self.v_experts = nn.ModuleList(
713 + [
714 + ColumnParallelLinear(
715 + hidden_size,
716 + self.kv_size * tp_size,
717 + bias=False,
718 + quant_config=quant_config,
719 + prefix=f"{prefix}.v_experts.{value}",
720 + )
721 + for value in range(num_experts)
722 + ]
723 + )
724 +
725 + self.o_proj = RowParallelLinear(
726 + self.total_num_heads * self.head_dim,
727 + hidden_size,
728 + bias=qkv_bias,
729 + quant_config=quant_config,
730 + prefix=f"{prefix}.o_proj",
731 + )
732 +
733 + self.rope_head_dim = rope_head_dim or self.head_dim
734 + self.rotary_emb = get_rope(
735 + self.rope_head_dim,
736 + max_position=max_position_embeddings,
737 + rope_parameters=rope_parameters,
738 + dual_chunk_attention_config=dual_chunk_attention_config,
739 + )
740 + self.attn = Attention(
741 + self.num_heads,
742 + self.head_dim,
743 + self.scaling,
744 + num_kv_heads=self.num_kv_heads,
745 + cache_config=cache_config,
746 + quant_config=quant_config,
747 + prefix=f"{prefix}.attn",
748 + **{
749 + "layer_idx": extract_layer_index(prefix),
750 + "dual_chunk_attention_config": dual_chunk_attention_config,
751 + }
752 + if dual_chunk_attention_config
753 + else {},
754 + )
755 +
756 + self.query_key_norm = query_key_norm
757 + if self.query_key_norm:
758 + self.q_norm = K2HorizonRMSNorm(
759 + hidden_size=self.q_size, n_groups=self.num_heads, eps=rms_norm_eps
760 + )
761 + self.k_norm = K2HorizonRMSNorm(
762 + hidden_size=self.kv_size, n_groups=self.num_kv_heads, eps=rms_norm_eps
763 + )
764 +
765 + self.gate_func = gate_func
766 + if self.gate_func is not None:
767 + self.gate_proj = ColumnParallelLinear(
768 + hidden_size,
769 + self.total_num_heads * self.head_dim,
770 + bias=False,
771 + quant_config=quant_config,
772 + prefix=f"{prefix}.gate_proj",
773 + )
774 +
775 + self.tp_size = get_tensor_model_parallel_world_size()
776 + self.tp_rank = get_tensor_model_parallel_rank()
777 +
778 + w1_shape = torch.Size(
779 + [
780 + self.num_experts,
781 + self.total_num_kv_heads * self.head_dim,
782 + hidden_size,
783 + ]
784 + )
785 + config_dtype = _get_config_dtype_str(
786 + use_fp8_w8a8=False,
787 + use_int8_w8a16=False,
788 + use_int4_w4a16=False,
789 + dtype=vllm_config.model_config.dtype,
790 + )
791 + self.fused_mova_config = try_get_optimal_moe_config(
792 + w1_shape=w1_shape,
793 + w2_shape=torch.Size([w1_shape[0], w1_shape[2], w1_shape[1]]),
794 + top_k=self.num_experts_per_tok,
795 + dtype=config_dtype,
796 + M=vllm_config.scheduler_config.max_num_batched_tokens,
797 + block_shape=None,
798 + )
799 +
800 + def compute_mova_v_sparse(self, hidden_states):
801 + router_logits, _ = self.v_router(hidden_states)
802 +
803 + routing_weights, selected_values = calc_router_weights(
804 + router_logits=router_logits,
805 + e_score_correction_bias=self.v_router.bias,
806 + score_func=self.router_score_func,
807 + top_k=self.num_experts_per_tok,
808 + scaling_factor=self.router_scaling_factor,
809 + )
810 +
811 + w1 = torch.stack(
812 + [expert.weight for expert in self.v_experts], dim=0
813 + ).contiguous()
814 +
815 + v = fused_mova_impl(
816 + config=self.fused_mova_config,
817 + hidden_states=hidden_states.contiguous(),
818 + w1=w1,
819 + topk_weights=routing_weights.contiguous(),
820 + topk_ids=selected_values.contiguous(),
821 + global_num_experts=self.num_experts,
822 + expert_map=None,
823 + )
824 +
825 + return v
826 +
827 + def forward(
828 + self,
829 + positions: torch.Tensor,
830 + hidden_states: torch.Tensor,
831 + ) -> torch.Tensor:
832 + qk, _ = self.qk_proj(hidden_states)
833 + q, k = qk.split([self.q_size, self.kv_size], dim=-1)
834 + v = self.compute_mova_v_sparse(hidden_states)
835 +
836 + if self.query_key_norm:
837 + q = self.q_norm(q)
838 + k = self.k_norm(k)
839 +
840 + if self.rope_head_dim == self.head_dim:
841 + q, k = self.rotary_emb(positions, q, k)
842 + else:
843 + q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
844 + k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
845 +
846 + q_rope, q_nope = torch.split(
847 + split_to_interleaved(q),
848 + split_size_or_sections=[
849 + self.rope_head_dim,
850 + self.head_dim - self.rope_head_dim,
851 + ],
852 + dim=-1,
853 + )
854 + k_rope, k_nope = torch.split(
855 + split_to_interleaved(k),
856 + split_size_or_sections=[
857 + self.rope_head_dim,
858 + self.head_dim - self.rope_head_dim,
859 + ],
860 + dim=-1,
861 + )
862 +
863 + q_rope, k_rope = self.rotary_emb(
864 + positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
865 + )
866 +
867 + q = interleaved_to_split(
868 + torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
869 + ).reshape(*q.shape[:-2], -1)
870 +
871 + k = interleaved_to_split(
872 + torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
873 + ).reshape(*k.shape[:-2], -1)
874 +
875 + attn_output = self.attn(q, k, v)
876 +
877 + if self.gate_func is not None:
878 + gate, _ = self.gate_proj(hidden_states)
879 + if self.gate_func == "silu":
880 + gate = F.silu(gate)
881 + else:
882 + assert self.gate_func == "softplus"
883 + gate = F.softplus(gate, beta=math.log(2))
884 +
885 + attn_output = attn_output * gate
886 +
887 + output, _ = self.o_proj(attn_output)
888 + return output
889 +
890 +
891 +class K2HorizonDecoderLayer(nn.Module):
892 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
893 + super().__init__()
894 +
895 + config = vllm_config.model_config.hf_text_config
896 + cache_config = vllm_config.cache_config
897 + quant_config = vllm_config.quant_config
898 +
899 + self.hidden_size = config.hidden_size
900 + max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
901 + dual_chunk_attention_config = getattr(
902 + config, "dual_chunk_attention_config", None
903 + )
904 +
905 + # `mlp_only_layers` in the config.
906 + layer_idx = extract_layer_index(prefix)
907 + mlp_only_layers = (
908 + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
909 + )
910 + is_sparse_layer = (layer_idx not in mlp_only_layers) and (
911 + config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
912 + )
913 +
914 + if is_sparse_layer and config.mova_num_experts > 0:
915 + self.self_attn = K2HorizonMoVAAttention(
916 + vllm_config=vllm_config,
917 + hidden_size=self.hidden_size,
918 + num_heads=config.num_attention_heads,
919 + num_kv_heads=config.num_key_value_heads,
920 + rope_parameters=config.rope_parameters,
921 + rope_head_dim=config.rope_head_dim,
922 + query_key_norm=config.query_key_norm,
923 + max_position_embeddings=max_position_embeddings,
924 + rms_norm_eps=config.rms_norm_eps,
925 + qkv_bias=config.attention_bias,
926 + head_dim=config.head_dim,
927 + gate_func=config.attention_gate_func,
928 + cache_config=cache_config,
929 + quant_config=quant_config,
930 + prefix=f"{prefix}.self_attn",
931 + dual_chunk_attention_config=dual_chunk_attention_config,
932 + num_experts=config.mova_num_experts,
933 + num_experts_per_tok=config.mova_num_experts_per_tok,
934 + moe_gate_bias=config.moe_gate_bias,
935 + router_score_func=config.router_score_func,
936 + router_scaling_factor=config.router_scaling_factor,
937 + )
938 + else:
939 + self.self_attn = K2HorizonAttention(
940 + hidden_size=self.hidden_size,
941 + num_heads=config.num_attention_heads,
942 + num_kv_heads=config.num_key_value_heads,
943 + rope_parameters=config.rope_parameters,
944 + rope_head_dim=config.rope_head_dim,
945 + query_key_norm=config.query_key_norm,
946 + max_position_embeddings=max_position_embeddings,
947 + rms_norm_eps=config.rms_norm_eps,
948 + qkv_bias=config.attention_bias,
949 + head_dim=config.head_dim,
950 + gate_func=config.attention_gate_func,
951 + cache_config=cache_config,
952 + quant_config=quant_config,
953 + prefix=f"{prefix}.self_attn",
954 + dual_chunk_attention_config=dual_chunk_attention_config,
955 + )
956 +
957 + if is_sparse_layer:
958 + self.mlp = K2HorizonSparseMoeBlock(
959 + vllm_config=vllm_config, prefix=f"{prefix}.mlp"
960 + )
961 + else:
962 + self.mlp = K2HorizonMLP(
963 + hidden_size=config.hidden_size,
964 + intermediate_size=config.intermediate_size,
965 + hidden_act=config.hidden_act,
966 + quant_config=quant_config,
967 + prefix=f"{prefix}.mlp",
968 + )
969 +
970 + self.input_layernorm = K2HorizonRMSNorm(
971 + hidden_size=config.hidden_size,
972 + n_groups=config.layernorm_num_groups,
973 + eps=config.rms_norm_eps,
974 + )
975 +
976 + self.post_attention_layernorm = K2HorizonRMSNorm(
977 + hidden_size=config.hidden_size,
978 + n_groups=config.layernorm_num_groups,
979 + eps=config.rms_norm_eps,
980 + )
981 +
982 + def forward(
983 + self,
984 + positions: torch.Tensor,
985 + hidden_states: torch.Tensor,
986 + residual: torch.Tensor | None,
987 + ) -> tuple[torch.Tensor, torch.Tensor]:
988 + # Self Attention
989 + if residual is None:
990 + residual = hidden_states
991 + hidden_states = self.input_layernorm(hidden_states)
992 + else:
993 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
994 +
995 + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
996 +
997 + # Fully Connected
998 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
999 + hidden_states = self.mlp(hidden_states)
1000 + return hidden_states, residual
1001 +
1002 +
1003 +@support_torch_compile
1004 +class K2HorizonModel(nn.Module, EagleModelMixin):
1005 + def __init__(
1006 + self,
1007 + *,
1008 + vllm_config: VllmConfig,
1009 + prefix: str = "",
1010 + decoder_layer_type: type[torch.nn.Module] = K2HorizonDecoderLayer,
1011 + ):
1012 + super().__init__()
1013 +
1014 + config = vllm_config.model_config.hf_text_config
1015 + quant_config = vllm_config.quant_config
1016 + parallel_config = vllm_config.parallel_config
1017 + eplb_config = parallel_config.eplb_config
1018 + self.num_redundant_experts = eplb_config.num_redundant_experts
1019 +
1020 + self.vocab_size = config.vocab_size
1021 + self.config = config
1022 + self.quant_config = quant_config
1023 + self.embed_tokens = VocabParallelEmbedding(
1024 + config.vocab_size,
1025 + config.hidden_size,
1026 + quant_config=quant_config,
1027 + prefix=f"{prefix}.embed_tokens",
1028 + )
1029 + self.start_layer, self.end_layer, self.layers = make_layers(
1030 + config.num_hidden_layers,
1031 + lambda prefix: decoder_layer_type(vllm_config=vllm_config, prefix=prefix),
1032 + prefix=f"{prefix}.layers",
1033 + )
1034 +
1035 + self.norm = K2HorizonRMSNorm(
1036 + hidden_size=config.hidden_size,
1037 + n_groups=config.layernorm_num_groups,
1038 + eps=config.rms_norm_eps,
1039 + )
1040 +
1041 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
1042 + ["hidden_states", "residual"], config.hidden_size
1043 + )
1044 +
1045 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1046 + return self.embed_tokens(input_ids)
1047 +
1048 + def forward(
1049 + self,
1050 + input_ids: torch.Tensor | None,
1051 + positions: torch.Tensor,
1052 + intermediate_tensors: IntermediateTensors | None = None,
1053 + inputs_embeds: torch.Tensor | None = None,
1054 + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
1055 + if get_pp_group().is_first_rank:
1056 + if inputs_embeds is not None:
1057 + hidden_states = inputs_embeds
1058 + else:
1059 + hidden_states = self.embed_input_ids(input_ids)
1060 + residual = None
1061 + else:
1062 + assert intermediate_tensors is not None
1063 + hidden_states = intermediate_tensors["hidden_states"]
1064 + residual = intermediate_tensors["residual"]
1065 +
1066 + aux_hidden_states = self._maybe_add_hidden_state(
1067 + [], self.start_layer, hidden_states, residual
1068 + )
1069 + for layer_idx, layer in enumerate(
1070 + islice(self.layers, self.start_layer, self.end_layer),
1071 + start=self.start_layer,
1072 + ):
1073 + hidden_states, residual = layer(positions, hidden_states, residual)
1074 + self._maybe_add_hidden_state(
1075 + aux_hidden_states, layer_idx + 1, hidden_states, residual
1076 + )
1077 +
1078 + if not get_pp_group().is_last_rank:
1079 + return IntermediateTensors(
1080 + {"hidden_states": hidden_states, "residual": residual}
1081 + )
1082 + hidden_states, _ = self.norm(hidden_states, residual)
1083 +
1084 + # Return auxiliary hidden states if collected
1085 + if len(aux_hidden_states) > 0:
1086 + return hidden_states, aux_hidden_states
1087 + return hidden_states
1088 +
1089 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1090 + # Params for weights, fp8 weight scales, fp8 activation scales
1091 + # (param_name, weight_name, expert_id, shard_id)
1092 + return fused_moe_make_expert_params_mapping(
1093 + self,
1094 + ckpt_gate_proj_name="gate_proj",
1095 + ckpt_down_proj_name="down_proj",
1096 + ckpt_up_proj_name="up_proj",
1097 + num_experts=self.config.num_experts,
1098 + num_redundant_experts=self.num_redundant_experts,
1099 + )
1100 +
1101 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1102 + stacked_params_mapping = [
1103 + # (param_name, shard_name, shard_id)
1104 + ("qkv_proj", "q_proj", "q"),
1105 + ("qkv_proj", "k_proj", "k"),
1106 + ("qkv_proj", "v_proj", "v"),
1107 + ("gate_up_proj", "gate_proj", 0),
1108 + ("gate_up_proj", "up_proj", 1),
1109 + ]
1110 +
1111 + # Skip loading extra parameters for GPTQ/modelopt models.
1112 + # ignore_suffixes = (
1113 + # ".bias",
1114 + # "_bias",
1115 + # ".weight_scale",
1116 + # "_weight_scale",
1117 + # ".input_scale",
1118 + # "_input_scale",
1119 + # )
1120 + ignore_suffixes = ()
1121 +
1122 + params_dict = dict(self.named_parameters())
1123 +
1124 + loaded_params: set[str] = set()
1125 + expert_params_mapping = self.get_expert_mapping()
1126 + for name, loaded_weight in weights:
1127 + if "scale" in name or "zero_point" in name:
1128 + name = maybe_remap_kv_scale_name(name, params_dict)
1129 + if name is None:
1130 + continue
1131 +
1132 + # QK norm weights
1133 + if name.endswith(".self_attn.q_norm.weight") or name.endswith(
1134 + ".self_attn.k_norm.weight"
1135 + ):
1136 + if is_pp_missing_parameter(name, self):
1137 + continue
1138 + if name not in params_dict:
1139 + continue
1140 +
1141 + param = params_dict[name]
1142 + tp_rank = get_tensor_model_parallel_rank()
1143 + tp_size = get_tensor_model_parallel_world_size()
1144 +
1145 + if loaded_weight.shape != param.shape:
1146 + if name.endswith(".self_attn.q_norm.weight"):
1147 + loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
1148 + else:
1149 + num_kv_heads = self.config.num_key_value_heads
1150 + if num_kv_heads >= tp_size:
1151 + loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
1152 + else:
1153 + num_kv_head_replicas = tp_size // num_kv_heads
1154 + kv_rank = tp_rank // num_kv_head_replicas
1155 + loaded_weight = loaded_weight.chunk(num_kv_heads, dim=0)[
1156 + kv_rank
1157 + ]
1158 +
1159 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
1160 + weight_loader(param, loaded_weight)
1161 + loaded_params.add(name)
1162 + continue
1163 +
1164 + # MoVA v experts
1165 + if ".self_attn.v_experts." in name:
1166 + if is_pp_missing_parameter(name, self):
1167 + continue
1168 + if name not in params_dict:
1169 + continue
1170 +
1171 + param = params_dict[name]
1172 + tp_rank = get_tensor_model_parallel_rank()
1173 + tp_size = get_tensor_model_parallel_world_size()
1174 + num_kv_heads = self.config.num_key_value_heads
1175 +
1176 + if loaded_weight.shape != param.shape:
1177 + if num_kv_heads >= tp_size:
1178 + loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
1179 + else:
1180 + num_kv_head_replicas = tp_size // num_kv_heads
1181 + kv_rank = tp_rank // num_kv_head_replicas
1182 + loaded_weight = loaded_weight.chunk(num_kv_heads, dim=0)[
1183 + kv_rank
1184 + ]
1185 +
1186 + default_weight_loader(param, loaded_weight)
1187 + loaded_params.add(name)
1188 + continue
1189 +
1190 + for param_name, weight_name, shard_id in stacked_params_mapping:
1191 + # Skip non-stacked layers and experts (experts handled below).
1192 + if weight_name not in name:
1193 + continue
1194 + # We have mlp.experts[0].gate_proj in the checkpoint.
1195 + # Since we handle the experts below in expert_params_mapping,
1196 + # we need to skip here BEFORE we update the name, otherwise
1197 + # name will be updated to mlp.experts[0].gate_up_proj, which
1198 + # will then be updated below in expert_params_mapping
1199 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
1200 + if "mlp.experts" in name:
1201 + continue
1202 + if ".self_attn." in name and param_name == "gate_up_proj":
1203 + assert weight_name == "gate_proj"
1204 + continue
1205 + if (
1206 + ".self_attn." in name
1207 + and param_name == "qkv_proj"
1208 + and name.replace(weight_name, param_name) not in params_dict
1209 + ):
1210 + qk_name = name.replace(weight_name, "qk_proj")
1211 + if is_pp_missing_parameter(qk_name, self):
1212 + continue
1213 + if qk_name not in params_dict:
1214 + continue
1215 +
1216 + assert weight_name in ["q_proj", "k_proj"]
1217 + if weight_name == "k_proj":
1218 + tp_size = get_tensor_model_parallel_world_size()
1219 + num_kv_heads = self.config.num_key_value_heads
1220 + if num_kv_heads < tp_size:
1221 + num_kv_head_replicas = tp_size // num_kv_heads
1222 + loaded_weight = (
1223 + loaded_weight.reshape(
1224 + num_kv_heads, -1, *loaded_weight.shape[1:]
1225 + )
1226 + .repeat_interleave(num_kv_head_replicas, dim=0)
1227 + .reshape(-1, *loaded_weight.shape[1:])
1228 + )
1229 +
1230 + param = params_dict[qk_name]
1231 + weight_loader = getattr(
1232 + param, "weight_loader", default_weight_loader
1233 + )
1234 + weight_loader(
1235 + param, loaded_weight, 0 if weight_name == "q_proj" else 1
1236 + )
1237 + loaded_params.add(qk_name)
1238 + break
1239 +
1240 + name = name.replace(weight_name, param_name)
1241 +
1242 + # Skip loading extra parameters for GPTQ/modelopt models.
1243 + if name.endswith(ignore_suffixes) and name not in params_dict:
1244 + continue
1245 +
1246 + # Skip layers on other devices.
1247 + if is_pp_missing_parameter(name, self):
1248 + continue
1249 + if name.endswith("scale"):
1250 + # Remapping the name of FP8 kv-scale.
1251 + name = maybe_remap_kv_scale_name(name, params_dict)
1252 + if name is None:
1253 + continue
1254 + if name not in params_dict:
1255 + continue
1256 +
1257 + param = params_dict[name]
1258 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
1259 + if weight_loader == default_weight_loader:
1260 + weight_loader(param, loaded_weight)
1261 + else:
1262 + weight_loader(param, loaded_weight, shard_id)
1263 + break
1264 + else:
1265 + is_expert_weight = False
1266 + for mapping in expert_params_mapping:
1267 + param_name, weight_name, expert_id, shard_id = mapping
1268 + if weight_name not in name:
1269 + continue
1270 +
1271 + # Anyway, this is an expert weight and should not be
1272 + # attempted to load as other weights later
1273 + is_expert_weight = True
1274 +
1275 + # Do not modify `name` since the loop may continue here
1276 + # Instead, create a new variable
1277 + name_mapped = name.replace(weight_name, param_name)
1278 +
1279 + if is_pp_missing_parameter(name_mapped, self):
1280 + continue
1281 +
1282 + # Skip loading extra parameters for GPTQ/modelopt models.
1283 + if (
1284 + name_mapped.endswith(ignore_suffixes)
1285 + and name_mapped not in params_dict
1286 + ):
1287 + continue
1288 +
1289 + param = params_dict[name_mapped]
1290 + # We should ask the weight loader to return success or not
1291 + # here since otherwise we may skip experts with other
1292 + # available replicas.
1293 + weight_loader = typing.cast(
1294 + Callable[..., bool], param.weight_loader
1295 + )
1296 + success = weight_loader(
1297 + param,
1298 + loaded_weight,
1299 + name_mapped,
1300 + shard_id=shard_id,
1301 + expert_id=expert_id,
1302 + return_success=True,
1303 + )
1304 + if success:
1305 + name = name_mapped
1306 + break
1307 + else:
1308 + if is_expert_weight:
1309 + # We've checked that this is an expert weight
1310 + # However it's not mapped locally to this rank
1311 + # So we simply skip it
1312 + continue
1313 +
1314 + # Skip loading extra parameters for GPTQ/modelopt models.
1315 + if name.endswith(ignore_suffixes) and name not in params_dict:
1316 + continue
1317 + # Skip layers on other devices.
1318 + if is_pp_missing_parameter(name, self):
1319 + continue
1320 + if name not in params_dict:
1321 + continue
1322 + param = params_dict[name]
1323 + weight_loader = getattr(
1324 + param, "weight_loader", default_weight_loader
1325 + )
1326 + weight_loader(param, loaded_weight)
1327 + loaded_params.add(name)
1328 +
1329 + return loaded_params
1330 +
1331 +
1332 +class K2HorizonForCausalLM(
1333 + nn.Module, SupportsPP, SupportsLoRA, SupportsEagle, SupportsEagle3, MixtureOfExperts
1334 +):
1335 + packed_modules_mapping = {
1336 + "qkv_proj": [
1337 + "q_proj",
1338 + "k_proj",
1339 + "v_proj",
1340 + ],
1341 + "qk_proj": [
1342 + "q_proj",
1343 + "k_proj",
1344 + ],
1345 + }
1346 +
1347 + embedding_modules = {
1348 + "embed_tokens": "input_embeddings",
1349 + "lm_head": "output_embeddings",
1350 + }
1351 +
1352 + fall_back_to_pt_during_load = False
1353 +
1354 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1355 + super().__init__()
1356 + config = vllm_config.model_config.hf_text_config
1357 + quant_config = vllm_config.quant_config
1358 + self.config = config
1359 + self.quant_config = quant_config
1360 + # Only perform the following mapping when K2HorizonMLP exists
1361 + if getattr(config, "mlp_only_layers", []):
1362 + self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
1363 + self.model = K2HorizonModel(
1364 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
1365 + )
1366 + self.lm_head = ParallelLMHead(
1367 + config.vocab_size,
1368 + config.hidden_size,
1369 + quant_config=quant_config,
1370 + prefix=maybe_prefix(prefix, "lm_head"),
1371 + )
1372 + if self.config.tie_word_embeddings:
1373 + self.lm_head.weight = self.model.embed_tokens.weight
1374 + self.logits_processor = LogitsProcessor(config.vocab_size)
1375 + self.make_empty_intermediate_tensors = (
1376 + self.model.make_empty_intermediate_tensors
1377 + )
1378 +
1379 + # Set MoE hyperparameters
1380 + self.expert_weights = []
1381 +
1382 + self.moe_layers = []
1383 + example_layer = None
1384 + for layer in self.model.layers:
1385 + if isinstance(layer, PPMissingLayer):
1386 + continue
1387 +
1388 + assert isinstance(layer, K2HorizonDecoderLayer)
1389 + if isinstance(layer.mlp, K2HorizonSparseMoeBlock):
1390 + example_layer = layer.mlp
1391 + self.moe_layers.append(layer.mlp.experts)
1392 +
1393 + self.num_moe_layers = len(self.moe_layers)
1394 + if example_layer is not None:
1395 + self.num_logical_experts = example_layer.n_logical_experts
1396 + self.num_physical_experts = example_layer.n_physical_experts
1397 + self.num_local_physical_experts = example_layer.n_local_physical_experts
1398 + self.num_routed_experts = example_layer.n_routed_experts
1399 + self.num_redundant_experts = example_layer.n_redundant_experts
1400 +
1401 + def update_physical_experts_metadata(
1402 + self,
1403 + num_physical_experts: int,
1404 + num_local_physical_experts: int,
1405 + ) -> None:
1406 + assert self.num_local_physical_experts == num_local_physical_experts
1407 + self.num_physical_experts = num_physical_experts
1408 + self.num_local_physical_experts = num_local_physical_experts
1409 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1410 + for layer in self.model.layers:
1411 + if isinstance(layer.mlp, K2HorizonSparseMoeBlock):
1412 + moe = layer.mlp
1413 + moe.n_local_physical_experts = num_local_physical_experts
1414 + moe.n_physical_experts = num_physical_experts
1415 + moe.n_redundant_experts = self.num_redundant_experts
1416 + moe.experts.update_expert_map()
1417 +
1418 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1419 + return self.model.embed_input_ids(input_ids)
1420 +
1421 + def forward(
1422 + self,
1423 + input_ids: torch.Tensor | None,
1424 + positions: torch.Tensor,
1425 + intermediate_tensors: IntermediateTensors | None = None,
1426 + inputs_embeds: torch.Tensor | None = None,
1427 + ) -> torch.Tensor | IntermediateTensors:
1428 + hidden_states = self.model(
1429 + input_ids, positions, intermediate_tensors, inputs_embeds
1430 + )
1431 + return hidden_states
1432 +
1433 + def compute_logits(
1434 + self,
1435 + hidden_states: torch.Tensor,
1436 + ) -> torch.Tensor | None:
1437 + logits = self.logits_processor(self.lm_head, hidden_states)
1438 + return logits
1439 +
1440 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1441 + loader = AutoWeightsLoader(self)
1442 + return loader.load_weights(weights)
1443 +
1444 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1445 + return self.model.get_expert_mapping()
vllm/model_executor/models/registry.py
+1
@@ -141,6 +141,7 @@ _TEXT_GENERATION_MODELS = {
141 "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
142 "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
143 "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
144 + "K2HorizonForCausalLM": ("k2_horizon", "K2HorizonForCausalLM"),
145 "KimiLinearForCausalLM": (
146 "vllm.models.kimi_k3",
147 "KimiLinearForCausalLM",
vllm/reasoning/__init__.py
+4
@@ -96,6 +96,10 @@ _REASONING_PARSERS_TO_REGISTER = {
96 "kimi_k3_reasoning_parser",
97 "KimiK3ReasoningParser",
98 ),
99 + "k2_horizon": (
100 + "k2_horizon_reasoning_parser",
101 + "K2HorizonReasoningParser",
102 + ),
103 "mimo": (
104 "qwen3_engine_reasoning_parser",
105 "Qwen3ParserReasoningAdapter",
vllm/reasoning/k2_horizon_reasoning_parser.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable, Sequence
5 +
6 +from vllm.entrypoints.generate.base.protocol import DeltaMessage
7 +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
8 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
9 +from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
10 +from vllm.tokenizers import TokenizerLike
11 +from vllm.tool_parsers.utils import partial_tag_overlap
12 +
13 +
14 +class K2HorizonReasoningParser(DeepSeekR1ReasoningParser):
15 + """Reasoning parser for K2 Horizon."""
16 +
17 + _TOOL_CALLS_START = "<ifm|tool_calls>"
18 + _EFFORT_TOKENS: dict[str, tuple[str, str]] = {
19 + "high": ("<ifm|think>", "</ifm|think>"),
20 + "medium": ("<ifm|think_fast>", "</ifm|think_fast>"),
21 + "low": ("<ifm|think_faster>", "</ifm|think_faster>"),
22 + }
23 +
24 + def __init__(self, tokenizer: TokenizerLike, *args, **kwargs) -> None:
25 + chat_template_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
26 + effort = chat_template_kwargs.get("reasoning_effort", "high")
27 + if not isinstance(effort, str) or effort not in self._EFFORT_TOKENS:
28 + supported = ", ".join(self._EFFORT_TOKENS)
29 + raise ValueError(
30 + f"Unsupported reasoning_effort {effort!r}. "
31 + f"Supported values: {supported}."
32 + )
33 +
34 + self._start_token, self._end_token = self._EFFORT_TOKENS[effort]
35 + super().__init__(tokenizer, *args, **kwargs)
36 + self._tool_calls_start_id = self.vocab.get(self._TOOL_CALLS_START)
37 + self._stream_buffer = ""
38 + self._stream_at_start = True
39 + self._stream_finished = False
40 +
41 + @property
42 + def start_token(self) -> str:
43 + return self._start_token
44 +
45 + @property
46 + def end_token(self) -> str:
47 + return self._end_token
48 +
49 + def _without_generated_start(self, model_output: str) -> str:
50 + return model_output.removeprefix(self.start_token)
51 +
52 + def _split_model_output(self, model_output: str) -> tuple[str, str | None]:
53 + output = self._without_generated_start(model_output)
54 + if self.end_token in output:
55 + reasoning, _, content = output.partition(self.end_token)
56 + return reasoning, content or None
57 +
58 + tool_start = output.find(self._TOOL_CALLS_START)
59 + if tool_start != -1:
60 + return output[:tool_start], output[tool_start:]
61 +
62 + return "", output or None
63 +
64 + def extract_reasoning(
65 + self,
66 + model_output: str,
67 + request: ChatCompletionRequest | ResponsesRequest,
68 + ) -> tuple[str, str | None]:
69 + del request
70 + return self._split_model_output(model_output)
71 +
72 + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
73 + for token_id in reversed(input_ids):
74 + if token_id == self.start_token_id:
75 + return False
76 + if token_id == self.end_token_id or token_id == self._tool_calls_start_id:
77 + return True
78 + return False
79 +
80 + def is_reasoning_end_streaming(
81 + self,
82 + input_ids: Sequence[int],
83 + delta_ids: Iterable[int],
84 + ) -> bool:
85 + del input_ids
86 + return self._stream_finished or super().is_reasoning_end_streaming(
87 + (), delta_ids
88 + )
89 +
90 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:
91 + if self.end_token_id in input_ids:
92 + return input_ids[input_ids.index(self.end_token_id) + 1 :]
93 + if (
94 + self._tool_calls_start_id is not None
95 + and self._tool_calls_start_id in input_ids
96 + ):
97 + return input_ids[input_ids.index(self._tool_calls_start_id) :]
98 + return input_ids if self._stream_finished else []
99 +
100 + def _strip_stream_start(self) -> bool:
101 + if not self._stream_at_start:
102 + return True
103 + if self.start_token.startswith(self._stream_buffer):
104 + if self._stream_buffer != self.start_token:
105 + return False
106 + self._stream_buffer = ""
107 + elif self._stream_buffer.startswith(self.start_token):
108 + self._stream_buffer = self._stream_buffer[len(self.start_token) :]
109 + self._stream_at_start = False
110 + return True
111 +
112 + def extract_reasoning_streaming(
113 + self,
114 + previous_text: str,
115 + current_text: str,
116 + delta_text: str,
117 + previous_token_ids: Sequence[int],
118 + current_token_ids: Sequence[int],
119 + delta_token_ids: Sequence[int],
120 + ) -> DeltaMessage | None:
121 + del (
122 + previous_text,
123 + current_text,
124 + previous_token_ids,
125 + current_token_ids,
126 + delta_token_ids,
127 + )
128 +
129 + if self._stream_finished:
130 + return DeltaMessage(content=delta_text) if delta_text else None
131 +
132 + self._stream_buffer += delta_text
133 + if not self._strip_stream_start():
134 + return None
135 +
136 + explicit_end = self._stream_buffer.find(self.end_token)
137 + tool_start = self._stream_buffer.find(self._TOOL_CALLS_START)
138 + if explicit_end != -1:
139 + reasoning = self._stream_buffer[:explicit_end]
140 + content = self._stream_buffer[explicit_end + len(self.end_token) :]
141 + elif tool_start != -1:
142 + reasoning = self._stream_buffer[:tool_start]
143 + content = self._stream_buffer[tool_start:]
144 + else:
145 + overlap = max(
146 + partial_tag_overlap(self._stream_buffer, self.end_token),
147 + partial_tag_overlap(self._stream_buffer, self._TOOL_CALLS_START),
148 + )
149 + sendable_end = len(self._stream_buffer) - overlap
150 + reasoning = self._stream_buffer[:sendable_end]
151 + self._stream_buffer = self._stream_buffer[sendable_end:]
152 + return DeltaMessage(reasoning=reasoning) if reasoning else None
153 +
154 + self._stream_buffer = ""
155 + self._stream_finished = True
156 + return DeltaMessage(
157 + reasoning=reasoning,
158 + content=content or None,
159 + )
vllm/tool_parsers/__init__.py
+4
@@ -118,6 +118,10 @@ _TOOL_PARSERS_TO_REGISTER = {
118 "kimi_k3_tool_parser",
119 "KimiK3ToolParser",
120 ),
121 + "k2_horizon": (
122 + "k2_horizon_tool_parser",
123 + "K2HorizonToolParser",
124 + ),
125 "llama3_json": (
126 "llama_tool_parser",
127 "Llama3JsonToolParser",
vllm/tool_parsers/k2_horizon_tool_parser.py new
+462
@@ -0,0 +1,462 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from collections.abc import Sequence
6 +from typing import Any
7 +
8 +import regex as re
9 +from openai.types.responses import ToolChoiceFunction
10 +
11 +from vllm.entrypoints.generate.base.protocol import (
12 + DeltaFunctionCall,
13 + DeltaMessage,
14 + DeltaToolCall,
15 + ExtractedToolCallInformation,
16 + FunctionCall,
17 + ToolCall,
18 +)
19 +from vllm.entrypoints.openai.chat_completion.protocol import (
20 + ChatCompletionNamedToolChoiceParam,
21 + ChatCompletionRequest,
22 +)
23 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24 +from vllm.logger import init_logger
25 +from vllm.tokenizers import TokenizerLike
26 +from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
27 +from vllm.tool_parsers.utils import (
28 + coerce_to_schema_type,
29 + extract_types_from_schema,
30 + find_tool_name,
31 + find_tool_properties,
32 + partial_tag_overlap,
33 +)
34 +
35 +logger = init_logger(__name__)
36 +
37 +ToolParserRequest = ChatCompletionRequest | ResponsesRequest
38 +
39 +
40 +class K2HorizonToolParser(ToolParser):
41 + """Parser for K2 Horizon's IFM tool-call formats."""
42 +
43 + supports_required_and_named = False
44 + engine_based_streaming = True
45 +
46 + _SUPPORTED_TOOL_FORMATS = frozenset({"json", "xml", "xml_typed"})
47 +
48 + _TOOL_CALLS_START = "<ifm|tool_calls>"
49 + _TOOL_CALLS_END = "</ifm|tool_calls>"
50 + _TOOL_CALL_START = "<ifm|tool_call>"
51 + _TOOL_CALL_END = "</ifm|tool_call>"
52 + _ARG_KEY_START = "<ifm|arg_key>"
53 + _ARG_KEY_END = "</ifm|arg_key>"
54 + _ARG_TYPE_START = "<ifm|arg_type>"
55 + _ARG_TYPE_END = "</ifm|arg_type>"
56 + _ARG_VALUE_START = "<ifm|arg_value>"
57 + _ARG_VALUE_END = "</ifm|arg_value>"
58 +
59 + _TOOL_CALL_RE = re.compile(
60 + re.escape(_TOOL_CALL_START) + r"(.*?)" + re.escape(_TOOL_CALL_END),
61 + re.DOTALL,
62 + )
63 + _ARG_RE = re.compile(
64 + re.escape(_ARG_KEY_START)
65 + + r"(.*?)"
66 + + re.escape(_ARG_KEY_END)
67 + + r"\s*(?:"
68 + + re.escape(_ARG_TYPE_START)
69 + + r"(.*?)"
70 + + re.escape(_ARG_TYPE_END)
71 + + r"\s*)?"
72 + + re.escape(_ARG_VALUE_START)
73 + + r"(.*?)"
74 + + re.escape(_ARG_VALUE_END),
75 + re.DOTALL,
76 + )
77 + _ARG_MARKERS = (
78 + _ARG_KEY_START,
79 + _ARG_KEY_END,
80 + _ARG_TYPE_START,
81 + _ARG_TYPE_END,
82 + _ARG_VALUE_START,
83 + _ARG_VALUE_END,
84 + )
85 +
86 + def __init__(
87 + self,
88 + tokenizer: TokenizerLike,
89 + tools: list[Tool] | None = None,
90 + ) -> None:
91 + super().__init__(tokenizer, tools)
92 + self.tool_format = "xml"
93 + self._stream_text = ""
94 + self._stream_cursor = 0
95 + self._stream_group_start: int | None = None
96 + self._stream_resolved = False
97 + self._pending_content = ""
98 + self._content_seen = False
99 +
100 + @classmethod
101 + def _request_tool_format(cls, request: ToolParserRequest) -> str:
102 + chat_template_kwargs = request.chat_template_kwargs or {}
103 + tool_format = chat_template_kwargs.get("tool_call_format", "xml")
104 + if (
105 + not isinstance(tool_format, str)
106 + or tool_format not in cls._SUPPORTED_TOOL_FORMATS
107 + ):
108 + supported = ", ".join(sorted(cls._SUPPORTED_TOOL_FORMATS))
109 + raise ValueError(
110 + f"Unsupported tool_call_format {tool_format!r}. "
111 + f"Supported values: {supported}."
112 + )
113 + return tool_format
114 +
115 + def adjust_request(self, request: ToolParserRequest) -> ToolParserRequest:
116 + self._request_tool_format(request)
117 + return request
118 +
119 + @staticmethod
120 + def _content_or_none(content: str) -> str | None:
121 + return content if content.strip() else None
122 +
123 + @classmethod
124 + def _no_tool_calls(cls, content: str) -> ExtractedToolCallInformation:
125 + return ExtractedToolCallInformation(
126 + tools_called=False,
127 + tool_calls=[],
128 + content=cls._content_or_none(content),
129 + )
130 +
131 + @staticmethod
132 + def _make_tool_call(name: str, arguments: dict[str, Any]) -> ToolCall:
133 + return ToolCall(
134 + type="function",
135 + function=FunctionCall(
136 + name=name,
137 + arguments=json.dumps(
138 + arguments,
139 + ensure_ascii=False,
140 + allow_nan=False,
141 + ),
142 + ),
143 + )
144 +
145 + @staticmethod
146 + def _base_arg_type(arg_type: str | None) -> str | None:
147 + if not arg_type:
148 + return None
149 + return arg_type.split("[", 1)[0].strip().lower()
150 +
151 + @staticmethod
152 + def _coerce_json_value(value: Any, schema: Any) -> Any:
153 + if not isinstance(schema, dict):
154 + return value
155 + raw_value = (
156 + value
157 + if isinstance(value, str)
158 + else json.dumps(value, ensure_ascii=False, allow_nan=False)
159 + )
160 + return coerce_to_schema_type(raw_value, extract_types_from_schema(schema))
161 +
162 + @classmethod
163 + def _coerce_xml_value(
164 + cls,
165 + value: str,
166 + schema: Any,
167 + explicit_type: str | None,
168 + ) -> Any:
169 + explicit_base = cls._base_arg_type(explicit_type)
170 + target_types = extract_types_from_schema(schema)
171 + if isinstance(schema, dict):
172 + if explicit_base in target_types and len(set(target_types)) > 1:
173 + target_types = [explicit_base]
174 + elif explicit_base == "any":
175 + return value
176 + elif explicit_base:
177 + target_types = [explicit_base]
178 +
179 + raw_value = value if "string" in target_types else value.strip()
180 + return coerce_to_schema_type(raw_value, target_types)
181 +
182 + @staticmethod
183 + def _named_tool(request: ToolParserRequest) -> str | None:
184 + tool_choice = request.tool_choice
185 + if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam):
186 + return tool_choice.function.name
187 + if isinstance(tool_choice, ToolChoiceFunction):
188 + return tool_choice.name
189 + return None
190 +
191 + @classmethod
192 + def _validate_tool_name(
193 + cls,
194 + name: Any,
195 + request: ToolParserRequest,
196 + ) -> str:
197 + if (
198 + not isinstance(name, str)
199 + or not name
200 + or any(char.isspace() for char in name)
201 + ):
202 + raise ValueError("IFM tool call has an invalid function name.")
203 + if request.tools and not find_tool_name(request.tools, name):
204 + raise ValueError(f"Unknown IFM tool name: {name}.")
205 + named_tool = cls._named_tool(request)
206 + if named_tool is not None and name != named_tool:
207 + raise ValueError(f"Unexpected IFM tool name: {name}.")
208 + return name
209 +
210 + @classmethod
211 + def _parse_json_call(
212 + cls,
213 + body: str,
214 + request: ToolParserRequest,
215 + ) -> tuple[str, dict[str, Any]]:
216 + raw_call = json.loads(body.strip())
217 + if not isinstance(raw_call, dict):
218 + raise ValueError("IFM JSON tool call must be an object.")
219 +
220 + name = cls._validate_tool_name(raw_call.get("name"), request)
221 + arguments = raw_call.get("arguments", {})
222 + if not isinstance(arguments, dict):
223 + raise ValueError("IFM JSON tool-call arguments must be an object.")
224 +
225 + properties = find_tool_properties(request.tools, name)
226 + arguments = {
227 + arg_name: cls._coerce_json_value(arg_value, properties.get(arg_name))
228 + for arg_name, arg_value in arguments.items()
229 + }
230 + return name, arguments
231 +
232 + @classmethod
233 + def _parse_xml_call(
234 + cls,
235 + body: str,
236 + request: ToolParserRequest,
237 + tool_format: str,
238 + ) -> tuple[str, dict[str, Any]]:
239 + first_arg = body.find(cls._ARG_KEY_START)
240 + if first_arg == -1:
241 + if any(marker in body for marker in cls._ARG_MARKERS):
242 + raise ValueError("Malformed IFM XML argument tags.")
243 + name = cls._validate_tool_name(body.strip(), request)
244 + return name, {}
245 +
246 + name = cls._validate_tool_name(body[:first_arg].strip(), request)
247 + properties = find_tool_properties(request.tools, name)
248 + arguments: dict[str, Any] = {}
249 + position = first_arg
250 + for match in cls._ARG_RE.finditer(body, first_arg):
251 + if body[position : match.start()].strip():
252 + raise ValueError("Malformed IFM XML argument tags.")
253 +
254 + arg_name = match.group(1).strip()
255 + if not arg_name:
256 + raise ValueError("IFM XML argument is missing a name.")
257 + if arg_name in arguments:
258 + raise ValueError(f"Duplicate IFM XML argument: {arg_name}.")
259 +
260 + explicit_type = match.group(2)
261 + if explicit_type is not None:
262 + explicit_type = explicit_type.strip()
263 + if tool_format == "xml" and explicit_type is not None:
264 + raise ValueError("Unexpected type tag in IFM XML tool call.")
265 + if tool_format == "xml_typed" and not explicit_type:
266 + raise ValueError("IFM typed XML argument is missing a type.")
267 + arguments[arg_name] = cls._coerce_xml_value(
268 + match.group(3),
269 + properties.get(arg_name),
270 + explicit_type,
271 + )
272 + position = match.end()
273 +
274 + if position == first_arg or body[position:].strip():
275 + raise ValueError("Malformed IFM XML argument tags.")
276 + return name, arguments
277 +
278 + @classmethod
279 + def _parse_call(
280 + cls,
281 + body: str,
282 + request: ToolParserRequest,
283 + tool_format: str,
284 + ) -> tuple[str, dict[str, Any]]:
285 + if tool_format == "json":
286 + return cls._parse_json_call(body, request)
287 + return cls._parse_xml_call(body, request, tool_format)
288 +
289 + def extract_tool_calls(
290 + self,
291 + model_output: str,
292 + request: ToolParserRequest,
293 + ) -> ExtractedToolCallInformation:
294 + tool_format = self._request_tool_format(request)
295 + if request.tool_choice == "none":
296 + return self._no_tool_calls(model_output)
297 +
298 + group_start = model_output.find(self._TOOL_CALLS_START)
299 + if group_start == -1:
300 + return self._no_tool_calls(model_output)
301 +
302 + body_start = group_start + len(self._TOOL_CALLS_START)
303 + group_end = model_output.find(self._TOOL_CALLS_END, body_start)
304 + if group_end == -1:
305 + return self._no_tool_calls(model_output)
306 +
307 + group_body = model_output[body_start:group_end]
308 + matches = list(self._TOOL_CALL_RE.finditer(group_body))
309 + if not matches or self._TOOL_CALL_RE.sub("", group_body).strip():
310 + logger.warning("Malformed K2 Horizon IFM tool-call group.")
311 + return self._no_tool_calls(model_output)
312 +
313 + try:
314 + tool_calls = [
315 + self._make_tool_call(
316 + *self._parse_call(match.group(1), request, tool_format)
317 + )
318 + for match in matches
319 + ]
320 + except (json.JSONDecodeError, TypeError, ValueError):
321 + logger.warning("Failed to parse K2 Horizon IFM tool calls.", exc_info=True)
322 + return self._no_tool_calls(model_output)
323 +
324 + suffix = model_output[group_end + len(self._TOOL_CALLS_END) :]
325 + content = model_output[:group_start] + suffix
326 + return ExtractedToolCallInformation(
327 + tools_called=True,
328 + tool_calls=tool_calls,
329 + content=self._content_or_none(content),
330 + )
331 +
332 + def _emit_content(self, addition: str, *, final: bool = False) -> str | None:
333 + if addition:
334 + if self._content_seen:
335 + return addition
336 + self._pending_content += addition
337 + if self._pending_content.strip():
338 + self._content_seen = True
339 + content = self._pending_content
340 + self._pending_content = ""
341 + return content
342 + if final:
343 + self._pending_content = ""
344 + return None
345 +
346 + def _streaming_tool_calls(
347 + self,
348 + tool_calls: list[ToolCall],
349 + ) -> list[DeltaToolCall]:
350 + deltas: list[DeltaToolCall] = []
351 + for index, tool_call in enumerate(tool_calls):
352 + arguments_json = tool_call.function.arguments
353 + arguments = json.loads(arguments_json)
354 + self.prev_tool_call_arr.append(
355 + {"name": tool_call.function.name, "arguments": arguments}
356 + )
357 + self.streamed_args_for_tool.append(arguments_json)
358 + deltas.append(
359 + DeltaToolCall(
360 + index=index,
361 + id=tool_call.id,
362 + type="function",
363 + function=DeltaFunctionCall(
364 + name=tool_call.function.name,
365 + arguments=arguments_json,
366 + ),
367 + )
368 + )
369 + return deltas
370 +
371 + def extract_tool_calls_streaming(
372 + self,
373 + previous_text: str,
374 + current_text: str,
375 + delta_text: str,
376 + previous_token_ids: Sequence[int],
377 + current_token_ids: Sequence[int],
378 + delta_token_ids: Sequence[int],
379 + request: ToolParserRequest,
380 + ) -> DeltaMessage | None:
381 + del (
382 + previous_text,
383 + current_text,
384 + previous_token_ids,
385 + current_token_ids,
386 + delta_token_ids,
387 + )
388 +
389 + self._request_tool_format(request)
390 + if request.tool_choice == "none":
391 + emitted_content = self._emit_content(delta_text)
392 + return (
393 + DeltaMessage(content=emitted_content)
394 + if emitted_content is not None
395 + else None
396 + )
397 +
398 + self._stream_text += delta_text
399 + if self._stream_resolved:
400 + emitted_content = self._emit_content(delta_text)
401 + return (
402 + DeltaMessage(content=emitted_content)
403 + if emitted_content is not None
404 + else None
405 + )
406 +
407 + content: str | None = None
408 + if self._stream_group_start is None:
409 + group_start = self._stream_text.find(
410 + self._TOOL_CALLS_START, self._stream_cursor
411 + )
412 + if group_start == -1:
413 + pending = self._stream_text[self._stream_cursor :]
414 + overlap = partial_tag_overlap(pending, self._TOOL_CALLS_START)
415 + sendable_end = len(self._stream_text) - overlap
416 + content = self._emit_content(
417 + self._stream_text[self._stream_cursor : sendable_end]
418 + )
419 + self._stream_cursor = sendable_end
420 + return DeltaMessage(content=content) if content is not None else None
421 +
422 + self._stream_group_start = group_start
423 + content = self._emit_content(
424 + self._stream_text[self._stream_cursor : group_start]
425 + )
426 + self._stream_cursor = group_start
427 +
428 + group_start = self._stream_group_start
429 + body_start = group_start + len(self._TOOL_CALLS_START)
430 + group_end = self._stream_text.find(self._TOOL_CALLS_END, body_start)
431 + if group_end == -1:
432 + return DeltaMessage(content=content) if content is not None else None
433 +
434 + group_end += len(self._TOOL_CALLS_END)
435 + group_text = self._stream_text[group_start:group_end]
436 + parsed = self.extract_tool_calls(group_text, request)
437 + self._stream_resolved = True
438 +
439 + if not parsed.tools_called:
440 + failed_content = self._emit_content(
441 + self._stream_text[self._stream_cursor :]
442 + )
443 + self._stream_cursor = len(self._stream_text)
444 + combined = (content or "") + (failed_content or "")
445 + return DeltaMessage(content=combined) if combined else None
446 +
447 + tool_calls = self._streaming_tool_calls(parsed.tool_calls)
448 + suffix = self._emit_content(self._stream_text[group_end:])
449 + self._stream_cursor = len(self._stream_text)
450 + combined = (content or "") + (suffix or "")
451 + return DeltaMessage(content=combined or None, tool_calls=tool_calls)
452 +
453 + def finish_streaming(self) -> DeltaMessage | None:
454 + if not self._stream_resolved:
455 + content = self._emit_content(
456 + self._stream_text[self._stream_cursor :], final=True
457 + )
458 + self._stream_cursor = len(self._stream_text)
459 + self._stream_resolved = True
460 + else:
461 + content = self._emit_content("", final=True)
462 + return DeltaMessage(content=content) if content is not None else None