diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 03c0b8e307..de896a6fc8 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -414,6 +414,7 @@ th {
| `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
| `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
+| `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ |
| `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
| `Lfm2ForCausalLM` | LFM2 | `LiquidAI/LFM2-1.2B`, `LiquidAI/LFM2-700M`, `LiquidAI/LFM2-350M`, etc. | ✅︎ | ✅︎ |
| `Lfm2MoeForCausalLM` | LFM2MoE | `LiquidAI/LFM2-8B-A1B-preview`, etc. | ✅︎ | ✅︎ |
diff --git a/tests/models/registry.py b/tests/models/registry.py
index 513c9a6760..035e7a1c40 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -364,6 +364,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"random": "ai21labs/Jamba-tiny-random",
},
),
+ "K2HorizonForCausalLM": _HfExamplesInfo(
+ "IFM/K2-Horizon-36B",
+ trust_remote_code=True,
+ is_available_online=False,
+ ),
"KimiLinearForCausalLM": _HfExamplesInfo(
"moonshotai/Kimi-Linear-48B-A3B-Instruct", trust_remote_code=True
),
diff --git a/tests/reasoning/test_k2_horizon_reasoning_parser.py b/tests/reasoning/test_k2_horizon_reasoning_parser.py
new file mode 100644
index 0000000000..1a554770f1
--- /dev/null
+++ b/tests/reasoning/test_k2_horizon_reasoning_parser.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import string
+from collections.abc import Sequence
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser.parser_manager import ParserManager
+from vllm.reasoning import ReasoningParserManager
+from vllm.reasoning.k2_horizon_reasoning_parser import K2HorizonReasoningParser
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+EFFORT_TOKENS = {
+ "high": ("", ""),
+ "medium": ("", ""),
+ "low": ("", ""),
+}
+TOOL_CALLS_START = ""
+
+
+class K2Tokenizer:
+ def __init__(self) -> None:
+ special_tokens = [
+ token for token_pair in EFFORT_TOKENS.values() for token in token_pair
+ ] + [TOOL_CALLS_START]
+ tokens = special_tokens + list(string.printable)
+ self._vocab = {token: index for index, token in enumerate(tokens, start=1)}
+
+ def get_vocab(self) -> dict[str, int]:
+ return dict(self._vocab)
+
+ def encode(self, text: str) -> list[int]:
+ tokens: list[int] = []
+ position = 0
+ specials = sorted(EFFORT_TOKENS["high"] + (TOOL_CALLS_START,), key=len)
+ while position < len(text):
+ for special in specials:
+ if text.startswith(special, position):
+ tokens.append(self._vocab[special])
+ position += len(special)
+ break
+ else:
+ tokens.append(self._vocab[text[position]])
+ position += 1
+ return tokens
+
+ def decode(
+ self,
+ ids: Sequence[int] | int,
+ skip_special_tokens: bool = False,
+ ) -> str:
+ del skip_special_tokens
+ if isinstance(ids, int):
+ ids = [ids]
+ inverse_vocab = {token_id: token for token, token_id in self._vocab.items()}
+ return "".join(inverse_vocab[token_id] for token_id in ids)
+
+
+def _request(**kwargs) -> ChatCompletionRequest:
+ return ChatCompletionRequest(
+ model="k2-horizon",
+ messages=[{"role": "user", "content": "test"}],
+ **kwargs,
+ )
+
+
+@pytest.mark.parametrize("effort", ["high", "medium", "low"])
+def test_reasoning_effort_selects_tokens(effort: str):
+ parser = K2HorizonReasoningParser(
+ K2Tokenizer(),
+ chat_template_kwargs={"reasoning_effort": effort},
+ )
+
+ assert (parser.start_token, parser.end_token) == EFFORT_TOKENS[effort]
+
+
+def test_reasoning_effort_precedence_and_default():
+ tokenizer = K2Tokenizer()
+ request = _request(
+ reasoning_effort="medium",
+ chat_template_kwargs={"reasoning_effort": "low"},
+ )
+ resolved_kwargs = (
+ request.build_chat_params(None, "auto")
+ .with_defaults({"reasoning_effort": "high"})
+ .chat_template_kwargs
+ )
+
+ parser = K2HorizonReasoningParser(
+ tokenizer,
+ chat_template_kwargs=resolved_kwargs,
+ )
+ default_parser = K2HorizonReasoningParser(tokenizer)
+
+ assert parser.start_token == EFFORT_TOKENS["medium"][0]
+ assert default_parser.start_token == EFFORT_TOKENS["high"][0]
+
+
+@pytest.mark.parametrize("effort", ["none", "unknown", None])
+def test_invalid_reasoning_effort_raises(effort: str | None):
+ with pytest.raises(ValueError, match="Unsupported reasoning_effort"):
+ K2HorizonReasoningParser(
+ K2Tokenizer(),
+ chat_template_kwargs={"reasoning_effort": effort},
+ )
+
+
+def test_non_streaming_explicit_and_implicit_boundaries():
+ parser = K2HorizonReasoningParser(K2Tokenizer())
+ request = _request()
+
+ assert parser.extract_reasoning("plananswer", request) == (
+ "plan",
+ "answer",
+ )
+ assert parser.extract_reasoning(
+ f"plan{TOOL_CALLS_START}ping", request
+ ) == (
+ "plan",
+ f"{TOOL_CALLS_START}ping",
+ )
+
+
+@pytest.mark.parametrize("reasoning", ["plan", " ", "\n", " \n"])
+def test_character_split_reasoning_stream(reasoning: str):
+ parser = K2HorizonReasoningParser(K2Tokenizer())
+ reasoning_parts: list[str] = []
+ content_parts: list[str] = []
+
+ for char in f"{reasoning}answer":
+ delta = parser.extract_reasoning_streaming(
+ previous_text="",
+ current_text="",
+ delta_text=char,
+ previous_token_ids=[],
+ current_token_ids=[],
+ delta_token_ids=[],
+ )
+ if delta is not None:
+ if delta.reasoning is not None:
+ reasoning_parts.append(delta.reasoning)
+ if delta.content is not None:
+ content_parts.append(delta.content)
+
+ assert "".join(reasoning_parts) == reasoning
+ assert "".join(content_parts) == "answer"
+
+
+def test_latest_boundary_controls_reasoning_state():
+ tokenizer = K2Tokenizer()
+ parser = K2HorizonReasoningParser(tokenizer)
+ vocab = tokenizer.get_vocab()
+
+ assert parser.is_reasoning_end([vocab[""], vocab[TOOL_CALLS_START]])
+ assert not parser.is_reasoning_end([vocab[TOOL_CALLS_START], vocab[""]])
+
+
+def test_composed_streaming_reasoning_to_tool_handoff():
+ tokenizer = K2Tokenizer()
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "ping",
+ "parameters": {"type": "object", "properties": {}},
+ },
+ }
+ ]
+ request = _request(tools=tools, tool_choice="auto")
+ parser_cls = ParserManager.get_parser(
+ tool_parser_name="k2_horizon",
+ reasoning_parser_name="k2_horizon",
+ enable_auto_tools=True,
+ )
+ assert parser_cls is not None
+ parser = parser_cls(
+ tokenizer,
+ request.tools,
+ chat_template_kwargs={"reasoning_effort": "high"},
+ )
+ output = "planping"
+ reasoning_parts: list[str] = []
+ tool_names: list[str] = []
+
+ for index, char in enumerate(output):
+ delta = parser.parse_delta(
+ delta_text=char,
+ delta_token_ids=tokenizer.encode(char),
+ request=request,
+ prompt_token_ids=[] if index == 0 else None,
+ finished=index == len(output) - 1,
+ )
+ if delta is not None:
+ if delta.reasoning is not None:
+ reasoning_parts.append(delta.reasoning)
+ tool_names.extend(
+ tool_call.function.name
+ for tool_call in delta.tool_calls
+ if tool_call.function is not None
+ and tool_call.function.name is not None
+ )
+
+ assert "".join(reasoning_parts) == "plan"
+ assert tool_names == ["ping"]
+
+
+def test_reasoning_parser_registered():
+ assert (
+ ReasoningParserManager.get_reasoning_parser("k2_horizon")
+ is K2HorizonReasoningParser
+ )
diff --git a/tests/tool_parsers/test_k2_horizon_tool_parser.py b/tests/tool_parsers/test_k2_horizon_tool_parser.py
new file mode 100644
index 0000000000..6fd4844aeb
--- /dev/null
+++ b/tests/tool_parsers/test_k2_horizon_tool_parser.py
@@ -0,0 +1,318 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.parser.parser_manager import ParserManager
+from vllm.tool_parsers import ToolParserManager
+from vllm.tool_parsers.k2_horizon_tool_parser import K2HorizonToolParser
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+TOOLS = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "query": {"type": "string"},
+ "limit": {"type": "integer"},
+ "enabled": {"type": "boolean"},
+ },
+ },
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "ping",
+ "parameters": {"type": "object", "properties": {}},
+ },
+ },
+]
+
+
+@pytest.fixture
+def tokenizer():
+ tokenizer = MagicMock()
+ tokenizer.get_vocab.return_value = {}
+ return tokenizer
+
+
+def _request(
+ tool_format: str | None = None,
+ tool_choice="auto",
+) -> ChatCompletionRequest:
+ chat_template_kwargs = (
+ {"tool_call_format": tool_format} if tool_format is not None else None
+ )
+ return ChatCompletionRequest(
+ model="k2-horizon",
+ messages=[{"role": "user", "content": "test"}],
+ tools=TOOLS,
+ tool_choice=tool_choice,
+ chat_template_kwargs=chat_template_kwargs,
+ )
+
+
+def _group(*calls: str, prefix: str = "", suffix: str = "") -> str:
+ return prefix + "" + "\n".join(calls) + "" + suffix
+
+
+def _json_call(name: str, arguments: dict) -> str:
+ payload = json.dumps({"name": name, "arguments": arguments})
+ return f"{payload}"
+
+
+def _xml_call(
+ name: str,
+ arguments: list[tuple[str, str]],
+ *,
+ typed: bool = False,
+) -> str:
+ parts = [f"{name}"]
+ for key, value in arguments:
+ parts.append(f"{key}")
+ if typed:
+ arg_type = "integer" if key == "limit" else "string"
+ parts.append(f"{arg_type}")
+ parts.append(f"{value}")
+ parts.append("")
+ return "".join(parts)
+
+
+def _collect_stream(
+ parser: K2HorizonToolParser,
+ request: ChatCompletionRequest,
+ output: str,
+) -> tuple[str, list[tuple[str, dict]]]:
+ messages: list[DeltaMessage | None] = []
+ for char in output:
+ messages.append(
+ parser.extract_tool_calls_streaming(
+ previous_text="",
+ current_text="",
+ delta_text=char,
+ previous_token_ids=[],
+ current_token_ids=[],
+ delta_token_ids=[],
+ request=request,
+ )
+ )
+ messages.append(parser.finish_streaming())
+
+ content: list[str] = []
+ calls: list[tuple[str, dict]] = []
+ for message in messages:
+ if message is None:
+ continue
+ if message.content is not None:
+ content.append(message.content)
+ for tool_call in message.tool_calls:
+ assert tool_call.function is not None
+ assert tool_call.function.name is not None
+ assert tool_call.function.arguments is not None
+ calls.append(
+ (
+ tool_call.function.name,
+ json.loads(tool_call.function.arguments),
+ )
+ )
+ return "".join(content), calls
+
+
+def test_xml_default_parses_multiple_calls_and_coerces_schema(tokenizer):
+ parser = K2HorizonToolParser(tokenizer, _request().tools)
+ output = _group(
+ _xml_call("lookup", [("query", "123"), ("limit", "3")]),
+ _xml_call("ping", []),
+ prefix="Before ",
+ suffix=" after.",
+ )
+
+ result = parser.extract_tool_calls(output, _request())
+
+ assert parser.tool_format == "xml"
+ assert result.tools_called
+ assert result.content == "Before after."
+ assert [call.function.name for call in result.tool_calls] == ["lookup", "ping"]
+ assert json.loads(result.tool_calls[0].function.arguments) == {
+ "query": "123",
+ "limit": 3,
+ }
+
+
+@pytest.mark.parametrize(
+ ("tool_format", "call", "expected"),
+ [
+ ("json", _json_call("lookup", {"limit": "2"}), {"limit": 2}),
+ (
+ "xml_typed",
+ _xml_call("lookup", [("query", "weather"), ("limit", "2")], typed=True),
+ {"query": "weather", "limit": 2},
+ ),
+ ],
+)
+def test_request_level_format_overrides(tokenizer, tool_format, call, expected):
+ request = _request(tool_format)
+ parser = K2HorizonToolParser(tokenizer, request.tools)
+
+ result = parser.extract_tool_calls(_group(call), request)
+
+ assert result.tools_called
+ assert json.loads(result.tool_calls[0].function.arguments) == expected
+
+
+@pytest.mark.parametrize(
+ "tool_format",
+ ["yaml", "python", "xml_untyped", "", None, 1],
+)
+def test_invalid_request_format_rejected_before_generation(tokenizer, tool_format):
+ parser = K2HorizonToolParser(tokenizer, _request().tools)
+ request = _request().model_copy(
+ update={"chat_template_kwargs": {"tool_call_format": tool_format}}
+ )
+
+ with pytest.raises(ValueError, match="Unsupported tool_call_format"):
+ parser.adjust_request(request)
+
+
+@pytest.mark.parametrize(
+ ("tool_request", "call"),
+ [
+ (_request(), _json_call("ping", {})),
+ (_request("json"), _xml_call("ping", [])),
+ (_request("xml"), _xml_call("lookup", [("query", "x")], typed=True)),
+ (_request("xml_typed"), _xml_call("lookup", [("query", "x")])),
+ ],
+)
+def test_mismatched_formats_fail_closed(tokenizer, tool_request, call):
+ parser = K2HorizonToolParser(tokenizer, tool_request.tools)
+ output = _group(call)
+
+ result = parser.extract_tool_calls(output, tool_request)
+
+ assert not result.tools_called
+ assert result.content == output
+
+
+@pytest.mark.parametrize(
+ "tool_choice",
+ [
+ "auto",
+ "required",
+ {"type": "function", "function": {"name": "lookup"}},
+ ],
+)
+def test_auto_required_and_named_use_ifm_parser(tokenizer, tool_choice):
+ request = _request("json", tool_choice)
+ parser_cls = ParserManager.get_parser(
+ tool_parser_name="k2_horizon",
+ enable_auto_tools=True,
+ )
+ assert parser_cls is not None
+ parser = parser_cls(tokenizer, request.tools)
+
+ assert parser.adjust_request(request) is request
+ assert request.structured_outputs is None
+ _, content, tool_calls = parser.parse(
+ _group(_json_call("lookup", {"enabled": "true"})),
+ request,
+ enable_auto_tools=True,
+ )
+
+ assert content is None
+ assert tool_calls is not None
+ assert json.loads(tool_calls[0].arguments) == {"enabled": True}
+
+
+def test_unknown_or_wrong_named_tool_fails_closed(tokenizer):
+ named_request = _request(
+ "json", {"type": "function", "function": {"name": "lookup"}}
+ )
+ parser = K2HorizonToolParser(tokenizer, named_request.tools)
+
+ unknown = parser.extract_tool_calls(
+ _group(_json_call("missing", {})), _request("json")
+ )
+ wrong_named = parser.extract_tool_calls(
+ _group(_json_call("ping", {})), named_request
+ )
+
+ assert not unknown.tools_called
+ assert not wrong_named.tools_called
+
+
+@pytest.mark.parametrize(
+ "output",
+ [
+ "ping",
+ "ping",
+ _group("junk"),
+ ],
+)
+def test_malformed_and_incomplete_input_is_content(tokenizer, output):
+ request = _request()
+ parser = K2HorizonToolParser(tokenizer, request.tools)
+
+ result = parser.extract_tool_calls(output, request)
+
+ assert not result.tools_called
+ assert result.content == output
+
+
+@pytest.mark.parametrize(
+ ("tool_request", "call", "expected"),
+ [
+ (_request(), _xml_call("lookup", [("limit", "2")]), {"limit": 2}),
+ (_request("json"), _json_call("lookup", {"limit": "2"}), {"limit": 2}),
+ (
+ _request("xml_typed"),
+ _xml_call("lookup", [("limit", "2")], typed=True),
+ {"limit": 2},
+ ),
+ ],
+)
+def test_character_split_streaming(tokenizer, tool_request, call, expected):
+ parser = K2HorizonToolParser(tokenizer, tool_request.tools)
+ output = _group(call, prefix="Before", suffix="after")
+
+ content, calls = _collect_stream(parser, tool_request, output)
+
+ assert content == "Beforeafter"
+ assert calls == [("lookup", expected)]
+
+
+def test_streaming_flushes_incomplete_markup_as_content(tokenizer):
+ request = _request()
+ parser = K2HorizonToolParser(tokenizer, request.tools)
+ output = "Beforeping"
+
+ content, calls = _collect_stream(parser, request, output)
+
+ assert content == output
+ assert calls == []
+
+
+def test_whitespace_only_surrounding_content_is_not_preserved(tokenizer):
+ request = _request()
+ parser = K2HorizonToolParser(tokenizer, request.tools)
+
+ result = parser.extract_tool_calls(
+ _group(_xml_call("ping", []), prefix=" \n", suffix="\t "), request
+ )
+
+ assert result.tools_called
+ assert result.content is None
+
+
+def test_tool_parser_registered():
+ assert ToolParserManager.get_tool_parser("k2_horizon") is K2HorizonToolParser
+ assert K2HorizonToolParser.supports_required_and_named is False
diff --git a/vllm/model_executor/models/k2_horizon.py b/vllm/model_executor/models/k2_horizon.py
new file mode 100644
index 0000000000..47e87c0dcb
--- /dev/null
+++ b/vllm/model_executor/models/k2_horizon.py
@@ -0,0 +1,1445 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 MBZUAI-IFM.
+# Copyright 2024 The Qwen team.
+# Copyright 2023 The vLLM team.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only K2Horizon model compatible with HuggingFace weights."""
+
+import math
+import typing
+from collections.abc import Callable, Iterable
+from itertools import islice
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm import _custom_ops as ops
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import (
+ get_ep_group,
+ get_pp_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+ tensor_model_parallel_all_gather,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.fused_moe import (
+ FusedMoEFactory,
+ fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+ _get_config_dtype_str,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+ _get_config_quant_dtype,
+ _prepare_expert_assignment,
+ dispatch_fused_moe_kernel,
+ try_get_optimal_moe_config,
+)
+from vllm.model_executor.layers.fused_moe.utils import (
+ moe_kernel_quantize_input,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ MergedColumnParallelLinear,
+ QKVParallelLinear,
+ ReplicatedLinear,
+ RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+ ParallelLMHead,
+ VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+ default_weight_loader,
+ maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.interfaces import (
+ EagleModelMixin,
+ MixtureOfExperts,
+ SupportsEagle,
+ SupportsEagle3,
+ SupportsLoRA,
+ SupportsPP,
+)
+from vllm.model_executor.models.utils import (
+ AutoWeightsLoader,
+ PPMissingLayer,
+ extract_layer_index,
+ is_pp_missing_parameter,
+ make_empty_intermediate_tensors_factory,
+ make_layers,
+ maybe_prefix,
+ sequence_parallel_chunk,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.triton_utils import tl
+
+logger = init_logger(__name__)
+
+
+def split_to_interleaved(x):
+ # Split halves: x0 x1 x2 x3 ... y0 y1 y2 y3 ...
+ # Interleaved: x0 y0 x1 y1 x2 y2 x3 y3 ...
+ return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1)
+
+
+def interleaved_to_split(x):
+ # Interleaved: x0 y0 x1 y1 x2 y2 x3 y3 ...
+ # Split halves: x0 x1 x2 x3 ... y0 y1 y2 y3 ...
+ return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1)
+
+
+def calc_router_weights(
+ router_logits: torch.Tensor,
+ e_score_correction_bias: torch.Tensor | None,
+ score_func: str,
+ top_k: int,
+ scaling_factor: float | None,
+ renormalize: bool | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if score_func == "softmax":
+ routing_scores = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+ else:
+ assert score_func == "sigmoid", (
+ f"Unsupported router score function: {score_func}"
+ )
+ routing_scores = torch.sigmoid(router_logits.to(torch.float32))
+
+ selection_scores = routing_scores
+ if e_score_correction_bias is not None:
+ selection_scores = selection_scores + e_score_correction_bias.to(
+ selection_scores
+ )
+
+ selected_indices = torch.topk(selection_scores, top_k, dim=-1).indices
+ routing_weights = torch.gather(routing_scores, dim=-1, index=selected_indices)
+ if renormalize is None:
+ renormalize = top_k > 1
+ if renormalize:
+ routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
+ if scaling_factor is not None:
+ routing_weights = routing_weights * scaling_factor
+ return routing_weights.to(router_logits.dtype), selected_indices
+
+
+def fused_mova_impl(
+ config,
+ hidden_states: torch.Tensor,
+ w1: torch.Tensor,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ use_fp8_w8a8: bool = False,
+ use_int8_w8a8: bool = False,
+ use_int8_w8a16: bool = False,
+ use_int4_w4a16: bool = False,
+ ocp_mx_scheme: str | None = None,
+ per_channel_quant: bool = False,
+ global_num_experts: int = -1,
+ expert_map: torch.Tensor | None = None,
+ w1_scale: torch.Tensor | None = None,
+ w1_zp: torch.Tensor | None = None,
+ a1_scale: torch.Tensor | None = None,
+ block_shape: list[int] | None = None,
+ w1_bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+ if ocp_mx_scheme is not None:
+ raise NotImplementedError(
+ f"Using ocp_mx_scheme={ocp_mx_scheme} in functional fused_experts call is "
+ "deprecated. Please use OCP_MXQuantizationEmulationTritonExperts."
+ )
+
+ # Check constraints.
+ if use_int4_w4a16:
+ assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
+ else:
+ assert hidden_states.size(1) == w1.size(2), (
+ f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}"
+ )
+
+ assert topk_weights.size() == topk_ids.size(), "topk shape mismatch"
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
+ assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16]
+
+ num_tokens = hidden_states.size(0)
+ E, N, K = w1.size()
+
+ if global_num_experts == -1:
+ global_num_experts = E
+ top_k_num = topk_ids.size(1)
+ M = num_tokens
+
+ # Note: for use_int8_w8a16 or use_int4_w4a16, the activations are
+ # quantized prior to calling fused_experts.
+ quant_dtype = _get_config_quant_dtype(
+ use_fp8_w8a8=use_fp8_w8a8,
+ use_int8_w8a8=use_int8_w8a8,
+ )
+
+ intermediate_cache1 = torch.empty(
+ (M, top_k_num, N), device=hidden_states.device, dtype=hidden_states.dtype
+ )
+
+ if hidden_states.dtype == torch.bfloat16:
+ compute_type = tl.bfloat16
+ elif hidden_states.dtype == torch.float16:
+ compute_type = tl.float16
+ elif hidden_states.dtype == torch.float32:
+ compute_type = tl.float32
+ else:
+ raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
+
+ out_hidden_states = torch.empty(
+ (M, N),
+ device=hidden_states.device,
+ dtype=hidden_states.dtype,
+ )
+
+ qhidden_states, a1q_scale = moe_kernel_quantize_input(
+ A=hidden_states,
+ A_scale=a1_scale,
+ quant_dtype=quant_dtype,
+ per_act_token_quant=per_channel_quant,
+ block_shape=block_shape,
+ )
+
+ sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
+ topk_ids,
+ config,
+ num_tokens,
+ top_k_num,
+ global_num_experts,
+ expert_map,
+ use_int8_w8a16=use_int8_w8a16,
+ use_int4_w4a16=use_int4_w4a16,
+ block_shape=block_shape,
+ ignore_invalid_experts=True,
+ )
+
+ dispatch_fused_moe_kernel(
+ qhidden_states,
+ w1,
+ intermediate_cache1,
+ a1q_scale,
+ w1_scale,
+ w1_zp,
+ topk_weights,
+ sorted_token_ids,
+ expert_ids,
+ num_tokens_post_padded,
+ False, # apply_router_weight_on_input,
+ top_k_num,
+ config,
+ compute_type=compute_type,
+ use_fp8_w8a8=use_fp8_w8a8,
+ use_int8_w8a8=use_int8_w8a8,
+ use_int8_w8a16=use_int8_w8a16,
+ use_int4_w4a16=use_int4_w4a16,
+ per_channel_quant=per_channel_quant,
+ block_shape=block_shape,
+ B_bias=w1_bias,
+ )
+
+ intermediate_cache1 = F.silu(intermediate_cache1)
+ intermediate_cache1 *= topk_weights[:, :, None].to(intermediate_cache1.dtype)
+
+ ops.moe_sum(intermediate_cache1, out_hidden_states)
+
+ return out_hidden_states
+
+
+class K2HorizonRMSNorm(RMSNorm):
+ def __init__(self, hidden_size: int, n_groups: int, eps: float = 1e-6):
+ assert hidden_size % n_groups == 0
+ super().__init__(hidden_size, eps=eps)
+ self.hidden_size = hidden_size
+ self.n_groups = n_groups
+ self.group_size = hidden_size // n_groups
+
+ def forward(self, x, residual=None):
+ if self.n_groups == 1:
+ return super().forward(x, residual)
+
+ if residual is not None:
+ x = x + residual
+ residual = x
+
+ orig_shape = x.shape
+ x_grouped = x.reshape(-1, self.n_groups, self.group_size).reshape(
+ -1, self.group_size
+ )
+ y_grouped = torch.empty_like(x_grouped)
+
+ # Reuse vLLM's RMSNorm CUDA kernel on each group.
+ ops.rms_norm(
+ y_grouped,
+ x_grouped,
+ torch.ones(self.group_size, dtype=x.dtype, device=x.device),
+ self.variance_epsilon,
+ )
+
+ y = y_grouped.reshape(orig_shape) * self.weight
+ return (y, residual) if residual is not None else y
+
+
+class K2HorizonMLP(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ hidden_act: str,
+ quant_config: QuantizationConfig | None = None,
+ reduce_results: bool = True,
+ expert_gate: torch.nn.Linear | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.gate_up_proj = MergedColumnParallelLinear(
+ hidden_size,
+ [intermediate_size] * 2,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate_up_proj",
+ )
+ self.down_proj = RowParallelLinear(
+ intermediate_size,
+ hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ reduce_results=reduce_results,
+ prefix=f"{prefix}.down_proj",
+ )
+ if hidden_act != "silu":
+ raise ValueError(
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
+ )
+ self.act_fn = SiluAndMul()
+ self.expert_gate = expert_gate
+
+ def forward(self, x):
+ gate_up, _ = self.gate_up_proj(x)
+ out = self.act_fn(gate_up)
+ out, _ = self.down_proj(out)
+
+ if self.expert_gate is not None:
+ out = F.sigmoid(self.expert_gate(x)[0]) * out
+
+ return out
+
+
+class K2HorizonSparseMoeBlock(nn.Module):
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ prefix: str = "",
+ ):
+ super().__init__()
+
+ config = vllm_config.model_config.hf_text_config
+ parallel_config = vllm_config.parallel_config
+ quant_config = vllm_config.quant_config
+
+ self.tp_size = get_tensor_model_parallel_world_size()
+
+ self.ep_group = get_ep_group().device_group
+ self.ep_rank = get_ep_group().rank_in_group
+ self.ep_size = self.ep_group.size()
+ self.n_routed_experts = config.num_experts
+
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+
+ if self.tp_size > config.num_experts:
+ raise ValueError(
+ f"Tensor parallel size {self.tp_size} is greater than "
+ f"the number of experts {config.num_experts}."
+ )
+
+ # Load balancing settings.
+ vllm_config = get_current_vllm_config()
+ eplb_config = vllm_config.parallel_config.eplb_config
+ self.enable_eplb = parallel_config.enable_eplb
+
+ self.n_logical_experts = self.n_routed_experts
+ self.n_redundant_experts = eplb_config.num_redundant_experts
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
+ self.physical_expert_end = (
+ self.physical_expert_start + self.n_local_physical_experts
+ )
+
+ self.gate = ReplicatedLinear(
+ config.hidden_size,
+ config.num_experts,
+ bias=config.moe_gate_bias,
+ skip_bias_add=True,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate",
+ )
+ if self.gate.bias is not None:
+ self.gate.bias = nn.Parameter(self.gate.bias.float(), requires_grad=False)
+
+ self.num_shared_experts = config.num_shared_experts
+ if config.num_shared_experts > 0:
+ self.shared_experts = K2HorizonMLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.moe_intermediate_size
+ * config.num_shared_experts,
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ reduce_results=False,
+ expert_gate=None,
+ prefix=f"{prefix}.shared_experts",
+ )
+ else:
+ self.shared_experts = None
+
+ self.experts = FusedMoEFactory(
+ shared_experts=self.shared_experts,
+ gate=self.gate,
+ num_experts=self.n_routed_experts,
+ top_k=config.num_experts_per_tok,
+ hidden_size=config.hidden_size,
+ intermediate_size=config.moe_intermediate_size,
+ renormalize=config.norm_topk_prob,
+ quant_config=quant_config,
+ prefix=f"{prefix}.experts",
+ enable_eplb=self.enable_eplb,
+ num_redundant_experts=self.n_redundant_experts,
+ is_sequence_parallel=self.is_sequence_parallel,
+ scoring_func=config.router_score_func,
+ e_score_correction_bias=self.gate.bias,
+ routed_scaling_factor=config.router_scaling_factor,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ assert hidden_states.dim() <= 2, (
+ "K2HorizonSparseMoeBlock only supports 1D or 2D inputs"
+ )
+ is_input_1d = hidden_states.dim() == 1
+ num_tokens, hidden_dim = hidden_states.shape
+ hidden_states = hidden_states.view(-1, hidden_dim)
+
+ if self.is_sequence_parallel:
+ hidden_states = sequence_parallel_chunk(hidden_states)
+
+ final_hidden_states = self.experts(
+ hidden_states=hidden_states, router_logits=hidden_states
+ )
+
+ if self.is_sequence_parallel:
+ final_hidden_states = tensor_model_parallel_all_gather(
+ final_hidden_states, 0
+ )
+ final_hidden_states = final_hidden_states[:num_tokens]
+
+ # return to 1d if input is 1d
+ return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states
+
+
+class K2HorizonAttention(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ num_kv_heads: int,
+ rope_parameters: dict[str, Any],
+ rope_head_dim: int,
+ query_key_norm: bool,
+ max_position_embeddings: int = 8192,
+ head_dim: int | None = None,
+ rms_norm_eps: float = 1e-06,
+ qkv_bias: bool = False,
+ cache_config: CacheConfig | None = None,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ dual_chunk_attention_config: dict[str, Any] | None = None,
+ gate_func: str | None = None,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = hidden_size
+ tp_size = get_tensor_model_parallel_world_size()
+ self.total_num_heads = num_heads
+ assert self.total_num_heads % tp_size == 0
+ self.num_heads = self.total_num_heads // tp_size
+ self.total_num_kv_heads = num_kv_heads
+ if self.total_num_kv_heads >= tp_size:
+ # Number of KV heads is greater than TP size, so we partition
+ # the KV heads across multiple tensor parallel GPUs.
+ assert self.total_num_kv_heads % tp_size == 0
+ else:
+ # Number of KV heads is less than TP size, so we replicate
+ # the KV heads across multiple tensor parallel GPUs.
+ assert tp_size % self.total_num_kv_heads == 0
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+ self.q_size = self.num_heads * self.head_dim
+ self.kv_size = self.num_kv_heads * self.head_dim
+ self.scaling = self.head_dim**-0.5
+ self.max_position_embeddings = max_position_embeddings
+ self.dual_chunk_attention_config = dual_chunk_attention_config
+
+ self.qkv_proj = QKVParallelLinear(
+ hidden_size,
+ self.head_dim,
+ self.total_num_heads,
+ self.total_num_kv_heads,
+ bias=qkv_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.qkv_proj",
+ )
+
+ self.o_proj = RowParallelLinear(
+ self.total_num_heads * self.head_dim,
+ hidden_size,
+ bias=qkv_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.o_proj",
+ )
+
+ self.rope_head_dim = rope_head_dim or self.head_dim
+ self.rotary_emb = get_rope(
+ self.rope_head_dim,
+ max_position=max_position_embeddings,
+ rope_parameters=rope_parameters,
+ dual_chunk_attention_config=dual_chunk_attention_config,
+ )
+ self.attn = Attention(
+ self.num_heads,
+ self.head_dim,
+ self.scaling,
+ num_kv_heads=self.num_kv_heads,
+ cache_config=cache_config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.attn",
+ **{
+ "layer_idx": extract_layer_index(prefix),
+ "dual_chunk_attention_config": dual_chunk_attention_config,
+ }
+ if dual_chunk_attention_config
+ else {},
+ )
+
+ self.query_key_norm = query_key_norm
+ if self.query_key_norm:
+ self.q_norm = K2HorizonRMSNorm(
+ hidden_size=self.q_size, n_groups=self.num_heads, eps=rms_norm_eps
+ )
+ self.k_norm = K2HorizonRMSNorm(
+ hidden_size=self.kv_size, n_groups=self.num_kv_heads, eps=rms_norm_eps
+ )
+
+ self.gate_func = gate_func
+ if self.gate_func is not None:
+ self.gate_proj = ColumnParallelLinear(
+ hidden_size,
+ self.total_num_heads * self.head_dim,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate_proj",
+ )
+
+ self.tp_size = get_tensor_model_parallel_world_size()
+ self.tp_rank = get_tensor_model_parallel_rank()
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ ) -> torch.Tensor:
+ qkv, _ = self.qkv_proj(hidden_states)
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+ if self.query_key_norm:
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ if self.rope_head_dim == self.head_dim:
+ q, k = self.rotary_emb(positions, q, k)
+ else:
+ q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
+ k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
+
+ q_rope, q_nope = torch.split(
+ split_to_interleaved(q),
+ split_size_or_sections=[
+ self.rope_head_dim,
+ self.head_dim - self.rope_head_dim,
+ ],
+ dim=-1,
+ )
+ k_rope, k_nope = torch.split(
+ split_to_interleaved(k),
+ split_size_or_sections=[
+ self.rope_head_dim,
+ self.head_dim - self.rope_head_dim,
+ ],
+ dim=-1,
+ )
+
+ q_rope, k_rope = self.rotary_emb(
+ positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
+ )
+
+ q = interleaved_to_split(
+ torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
+ ).reshape(*q.shape[:-2], -1)
+
+ k = interleaved_to_split(
+ torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
+ ).reshape(*k.shape[:-2], -1)
+
+ attn_output = self.attn(q, k, v)
+
+ if self.gate_func is not None:
+ gate, _ = self.gate_proj(hidden_states)
+ if self.gate_func == "silu":
+ gate = F.silu(gate)
+ else:
+ assert self.gate_func == "softplus"
+ gate = F.softplus(gate, beta=math.log(2))
+
+ attn_output = attn_output * gate
+
+ output, _ = self.o_proj(attn_output)
+ return output
+
+
+class K2HorizonMoVAAttention(nn.Module):
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ hidden_size: int,
+ num_heads: int,
+ num_kv_heads: int,
+ rope_parameters: dict[str, Any],
+ rope_head_dim: int,
+ query_key_norm: bool,
+ num_experts: int,
+ num_experts_per_tok: int,
+ moe_gate_bias: bool,
+ router_score_func: str,
+ router_scaling_factor: float,
+ max_position_embeddings: int = 8192,
+ head_dim: int | None = None,
+ rms_norm_eps: float = 1e-06,
+ qkv_bias: bool = False,
+ cache_config: CacheConfig | None = None,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ dual_chunk_attention_config: dict[str, Any] | None = None,
+ gate_func: str | None = None,
+ ) -> None:
+ super().__init__()
+ self.hidden_size = hidden_size
+ tp_size = get_tensor_model_parallel_world_size()
+ self.total_num_heads = num_heads
+ assert self.total_num_heads % tp_size == 0
+ self.num_heads = self.total_num_heads // tp_size
+ self.total_num_kv_heads = num_kv_heads
+ if self.total_num_kv_heads >= tp_size:
+ # Number of KV heads is greater than TP size, so we partition
+ # the KV heads across multiple tensor parallel GPUs.
+ assert self.total_num_kv_heads % tp_size == 0
+ else:
+ # Number of KV heads is less than TP size, so we replicate
+ # the KV heads across multiple tensor parallel GPUs.
+ assert tp_size % self.total_num_kv_heads == 0
+
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
+ self.q_size = self.num_heads * self.head_dim
+ self.kv_size = self.num_kv_heads * self.head_dim
+ self.scaling = self.head_dim**-0.5
+ self.max_position_embeddings = max_position_embeddings
+ self.dual_chunk_attention_config = dual_chunk_attention_config
+
+ self.num_experts = num_experts
+ self.num_experts_per_tok = num_experts_per_tok
+ self.router_score_func = router_score_func
+ self.router_scaling_factor = router_scaling_factor
+
+ self.qk_proj = MergedColumnParallelLinear(
+ hidden_size,
+ [self.total_num_heads * self.head_dim, self.kv_size * tp_size],
+ bias=qkv_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.qk_proj",
+ )
+
+ self.v_router = ReplicatedLinear(
+ hidden_size,
+ num_experts,
+ bias=moe_gate_bias,
+ skip_bias_add=True,
+ quant_config=quant_config,
+ prefix=f"{prefix}.v_router",
+ )
+ if self.v_router.bias is not None:
+ self.v_router.bias = nn.Parameter(
+ self.v_router.bias.float(), requires_grad=False
+ )
+
+ self.v_experts = nn.ModuleList(
+ [
+ ColumnParallelLinear(
+ hidden_size,
+ self.kv_size * tp_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.v_experts.{value}",
+ )
+ for value in range(num_experts)
+ ]
+ )
+
+ self.o_proj = RowParallelLinear(
+ self.total_num_heads * self.head_dim,
+ hidden_size,
+ bias=qkv_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.o_proj",
+ )
+
+ self.rope_head_dim = rope_head_dim or self.head_dim
+ self.rotary_emb = get_rope(
+ self.rope_head_dim,
+ max_position=max_position_embeddings,
+ rope_parameters=rope_parameters,
+ dual_chunk_attention_config=dual_chunk_attention_config,
+ )
+ self.attn = Attention(
+ self.num_heads,
+ self.head_dim,
+ self.scaling,
+ num_kv_heads=self.num_kv_heads,
+ cache_config=cache_config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.attn",
+ **{
+ "layer_idx": extract_layer_index(prefix),
+ "dual_chunk_attention_config": dual_chunk_attention_config,
+ }
+ if dual_chunk_attention_config
+ else {},
+ )
+
+ self.query_key_norm = query_key_norm
+ if self.query_key_norm:
+ self.q_norm = K2HorizonRMSNorm(
+ hidden_size=self.q_size, n_groups=self.num_heads, eps=rms_norm_eps
+ )
+ self.k_norm = K2HorizonRMSNorm(
+ hidden_size=self.kv_size, n_groups=self.num_kv_heads, eps=rms_norm_eps
+ )
+
+ self.gate_func = gate_func
+ if self.gate_func is not None:
+ self.gate_proj = ColumnParallelLinear(
+ hidden_size,
+ self.total_num_heads * self.head_dim,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate_proj",
+ )
+
+ self.tp_size = get_tensor_model_parallel_world_size()
+ self.tp_rank = get_tensor_model_parallel_rank()
+
+ w1_shape = torch.Size(
+ [
+ self.num_experts,
+ self.total_num_kv_heads * self.head_dim,
+ hidden_size,
+ ]
+ )
+ config_dtype = _get_config_dtype_str(
+ use_fp8_w8a8=False,
+ use_int8_w8a16=False,
+ use_int4_w4a16=False,
+ dtype=vllm_config.model_config.dtype,
+ )
+ self.fused_mova_config = try_get_optimal_moe_config(
+ w1_shape=w1_shape,
+ w2_shape=torch.Size([w1_shape[0], w1_shape[2], w1_shape[1]]),
+ top_k=self.num_experts_per_tok,
+ dtype=config_dtype,
+ M=vllm_config.scheduler_config.max_num_batched_tokens,
+ block_shape=None,
+ )
+
+ def compute_mova_v_sparse(self, hidden_states):
+ router_logits, _ = self.v_router(hidden_states)
+
+ routing_weights, selected_values = calc_router_weights(
+ router_logits=router_logits,
+ e_score_correction_bias=self.v_router.bias,
+ score_func=self.router_score_func,
+ top_k=self.num_experts_per_tok,
+ scaling_factor=self.router_scaling_factor,
+ )
+
+ w1 = torch.stack(
+ [expert.weight for expert in self.v_experts], dim=0
+ ).contiguous()
+
+ v = fused_mova_impl(
+ config=self.fused_mova_config,
+ hidden_states=hidden_states.contiguous(),
+ w1=w1,
+ topk_weights=routing_weights.contiguous(),
+ topk_ids=selected_values.contiguous(),
+ global_num_experts=self.num_experts,
+ expert_map=None,
+ )
+
+ return v
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ ) -> torch.Tensor:
+ qk, _ = self.qk_proj(hidden_states)
+ q, k = qk.split([self.q_size, self.kv_size], dim=-1)
+ v = self.compute_mova_v_sparse(hidden_states)
+
+ if self.query_key_norm:
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ if self.rope_head_dim == self.head_dim:
+ q, k = self.rotary_emb(positions, q, k)
+ else:
+ q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
+ k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
+
+ q_rope, q_nope = torch.split(
+ split_to_interleaved(q),
+ split_size_or_sections=[
+ self.rope_head_dim,
+ self.head_dim - self.rope_head_dim,
+ ],
+ dim=-1,
+ )
+ k_rope, k_nope = torch.split(
+ split_to_interleaved(k),
+ split_size_or_sections=[
+ self.rope_head_dim,
+ self.head_dim - self.rope_head_dim,
+ ],
+ dim=-1,
+ )
+
+ q_rope, k_rope = self.rotary_emb(
+ positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
+ )
+
+ q = interleaved_to_split(
+ torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
+ ).reshape(*q.shape[:-2], -1)
+
+ k = interleaved_to_split(
+ torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
+ ).reshape(*k.shape[:-2], -1)
+
+ attn_output = self.attn(q, k, v)
+
+ if self.gate_func is not None:
+ gate, _ = self.gate_proj(hidden_states)
+ if self.gate_func == "silu":
+ gate = F.silu(gate)
+ else:
+ assert self.gate_func == "softplus"
+ gate = F.softplus(gate, beta=math.log(2))
+
+ attn_output = attn_output * gate
+
+ output, _ = self.o_proj(attn_output)
+ return output
+
+
+class K2HorizonDecoderLayer(nn.Module):
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+
+ config = vllm_config.model_config.hf_text_config
+ cache_config = vllm_config.cache_config
+ quant_config = vllm_config.quant_config
+
+ self.hidden_size = config.hidden_size
+ max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+ dual_chunk_attention_config = getattr(
+ config, "dual_chunk_attention_config", None
+ )
+
+ # `mlp_only_layers` in the config.
+ layer_idx = extract_layer_index(prefix)
+ mlp_only_layers = (
+ [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
+ )
+ is_sparse_layer = (layer_idx not in mlp_only_layers) and (
+ config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
+ )
+
+ if is_sparse_layer and config.mova_num_experts > 0:
+ self.self_attn = K2HorizonMoVAAttention(
+ vllm_config=vllm_config,
+ hidden_size=self.hidden_size,
+ num_heads=config.num_attention_heads,
+ num_kv_heads=config.num_key_value_heads,
+ rope_parameters=config.rope_parameters,
+ rope_head_dim=config.rope_head_dim,
+ query_key_norm=config.query_key_norm,
+ max_position_embeddings=max_position_embeddings,
+ rms_norm_eps=config.rms_norm_eps,
+ qkv_bias=config.attention_bias,
+ head_dim=config.head_dim,
+ gate_func=config.attention_gate_func,
+ cache_config=cache_config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.self_attn",
+ dual_chunk_attention_config=dual_chunk_attention_config,
+ num_experts=config.mova_num_experts,
+ num_experts_per_tok=config.mova_num_experts_per_tok,
+ moe_gate_bias=config.moe_gate_bias,
+ router_score_func=config.router_score_func,
+ router_scaling_factor=config.router_scaling_factor,
+ )
+ else:
+ self.self_attn = K2HorizonAttention(
+ hidden_size=self.hidden_size,
+ num_heads=config.num_attention_heads,
+ num_kv_heads=config.num_key_value_heads,
+ rope_parameters=config.rope_parameters,
+ rope_head_dim=config.rope_head_dim,
+ query_key_norm=config.query_key_norm,
+ max_position_embeddings=max_position_embeddings,
+ rms_norm_eps=config.rms_norm_eps,
+ qkv_bias=config.attention_bias,
+ head_dim=config.head_dim,
+ gate_func=config.attention_gate_func,
+ cache_config=cache_config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.self_attn",
+ dual_chunk_attention_config=dual_chunk_attention_config,
+ )
+
+ if is_sparse_layer:
+ self.mlp = K2HorizonSparseMoeBlock(
+ vllm_config=vllm_config, prefix=f"{prefix}.mlp"
+ )
+ else:
+ self.mlp = K2HorizonMLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp",
+ )
+
+ self.input_layernorm = K2HorizonRMSNorm(
+ hidden_size=config.hidden_size,
+ n_groups=config.layernorm_num_groups,
+ eps=config.rms_norm_eps,
+ )
+
+ self.post_attention_layernorm = K2HorizonRMSNorm(
+ hidden_size=config.hidden_size,
+ n_groups=config.layernorm_num_groups,
+ eps=config.rms_norm_eps,
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor | None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Self Attention
+ if residual is None:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ else:
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+ hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+
+ # Fully Connected
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+ hidden_states = self.mlp(hidden_states)
+ return hidden_states, residual
+
+
+@support_torch_compile
+class K2HorizonModel(nn.Module, EagleModelMixin):
+ def __init__(
+ self,
+ *,
+ vllm_config: VllmConfig,
+ prefix: str = "",
+ decoder_layer_type: type[torch.nn.Module] = K2HorizonDecoderLayer,
+ ):
+ super().__init__()
+
+ config = vllm_config.model_config.hf_text_config
+ quant_config = vllm_config.quant_config
+ parallel_config = vllm_config.parallel_config
+ eplb_config = parallel_config.eplb_config
+ self.num_redundant_experts = eplb_config.num_redundant_experts
+
+ self.vocab_size = config.vocab_size
+ self.config = config
+ self.quant_config = quant_config
+ self.embed_tokens = VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.embed_tokens",
+ )
+ self.start_layer, self.end_layer, self.layers = make_layers(
+ config.num_hidden_layers,
+ lambda prefix: decoder_layer_type(vllm_config=vllm_config, prefix=prefix),
+ prefix=f"{prefix}.layers",
+ )
+
+ self.norm = K2HorizonRMSNorm(
+ hidden_size=config.hidden_size,
+ n_groups=config.layernorm_num_groups,
+ eps=config.rms_norm_eps,
+ )
+
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
+ )
+
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return self.embed_tokens(input_ids)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
+ if get_pp_group().is_first_rank:
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+ else:
+ hidden_states = self.embed_input_ids(input_ids)
+ residual = None
+ else:
+ assert intermediate_tensors is not None
+ hidden_states = intermediate_tensors["hidden_states"]
+ residual = intermediate_tensors["residual"]
+
+ aux_hidden_states = self._maybe_add_hidden_state(
+ [], self.start_layer, hidden_states, residual
+ )
+ for layer_idx, layer in enumerate(
+ islice(self.layers, self.start_layer, self.end_layer),
+ start=self.start_layer,
+ ):
+ hidden_states, residual = layer(positions, hidden_states, residual)
+ self._maybe_add_hidden_state(
+ aux_hidden_states, layer_idx + 1, hidden_states, residual
+ )
+
+ if not get_pp_group().is_last_rank:
+ return IntermediateTensors(
+ {"hidden_states": hidden_states, "residual": residual}
+ )
+ hidden_states, _ = self.norm(hidden_states, residual)
+
+ # Return auxiliary hidden states if collected
+ if len(aux_hidden_states) > 0:
+ return hidden_states, aux_hidden_states
+ return hidden_states
+
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+ # Params for weights, fp8 weight scales, fp8 activation scales
+ # (param_name, weight_name, expert_id, shard_id)
+ return fused_moe_make_expert_params_mapping(
+ self,
+ ckpt_gate_proj_name="gate_proj",
+ ckpt_down_proj_name="down_proj",
+ ckpt_up_proj_name="up_proj",
+ num_experts=self.config.num_experts,
+ num_redundant_experts=self.num_redundant_experts,
+ )
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ stacked_params_mapping = [
+ # (param_name, shard_name, shard_id)
+ ("qkv_proj", "q_proj", "q"),
+ ("qkv_proj", "k_proj", "k"),
+ ("qkv_proj", "v_proj", "v"),
+ ("gate_up_proj", "gate_proj", 0),
+ ("gate_up_proj", "up_proj", 1),
+ ]
+
+ # Skip loading extra parameters for GPTQ/modelopt models.
+ # ignore_suffixes = (
+ # ".bias",
+ # "_bias",
+ # ".weight_scale",
+ # "_weight_scale",
+ # ".input_scale",
+ # "_input_scale",
+ # )
+ ignore_suffixes = ()
+
+ params_dict = dict(self.named_parameters())
+
+ loaded_params: set[str] = set()
+ expert_params_mapping = self.get_expert_mapping()
+ for name, loaded_weight in weights:
+ if "scale" in name or "zero_point" in name:
+ name = maybe_remap_kv_scale_name(name, params_dict)
+ if name is None:
+ continue
+
+ # QK norm weights
+ if name.endswith(".self_attn.q_norm.weight") or name.endswith(
+ ".self_attn.k_norm.weight"
+ ):
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name not in params_dict:
+ continue
+
+ param = params_dict[name]
+ tp_rank = get_tensor_model_parallel_rank()
+ tp_size = get_tensor_model_parallel_world_size()
+
+ if loaded_weight.shape != param.shape:
+ if name.endswith(".self_attn.q_norm.weight"):
+ loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
+ else:
+ num_kv_heads = self.config.num_key_value_heads
+ if num_kv_heads >= tp_size:
+ loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
+ else:
+ num_kv_head_replicas = tp_size // num_kv_heads
+ kv_rank = tp_rank // num_kv_head_replicas
+ loaded_weight = loaded_weight.chunk(num_kv_heads, dim=0)[
+ kv_rank
+ ]
+
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
+ weight_loader(param, loaded_weight)
+ loaded_params.add(name)
+ continue
+
+ # MoVA v experts
+ if ".self_attn.v_experts." in name:
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name not in params_dict:
+ continue
+
+ param = params_dict[name]
+ tp_rank = get_tensor_model_parallel_rank()
+ tp_size = get_tensor_model_parallel_world_size()
+ num_kv_heads = self.config.num_key_value_heads
+
+ if loaded_weight.shape != param.shape:
+ if num_kv_heads >= tp_size:
+ loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
+ else:
+ num_kv_head_replicas = tp_size // num_kv_heads
+ kv_rank = tp_rank // num_kv_head_replicas
+ loaded_weight = loaded_weight.chunk(num_kv_heads, dim=0)[
+ kv_rank
+ ]
+
+ default_weight_loader(param, loaded_weight)
+ loaded_params.add(name)
+ continue
+
+ for param_name, weight_name, shard_id in stacked_params_mapping:
+ # Skip non-stacked layers and experts (experts handled below).
+ if weight_name not in name:
+ continue
+ # We have mlp.experts[0].gate_proj in the checkpoint.
+ # Since we handle the experts below in expert_params_mapping,
+ # we need to skip here BEFORE we update the name, otherwise
+ # name will be updated to mlp.experts[0].gate_up_proj, which
+ # will then be updated below in expert_params_mapping
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+ if "mlp.experts" in name:
+ continue
+ if ".self_attn." in name and param_name == "gate_up_proj":
+ assert weight_name == "gate_proj"
+ continue
+ if (
+ ".self_attn." in name
+ and param_name == "qkv_proj"
+ and name.replace(weight_name, param_name) not in params_dict
+ ):
+ qk_name = name.replace(weight_name, "qk_proj")
+ if is_pp_missing_parameter(qk_name, self):
+ continue
+ if qk_name not in params_dict:
+ continue
+
+ assert weight_name in ["q_proj", "k_proj"]
+ if weight_name == "k_proj":
+ tp_size = get_tensor_model_parallel_world_size()
+ num_kv_heads = self.config.num_key_value_heads
+ if num_kv_heads < tp_size:
+ num_kv_head_replicas = tp_size // num_kv_heads
+ loaded_weight = (
+ loaded_weight.reshape(
+ num_kv_heads, -1, *loaded_weight.shape[1:]
+ )
+ .repeat_interleave(num_kv_head_replicas, dim=0)
+ .reshape(-1, *loaded_weight.shape[1:])
+ )
+
+ param = params_dict[qk_name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(
+ param, loaded_weight, 0 if weight_name == "q_proj" else 1
+ )
+ loaded_params.add(qk_name)
+ break
+
+ name = name.replace(weight_name, param_name)
+
+ # Skip loading extra parameters for GPTQ/modelopt models.
+ if name.endswith(ignore_suffixes) and name not in params_dict:
+ continue
+
+ # Skip layers on other devices.
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name.endswith("scale"):
+ # Remapping the name of FP8 kv-scale.
+ name = maybe_remap_kv_scale_name(name, params_dict)
+ if name is None:
+ continue
+ if name not in params_dict:
+ continue
+
+ param = params_dict[name]
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
+ if weight_loader == default_weight_loader:
+ weight_loader(param, loaded_weight)
+ else:
+ weight_loader(param, loaded_weight, shard_id)
+ break
+ else:
+ is_expert_weight = False
+ for mapping in expert_params_mapping:
+ param_name, weight_name, expert_id, shard_id = mapping
+ if weight_name not in name:
+ continue
+
+ # Anyway, this is an expert weight and should not be
+ # attempted to load as other weights later
+ is_expert_weight = True
+
+ # Do not modify `name` since the loop may continue here
+ # Instead, create a new variable
+ name_mapped = name.replace(weight_name, param_name)
+
+ if is_pp_missing_parameter(name_mapped, self):
+ continue
+
+ # Skip loading extra parameters for GPTQ/modelopt models.
+ if (
+ name_mapped.endswith(ignore_suffixes)
+ and name_mapped not in params_dict
+ ):
+ continue
+
+ param = params_dict[name_mapped]
+ # We should ask the weight loader to return success or not
+ # here since otherwise we may skip experts with other
+ # available replicas.
+ weight_loader = typing.cast(
+ Callable[..., bool], param.weight_loader
+ )
+ success = weight_loader(
+ param,
+ loaded_weight,
+ name_mapped,
+ shard_id=shard_id,
+ expert_id=expert_id,
+ return_success=True,
+ )
+ if success:
+ name = name_mapped
+ break
+ else:
+ if is_expert_weight:
+ # We've checked that this is an expert weight
+ # However it's not mapped locally to this rank
+ # So we simply skip it
+ continue
+
+ # Skip loading extra parameters for GPTQ/modelopt models.
+ if name.endswith(ignore_suffixes) and name not in params_dict:
+ continue
+ # Skip layers on other devices.
+ if is_pp_missing_parameter(name, self):
+ continue
+ if name not in params_dict:
+ continue
+ param = params_dict[name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(name)
+
+ return loaded_params
+
+
+class K2HorizonForCausalLM(
+ nn.Module, SupportsPP, SupportsLoRA, SupportsEagle, SupportsEagle3, MixtureOfExperts
+):
+ packed_modules_mapping = {
+ "qkv_proj": [
+ "q_proj",
+ "k_proj",
+ "v_proj",
+ ],
+ "qk_proj": [
+ "q_proj",
+ "k_proj",
+ ],
+ }
+
+ embedding_modules = {
+ "embed_tokens": "input_embeddings",
+ "lm_head": "output_embeddings",
+ }
+
+ fall_back_to_pt_during_load = False
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+ super().__init__()
+ config = vllm_config.model_config.hf_text_config
+ quant_config = vllm_config.quant_config
+ self.config = config
+ self.quant_config = quant_config
+ # Only perform the following mapping when K2HorizonMLP exists
+ if getattr(config, "mlp_only_layers", []):
+ self.packed_modules_mapping["gate_up_proj"] = ["gate_proj", "up_proj"]
+ self.model = K2HorizonModel(
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+ )
+ self.lm_head = ParallelLMHead(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=maybe_prefix(prefix, "lm_head"),
+ )
+ if self.config.tie_word_embeddings:
+ self.lm_head.weight = self.model.embed_tokens.weight
+ self.logits_processor = LogitsProcessor(config.vocab_size)
+ self.make_empty_intermediate_tensors = (
+ self.model.make_empty_intermediate_tensors
+ )
+
+ # Set MoE hyperparameters
+ self.expert_weights = []
+
+ self.moe_layers = []
+ example_layer = None
+ for layer in self.model.layers:
+ if isinstance(layer, PPMissingLayer):
+ continue
+
+ assert isinstance(layer, K2HorizonDecoderLayer)
+ if isinstance(layer.mlp, K2HorizonSparseMoeBlock):
+ example_layer = layer.mlp
+ self.moe_layers.append(layer.mlp.experts)
+
+ self.num_moe_layers = len(self.moe_layers)
+ if example_layer is not None:
+ self.num_logical_experts = example_layer.n_logical_experts
+ self.num_physical_experts = example_layer.n_physical_experts
+ self.num_local_physical_experts = example_layer.n_local_physical_experts
+ self.num_routed_experts = example_layer.n_routed_experts
+ self.num_redundant_experts = example_layer.n_redundant_experts
+
+ def update_physical_experts_metadata(
+ self,
+ num_physical_experts: int,
+ num_local_physical_experts: int,
+ ) -> None:
+ assert self.num_local_physical_experts == num_local_physical_experts
+ self.num_physical_experts = num_physical_experts
+ self.num_local_physical_experts = num_local_physical_experts
+ self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+ for layer in self.model.layers:
+ if isinstance(layer.mlp, K2HorizonSparseMoeBlock):
+ moe = layer.mlp
+ moe.n_local_physical_experts = num_local_physical_experts
+ moe.n_physical_experts = num_physical_experts
+ moe.n_redundant_experts = self.num_redundant_experts
+ moe.experts.update_expert_map()
+
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return self.model.embed_input_ids(input_ids)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor | IntermediateTensors:
+ hidden_states = self.model(
+ input_ids, positions, intermediate_tensors, inputs_embeds
+ )
+ return hidden_states
+
+ def compute_logits(
+ self,
+ hidden_states: torch.Tensor,
+ ) -> torch.Tensor | None:
+ logits = self.logits_processor(self.lm_head, hidden_states)
+ return logits
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ loader = AutoWeightsLoader(self)
+ return loader.load_weights(weights)
+
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+ return self.model.get_expert_mapping()
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index 9f77ea83e4..d5a490e7bf 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -141,6 +141,7 @@ _TEXT_GENERATION_MODELS = {
"IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
"Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
"JambaForCausalLM": ("jamba", "JambaForCausalLM"),
+ "K2HorizonForCausalLM": ("k2_horizon", "K2HorizonForCausalLM"),
"KimiLinearForCausalLM": (
"vllm.models.kimi_k3",
"KimiLinearForCausalLM",
diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py
index 18e38e74c1..70469f6695 100644
--- a/vllm/reasoning/__init__.py
+++ b/vllm/reasoning/__init__.py
@@ -96,6 +96,10 @@ _REASONING_PARSERS_TO_REGISTER = {
"kimi_k3_reasoning_parser",
"KimiK3ReasoningParser",
),
+ "k2_horizon": (
+ "k2_horizon_reasoning_parser",
+ "K2HorizonReasoningParser",
+ ),
"mimo": (
"qwen3_engine_reasoning_parser",
"Qwen3ParserReasoningAdapter",
diff --git a/vllm/reasoning/k2_horizon_reasoning_parser.py b/vllm/reasoning/k2_horizon_reasoning_parser.py
new file mode 100644
index 0000000000..4c24c6d1c4
--- /dev/null
+++ b/vllm/reasoning/k2_horizon_reasoning_parser.py
@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable, Sequence
+
+from vllm.entrypoints.generate.base.protocol import DeltaMessage
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
+from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.utils import partial_tag_overlap
+
+
+class K2HorizonReasoningParser(DeepSeekR1ReasoningParser):
+ """Reasoning parser for K2 Horizon."""
+
+ _TOOL_CALLS_START = ""
+ _EFFORT_TOKENS: dict[str, tuple[str, str]] = {
+ "high": ("", ""),
+ "medium": ("", ""),
+ "low": ("", ""),
+ }
+
+ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs) -> None:
+ chat_template_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
+ effort = chat_template_kwargs.get("reasoning_effort", "high")
+ if not isinstance(effort, str) or effort not in self._EFFORT_TOKENS:
+ supported = ", ".join(self._EFFORT_TOKENS)
+ raise ValueError(
+ f"Unsupported reasoning_effort {effort!r}. "
+ f"Supported values: {supported}."
+ )
+
+ self._start_token, self._end_token = self._EFFORT_TOKENS[effort]
+ super().__init__(tokenizer, *args, **kwargs)
+ self._tool_calls_start_id = self.vocab.get(self._TOOL_CALLS_START)
+ self._stream_buffer = ""
+ self._stream_at_start = True
+ self._stream_finished = False
+
+ @property
+ def start_token(self) -> str:
+ return self._start_token
+
+ @property
+ def end_token(self) -> str:
+ return self._end_token
+
+ def _without_generated_start(self, model_output: str) -> str:
+ return model_output.removeprefix(self.start_token)
+
+ def _split_model_output(self, model_output: str) -> tuple[str, str | None]:
+ output = self._without_generated_start(model_output)
+ if self.end_token in output:
+ reasoning, _, content = output.partition(self.end_token)
+ return reasoning, content or None
+
+ tool_start = output.find(self._TOOL_CALLS_START)
+ if tool_start != -1:
+ return output[:tool_start], output[tool_start:]
+
+ return "", output or None
+
+ def extract_reasoning(
+ self,
+ model_output: str,
+ request: ChatCompletionRequest | ResponsesRequest,
+ ) -> tuple[str, str | None]:
+ del request
+ return self._split_model_output(model_output)
+
+ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+ for token_id in reversed(input_ids):
+ if token_id == self.start_token_id:
+ return False
+ if token_id == self.end_token_id or token_id == self._tool_calls_start_id:
+ return True
+ return False
+
+ def is_reasoning_end_streaming(
+ self,
+ input_ids: Sequence[int],
+ delta_ids: Iterable[int],
+ ) -> bool:
+ del input_ids
+ return self._stream_finished or super().is_reasoning_end_streaming(
+ (), delta_ids
+ )
+
+ def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+ if self.end_token_id in input_ids:
+ return input_ids[input_ids.index(self.end_token_id) + 1 :]
+ if (
+ self._tool_calls_start_id is not None
+ and self._tool_calls_start_id in input_ids
+ ):
+ return input_ids[input_ids.index(self._tool_calls_start_id) :]
+ return input_ids if self._stream_finished else []
+
+ def _strip_stream_start(self) -> bool:
+ if not self._stream_at_start:
+ return True
+ if self.start_token.startswith(self._stream_buffer):
+ if self._stream_buffer != self.start_token:
+ return False
+ self._stream_buffer = ""
+ elif self._stream_buffer.startswith(self.start_token):
+ self._stream_buffer = self._stream_buffer[len(self.start_token) :]
+ self._stream_at_start = False
+ return True
+
+ def extract_reasoning_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int],
+ current_token_ids: Sequence[int],
+ delta_token_ids: Sequence[int],
+ ) -> DeltaMessage | None:
+ del (
+ previous_text,
+ current_text,
+ previous_token_ids,
+ current_token_ids,
+ delta_token_ids,
+ )
+
+ if self._stream_finished:
+ return DeltaMessage(content=delta_text) if delta_text else None
+
+ self._stream_buffer += delta_text
+ if not self._strip_stream_start():
+ return None
+
+ explicit_end = self._stream_buffer.find(self.end_token)
+ tool_start = self._stream_buffer.find(self._TOOL_CALLS_START)
+ if explicit_end != -1:
+ reasoning = self._stream_buffer[:explicit_end]
+ content = self._stream_buffer[explicit_end + len(self.end_token) :]
+ elif tool_start != -1:
+ reasoning = self._stream_buffer[:tool_start]
+ content = self._stream_buffer[tool_start:]
+ else:
+ overlap = max(
+ partial_tag_overlap(self._stream_buffer, self.end_token),
+ partial_tag_overlap(self._stream_buffer, self._TOOL_CALLS_START),
+ )
+ sendable_end = len(self._stream_buffer) - overlap
+ reasoning = self._stream_buffer[:sendable_end]
+ self._stream_buffer = self._stream_buffer[sendable_end:]
+ return DeltaMessage(reasoning=reasoning) if reasoning else None
+
+ self._stream_buffer = ""
+ self._stream_finished = True
+ return DeltaMessage(
+ reasoning=reasoning,
+ content=content or None,
+ )
diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py
index 06ed049beb..a536891bae 100644
--- a/vllm/tool_parsers/__init__.py
+++ b/vllm/tool_parsers/__init__.py
@@ -118,6 +118,10 @@ _TOOL_PARSERS_TO_REGISTER = {
"kimi_k3_tool_parser",
"KimiK3ToolParser",
),
+ "k2_horizon": (
+ "k2_horizon_tool_parser",
+ "K2HorizonToolParser",
+ ),
"llama3_json": (
"llama_tool_parser",
"Llama3JsonToolParser",
diff --git a/vllm/tool_parsers/k2_horizon_tool_parser.py b/vllm/tool_parsers/k2_horizon_tool_parser.py
new file mode 100644
index 0000000000..742bf7d140
--- /dev/null
+++ b/vllm/tool_parsers/k2_horizon_tool_parser.py
@@ -0,0 +1,462 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from collections.abc import Sequence
+from typing import Any
+
+import regex as re
+from openai.types.responses import ToolChoiceFunction
+
+from vllm.entrypoints.generate.base.protocol import (
+ DeltaFunctionCall,
+ DeltaMessage,
+ DeltaToolCall,
+ ExtractedToolCallInformation,
+ FunctionCall,
+ ToolCall,
+)
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionNamedToolChoiceParam,
+ ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.logger import init_logger
+from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
+from vllm.tool_parsers.utils import (
+ coerce_to_schema_type,
+ extract_types_from_schema,
+ find_tool_name,
+ find_tool_properties,
+ partial_tag_overlap,
+)
+
+logger = init_logger(__name__)
+
+ToolParserRequest = ChatCompletionRequest | ResponsesRequest
+
+
+class K2HorizonToolParser(ToolParser):
+ """Parser for K2 Horizon's IFM tool-call formats."""
+
+ supports_required_and_named = False
+ engine_based_streaming = True
+
+ _SUPPORTED_TOOL_FORMATS = frozenset({"json", "xml", "xml_typed"})
+
+ _TOOL_CALLS_START = ""
+ _TOOL_CALLS_END = ""
+ _TOOL_CALL_START = ""
+ _TOOL_CALL_END = ""
+ _ARG_KEY_START = ""
+ _ARG_KEY_END = ""
+ _ARG_TYPE_START = ""
+ _ARG_TYPE_END = ""
+ _ARG_VALUE_START = ""
+ _ARG_VALUE_END = ""
+
+ _TOOL_CALL_RE = re.compile(
+ re.escape(_TOOL_CALL_START) + r"(.*?)" + re.escape(_TOOL_CALL_END),
+ re.DOTALL,
+ )
+ _ARG_RE = re.compile(
+ re.escape(_ARG_KEY_START)
+ + r"(.*?)"
+ + re.escape(_ARG_KEY_END)
+ + r"\s*(?:"
+ + re.escape(_ARG_TYPE_START)
+ + r"(.*?)"
+ + re.escape(_ARG_TYPE_END)
+ + r"\s*)?"
+ + re.escape(_ARG_VALUE_START)
+ + r"(.*?)"
+ + re.escape(_ARG_VALUE_END),
+ re.DOTALL,
+ )
+ _ARG_MARKERS = (
+ _ARG_KEY_START,
+ _ARG_KEY_END,
+ _ARG_TYPE_START,
+ _ARG_TYPE_END,
+ _ARG_VALUE_START,
+ _ARG_VALUE_END,
+ )
+
+ def __init__(
+ self,
+ tokenizer: TokenizerLike,
+ tools: list[Tool] | None = None,
+ ) -> None:
+ super().__init__(tokenizer, tools)
+ self.tool_format = "xml"
+ self._stream_text = ""
+ self._stream_cursor = 0
+ self._stream_group_start: int | None = None
+ self._stream_resolved = False
+ self._pending_content = ""
+ self._content_seen = False
+
+ @classmethod
+ def _request_tool_format(cls, request: ToolParserRequest) -> str:
+ chat_template_kwargs = request.chat_template_kwargs or {}
+ tool_format = chat_template_kwargs.get("tool_call_format", "xml")
+ if (
+ not isinstance(tool_format, str)
+ or tool_format not in cls._SUPPORTED_TOOL_FORMATS
+ ):
+ supported = ", ".join(sorted(cls._SUPPORTED_TOOL_FORMATS))
+ raise ValueError(
+ f"Unsupported tool_call_format {tool_format!r}. "
+ f"Supported values: {supported}."
+ )
+ return tool_format
+
+ def adjust_request(self, request: ToolParserRequest) -> ToolParserRequest:
+ self._request_tool_format(request)
+ return request
+
+ @staticmethod
+ def _content_or_none(content: str) -> str | None:
+ return content if content.strip() else None
+
+ @classmethod
+ def _no_tool_calls(cls, content: str) -> ExtractedToolCallInformation:
+ return ExtractedToolCallInformation(
+ tools_called=False,
+ tool_calls=[],
+ content=cls._content_or_none(content),
+ )
+
+ @staticmethod
+ def _make_tool_call(name: str, arguments: dict[str, Any]) -> ToolCall:
+ return ToolCall(
+ type="function",
+ function=FunctionCall(
+ name=name,
+ arguments=json.dumps(
+ arguments,
+ ensure_ascii=False,
+ allow_nan=False,
+ ),
+ ),
+ )
+
+ @staticmethod
+ def _base_arg_type(arg_type: str | None) -> str | None:
+ if not arg_type:
+ return None
+ return arg_type.split("[", 1)[0].strip().lower()
+
+ @staticmethod
+ def _coerce_json_value(value: Any, schema: Any) -> Any:
+ if not isinstance(schema, dict):
+ return value
+ raw_value = (
+ value
+ if isinstance(value, str)
+ else json.dumps(value, ensure_ascii=False, allow_nan=False)
+ )
+ return coerce_to_schema_type(raw_value, extract_types_from_schema(schema))
+
+ @classmethod
+ def _coerce_xml_value(
+ cls,
+ value: str,
+ schema: Any,
+ explicit_type: str | None,
+ ) -> Any:
+ explicit_base = cls._base_arg_type(explicit_type)
+ target_types = extract_types_from_schema(schema)
+ if isinstance(schema, dict):
+ if explicit_base in target_types and len(set(target_types)) > 1:
+ target_types = [explicit_base]
+ elif explicit_base == "any":
+ return value
+ elif explicit_base:
+ target_types = [explicit_base]
+
+ raw_value = value if "string" in target_types else value.strip()
+ return coerce_to_schema_type(raw_value, target_types)
+
+ @staticmethod
+ def _named_tool(request: ToolParserRequest) -> str | None:
+ tool_choice = request.tool_choice
+ if isinstance(tool_choice, ChatCompletionNamedToolChoiceParam):
+ return tool_choice.function.name
+ if isinstance(tool_choice, ToolChoiceFunction):
+ return tool_choice.name
+ return None
+
+ @classmethod
+ def _validate_tool_name(
+ cls,
+ name: Any,
+ request: ToolParserRequest,
+ ) -> str:
+ if (
+ not isinstance(name, str)
+ or not name
+ or any(char.isspace() for char in name)
+ ):
+ raise ValueError("IFM tool call has an invalid function name.")
+ if request.tools and not find_tool_name(request.tools, name):
+ raise ValueError(f"Unknown IFM tool name: {name}.")
+ named_tool = cls._named_tool(request)
+ if named_tool is not None and name != named_tool:
+ raise ValueError(f"Unexpected IFM tool name: {name}.")
+ return name
+
+ @classmethod
+ def _parse_json_call(
+ cls,
+ body: str,
+ request: ToolParserRequest,
+ ) -> tuple[str, dict[str, Any]]:
+ raw_call = json.loads(body.strip())
+ if not isinstance(raw_call, dict):
+ raise ValueError("IFM JSON tool call must be an object.")
+
+ name = cls._validate_tool_name(raw_call.get("name"), request)
+ arguments = raw_call.get("arguments", {})
+ if not isinstance(arguments, dict):
+ raise ValueError("IFM JSON tool-call arguments must be an object.")
+
+ properties = find_tool_properties(request.tools, name)
+ arguments = {
+ arg_name: cls._coerce_json_value(arg_value, properties.get(arg_name))
+ for arg_name, arg_value in arguments.items()
+ }
+ return name, arguments
+
+ @classmethod
+ def _parse_xml_call(
+ cls,
+ body: str,
+ request: ToolParserRequest,
+ tool_format: str,
+ ) -> tuple[str, dict[str, Any]]:
+ first_arg = body.find(cls._ARG_KEY_START)
+ if first_arg == -1:
+ if any(marker in body for marker in cls._ARG_MARKERS):
+ raise ValueError("Malformed IFM XML argument tags.")
+ name = cls._validate_tool_name(body.strip(), request)
+ return name, {}
+
+ name = cls._validate_tool_name(body[:first_arg].strip(), request)
+ properties = find_tool_properties(request.tools, name)
+ arguments: dict[str, Any] = {}
+ position = first_arg
+ for match in cls._ARG_RE.finditer(body, first_arg):
+ if body[position : match.start()].strip():
+ raise ValueError("Malformed IFM XML argument tags.")
+
+ arg_name = match.group(1).strip()
+ if not arg_name:
+ raise ValueError("IFM XML argument is missing a name.")
+ if arg_name in arguments:
+ raise ValueError(f"Duplicate IFM XML argument: {arg_name}.")
+
+ explicit_type = match.group(2)
+ if explicit_type is not None:
+ explicit_type = explicit_type.strip()
+ if tool_format == "xml" and explicit_type is not None:
+ raise ValueError("Unexpected type tag in IFM XML tool call.")
+ if tool_format == "xml_typed" and not explicit_type:
+ raise ValueError("IFM typed XML argument is missing a type.")
+ arguments[arg_name] = cls._coerce_xml_value(
+ match.group(3),
+ properties.get(arg_name),
+ explicit_type,
+ )
+ position = match.end()
+
+ if position == first_arg or body[position:].strip():
+ raise ValueError("Malformed IFM XML argument tags.")
+ return name, arguments
+
+ @classmethod
+ def _parse_call(
+ cls,
+ body: str,
+ request: ToolParserRequest,
+ tool_format: str,
+ ) -> tuple[str, dict[str, Any]]:
+ if tool_format == "json":
+ return cls._parse_json_call(body, request)
+ return cls._parse_xml_call(body, request, tool_format)
+
+ def extract_tool_calls(
+ self,
+ model_output: str,
+ request: ToolParserRequest,
+ ) -> ExtractedToolCallInformation:
+ tool_format = self._request_tool_format(request)
+ if request.tool_choice == "none":
+ return self._no_tool_calls(model_output)
+
+ group_start = model_output.find(self._TOOL_CALLS_START)
+ if group_start == -1:
+ return self._no_tool_calls(model_output)
+
+ body_start = group_start + len(self._TOOL_CALLS_START)
+ group_end = model_output.find(self._TOOL_CALLS_END, body_start)
+ if group_end == -1:
+ return self._no_tool_calls(model_output)
+
+ group_body = model_output[body_start:group_end]
+ matches = list(self._TOOL_CALL_RE.finditer(group_body))
+ if not matches or self._TOOL_CALL_RE.sub("", group_body).strip():
+ logger.warning("Malformed K2 Horizon IFM tool-call group.")
+ return self._no_tool_calls(model_output)
+
+ try:
+ tool_calls = [
+ self._make_tool_call(
+ *self._parse_call(match.group(1), request, tool_format)
+ )
+ for match in matches
+ ]
+ except (json.JSONDecodeError, TypeError, ValueError):
+ logger.warning("Failed to parse K2 Horizon IFM tool calls.", exc_info=True)
+ return self._no_tool_calls(model_output)
+
+ suffix = model_output[group_end + len(self._TOOL_CALLS_END) :]
+ content = model_output[:group_start] + suffix
+ return ExtractedToolCallInformation(
+ tools_called=True,
+ tool_calls=tool_calls,
+ content=self._content_or_none(content),
+ )
+
+ def _emit_content(self, addition: str, *, final: bool = False) -> str | None:
+ if addition:
+ if self._content_seen:
+ return addition
+ self._pending_content += addition
+ if self._pending_content.strip():
+ self._content_seen = True
+ content = self._pending_content
+ self._pending_content = ""
+ return content
+ if final:
+ self._pending_content = ""
+ return None
+
+ def _streaming_tool_calls(
+ self,
+ tool_calls: list[ToolCall],
+ ) -> list[DeltaToolCall]:
+ deltas: list[DeltaToolCall] = []
+ for index, tool_call in enumerate(tool_calls):
+ arguments_json = tool_call.function.arguments
+ arguments = json.loads(arguments_json)
+ self.prev_tool_call_arr.append(
+ {"name": tool_call.function.name, "arguments": arguments}
+ )
+ self.streamed_args_for_tool.append(arguments_json)
+ deltas.append(
+ DeltaToolCall(
+ index=index,
+ id=tool_call.id,
+ type="function",
+ function=DeltaFunctionCall(
+ name=tool_call.function.name,
+ arguments=arguments_json,
+ ),
+ )
+ )
+ return deltas
+
+ def extract_tool_calls_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int],
+ current_token_ids: Sequence[int],
+ delta_token_ids: Sequence[int],
+ request: ToolParserRequest,
+ ) -> DeltaMessage | None:
+ del (
+ previous_text,
+ current_text,
+ previous_token_ids,
+ current_token_ids,
+ delta_token_ids,
+ )
+
+ self._request_tool_format(request)
+ if request.tool_choice == "none":
+ emitted_content = self._emit_content(delta_text)
+ return (
+ DeltaMessage(content=emitted_content)
+ if emitted_content is not None
+ else None
+ )
+
+ self._stream_text += delta_text
+ if self._stream_resolved:
+ emitted_content = self._emit_content(delta_text)
+ return (
+ DeltaMessage(content=emitted_content)
+ if emitted_content is not None
+ else None
+ )
+
+ content: str | None = None
+ if self._stream_group_start is None:
+ group_start = self._stream_text.find(
+ self._TOOL_CALLS_START, self._stream_cursor
+ )
+ if group_start == -1:
+ pending = self._stream_text[self._stream_cursor :]
+ overlap = partial_tag_overlap(pending, self._TOOL_CALLS_START)
+ sendable_end = len(self._stream_text) - overlap
+ content = self._emit_content(
+ self._stream_text[self._stream_cursor : sendable_end]
+ )
+ self._stream_cursor = sendable_end
+ return DeltaMessage(content=content) if content is not None else None
+
+ self._stream_group_start = group_start
+ content = self._emit_content(
+ self._stream_text[self._stream_cursor : group_start]
+ )
+ self._stream_cursor = group_start
+
+ group_start = self._stream_group_start
+ body_start = group_start + len(self._TOOL_CALLS_START)
+ group_end = self._stream_text.find(self._TOOL_CALLS_END, body_start)
+ if group_end == -1:
+ return DeltaMessage(content=content) if content is not None else None
+
+ group_end += len(self._TOOL_CALLS_END)
+ group_text = self._stream_text[group_start:group_end]
+ parsed = self.extract_tool_calls(group_text, request)
+ self._stream_resolved = True
+
+ if not parsed.tools_called:
+ failed_content = self._emit_content(
+ self._stream_text[self._stream_cursor :]
+ )
+ self._stream_cursor = len(self._stream_text)
+ combined = (content or "") + (failed_content or "")
+ return DeltaMessage(content=combined) if combined else None
+
+ tool_calls = self._streaming_tool_calls(parsed.tool_calls)
+ suffix = self._emit_content(self._stream_text[group_end:])
+ self._stream_cursor = len(self._stream_text)
+ combined = (content or "") + (suffix or "")
+ return DeltaMessage(content=combined or None, tool_calls=tool_calls)
+
+ def finish_streaming(self) -> DeltaMessage | None:
+ if not self._stream_resolved:
+ content = self._emit_content(
+ self._stream_text[self._stream_cursor :], final=True
+ )
+ self._stream_cursor = len(self._stream_text)
+ self._stream_resolved = True
+ else:
+ content = self._emit_content("", final=True)
+ return DeltaMessage(content=content) if content is not None else None