[Model] Add native IFM K2 Horizon serving support (#37654)
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
co-authored by
Xiaoyu Zhang
Xinyuan Tong
parent
02d9b3060a
commit
3bac084d4e
@@ -25,6 +25,7 @@ from sglang.srt.configs.interns2preview import InternS2PreviewConfig
|
||||
from sglang.srt.configs.janus_pro import MultiModalityConfig
|
||||
from sglang.srt.configs.jet_nemotron import JetNemotronConfig
|
||||
from sglang.srt.configs.jet_vlm import JetVLMConfig
|
||||
from sglang.srt.configs.k2_horizon import K2HorizonConfig, XllmConfig
|
||||
from sglang.srt.configs.kimi_k3 import KimiK3Config
|
||||
from sglang.srt.configs.kimi_k25 import KimiK25Config
|
||||
from sglang.srt.configs.kimi_linear import KimiLinearConfig
|
||||
@@ -79,6 +80,8 @@ __all__ = [
|
||||
"MultiModalityConfig",
|
||||
"KimiVLConfig",
|
||||
"MoonViTConfig",
|
||||
"K2HorizonConfig",
|
||||
"XllmConfig",
|
||||
"Step3VLConfig",
|
||||
"Step3TextConfig",
|
||||
"Step3VisionEncoderConfig",
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# 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.
|
||||
|
||||
"""Native configuration shells for the xLLM and K2 Horizon model families.
|
||||
|
||||
The released checkpoints persist the complete architecture in ``config.json``.
|
||||
Keeping these classes intentionally thin lets SGLang select its native runtime
|
||||
without importing checkpoint-provided Python code, while preserving every
|
||||
model-specific field for validation in ``sglang.srt.models.xllm``.
|
||||
"""
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
|
||||
class XllmConfig(PretrainedConfig):
|
||||
model_type = "xllm"
|
||||
|
||||
|
||||
class K2HorizonConfig(PretrainedConfig):
|
||||
model_type = "k2_horizon"
|
||||
@@ -576,6 +576,7 @@ class ModelConfig:
|
||||
self.hf_eos_token_id = self._get_hf_eos_token_id()
|
||||
# Set by scheduler when reasoning_parser is enabled
|
||||
self.think_end_ids: Optional[List[int]] = None
|
||||
self.request_selectable_think_end_id_sequences: Optional[List[List[int]]] = None
|
||||
|
||||
# multimodal
|
||||
self.image_token_id = getattr(
|
||||
|
||||
@@ -14,6 +14,9 @@ from sglang.srt.constrained.base_grammar_backend import (
|
||||
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
|
||||
from sglang.srt.distributed.communication_tags import P2PTag
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.sampling.sampling_params import (
|
||||
get_request_reasoning_end_token_ids,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.io_struct import AbortReq
|
||||
@@ -121,11 +124,21 @@ class GrammarManager:
|
||||
thinking_budget = custom_params.get("thinking_budget")
|
||||
return thinking_budget if isinstance(thinking_budget, int) else None
|
||||
|
||||
def _apply_request_reasoning_budget(self, req: Req) -> None:
|
||||
thinking_budget = self._get_request_thinking_budget(req)
|
||||
if thinking_budget is None:
|
||||
def _apply_request_reasoning_config(self, req: Req) -> None:
|
||||
if not isinstance(req.grammar, ReasonerGrammarObject):
|
||||
return
|
||||
if isinstance(req.grammar, ReasonerGrammarObject):
|
||||
think_end_ids = get_request_reasoning_end_token_ids(
|
||||
req.sampling_params.custom_params,
|
||||
allowed_sequences=getattr(
|
||||
self.scheduler.model_config,
|
||||
"request_selectable_think_end_id_sequences",
|
||||
None,
|
||||
),
|
||||
)
|
||||
if think_end_ids is not None:
|
||||
req.grammar.set_request_think_end_ids(think_end_ids)
|
||||
thinking_budget = self._get_request_thinking_budget(req)
|
||||
if thinking_budget is not None:
|
||||
req.grammar.max_think_tokens = thinking_budget
|
||||
|
||||
def process_req_with_grammar(self, req: Req) -> bool:
|
||||
@@ -167,14 +180,14 @@ class GrammarManager:
|
||||
)
|
||||
req.set_finish_with_abort(error_msg)
|
||||
else:
|
||||
self._apply_request_reasoning_budget(req)
|
||||
self._apply_request_reasoning_config(req)
|
||||
elif self._enable_strict_thinking:
|
||||
grammar_obj = self.grammar_backend.init_strict_reasoning_grammar(
|
||||
req.require_reasoning
|
||||
)
|
||||
if grammar_obj is not None:
|
||||
req.grammar = grammar_obj
|
||||
self._apply_request_reasoning_budget(req)
|
||||
self._apply_request_reasoning_config(req)
|
||||
|
||||
if add_to_grammar_queue:
|
||||
self.grammar_queue.append(req)
|
||||
@@ -284,7 +297,7 @@ class GrammarManager:
|
||||
)
|
||||
req.grammar = InvalidGrammarObject(f"Grammar compilation failed: {e}")
|
||||
self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
|
||||
self._apply_request_reasoning_budget(req)
|
||||
self._apply_request_reasoning_config(req)
|
||||
if isinstance(req.grammar, InvalidGrammarObject):
|
||||
error_msg = f"Failed to compile {req.grammar_key[0]} grammar: {req.grammar.error_message}"
|
||||
req.set_finish_with_abort(error_msg)
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
"""The baseclass of a backend for reasoner grammar-guided constrained decoding."""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional, Tuple, Union
|
||||
from typing import List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import torch
|
||||
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
|
||||
@@ -38,7 +38,7 @@ class ReasonerGrammarObject(BaseGrammarObject):
|
||||
def __init__(
|
||||
self,
|
||||
grammar: Optional[BaseGrammarObject],
|
||||
think_end_ids: List[int],
|
||||
think_end_ids: Sequence[int],
|
||||
think_excluded_token_ids: Optional[List[int]] = None,
|
||||
max_think_tokens: int = -1,
|
||||
enable_token_filter: bool = False,
|
||||
@@ -74,6 +74,26 @@ class ReasonerGrammarObject(BaseGrammarObject):
|
||||
self.tokens_in_think = -1
|
||||
self.tokens_after_end = 0
|
||||
|
||||
def set_request_think_end_ids(self, think_end_ids: Sequence[int]) -> None:
|
||||
"""Select one request-specific reasoning terminator before decoding."""
|
||||
token_ids = tuple(think_end_ids)
|
||||
if not token_ids:
|
||||
raise ValueError("request reasoning terminator must not be empty")
|
||||
if any(type(token_id) is not int or token_id < 0 for token_id in token_ids):
|
||||
raise ValueError("request reasoning terminator must contain token IDs")
|
||||
if (
|
||||
self.current_token is not None
|
||||
or self._thinking_match_history
|
||||
or (self.tokens_in_think, self.tokens_after_end) not in ((0, -1), (-1, 0))
|
||||
):
|
||||
raise RuntimeError(
|
||||
"request reasoning terminator must be selected before decoding"
|
||||
)
|
||||
|
||||
self.think_end_ids = token_ids
|
||||
self._think_end_matcher = TokenSequenceMatcher(token_ids)
|
||||
self._matched_think_end_tokens = 0
|
||||
|
||||
def _is_thinking(self):
|
||||
return self.tokens_in_think >= 0 and self.tokens_after_end == -1
|
||||
|
||||
|
||||
@@ -2157,6 +2157,16 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
|
||||
else:
|
||||
committed_output_id = output_id[0].item()
|
||||
decode_req.req.output_ids.append(committed_output_id)
|
||||
if not replayed_boundary:
|
||||
# The handoff token is generated on the prefill worker, so it does
|
||||
# not pass through the decode worker's normal batch-result path.
|
||||
# Account for it here using the same request-selected reasoning
|
||||
# terminator matcher as subsequent decode tokens. A rebootstrap
|
||||
# boundary has already been accounted for before retraction and
|
||||
# must not be consumed twice.
|
||||
self.scheduler.batch_result_processor._maybe_update_reasoning_tokens(
|
||||
decode_req.req, committed_output_id
|
||||
)
|
||||
decode_req.req.cached_tokens = cached_tokens[0].item()
|
||||
# The prefill node already reported its prefix-cache hit in
|
||||
# cached_tokens[0]. Seed already_computed with it so that
|
||||
|
||||
@@ -2029,6 +2029,7 @@ class MessageProcessingResult:
|
||||
tool_call_constraint: Optional[ToolCallConstraint] = None
|
||||
skip_special_tokens: bool = True
|
||||
require_reasoning: bool = False
|
||||
reasoning_end_token_ids: Optional[List[int]] = None
|
||||
|
||||
|
||||
class ToolCallProcessingResult(NamedTuple):
|
||||
|
||||
@@ -94,6 +94,9 @@ from sglang.srt.parser.jinja_template_utils import (
|
||||
process_content_for_template_format,
|
||||
)
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
from sglang.srt.sampling.sampling_params import (
|
||||
set_request_reasoning_end_token_ids,
|
||||
)
|
||||
from sglang.srt.utils.weight_versions import build_endpoint_weight_version_metadata
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -1073,6 +1076,9 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
tool_call_constraint=processed_messages.tool_call_constraint,
|
||||
renderer_handles_response_format=self.chat_encoding_spec == "kimi_k3",
|
||||
)
|
||||
set_request_reasoning_end_token_ids(
|
||||
sampling_params, processed_messages.reasoning_end_token_ids
|
||||
)
|
||||
|
||||
# Handle single vs multiple requests
|
||||
if request.input_ids is not None:
|
||||
@@ -1264,6 +1270,31 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
result.tool_call_constraint = tool_call_constraint
|
||||
result.require_reasoning = thinking_mode
|
||||
result.skip_special_tokens = request.skip_special_tokens
|
||||
if self.reasoning_parser == "k2_horizon" and thinking_mode:
|
||||
parser = ReasoningParser(
|
||||
model_type=self.reasoning_parser,
|
||||
stream_reasoning=False,
|
||||
force_reasoning=True,
|
||||
request=request,
|
||||
tokenizer=self.tokenizer_manager.tokenizer,
|
||||
)
|
||||
token_ids = self.tokenizer_manager.tokenizer.encode(
|
||||
parser.detector.think_end_token,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
if hasattr(token_ids, "tolist"):
|
||||
token_ids = token_ids.tolist()
|
||||
if (
|
||||
not isinstance(token_ids, list)
|
||||
or not token_ids
|
||||
or any(
|
||||
type(token_id) is not int or token_id < 0 for token_id in token_ids
|
||||
)
|
||||
):
|
||||
raise ValueError(
|
||||
"The selected K2 reasoning terminator could not be encoded"
|
||||
)
|
||||
result.reasoning_end_token_ids = list(token_ids)
|
||||
return result
|
||||
|
||||
def _apply_jinja_template(
|
||||
|
||||
@@ -73,6 +73,9 @@ from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
from sglang.srt.function_call.json_array_parser import JsonArrayParser
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
from sglang.srt.sampling.sampling_params import (
|
||||
set_request_reasoning_end_token_ids,
|
||||
)
|
||||
from sglang.srt.utils import random_uuid
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -393,6 +396,11 @@ class OpenAIServingResponses(OpenAIServingChat):
|
||||
else None
|
||||
),
|
||||
)
|
||||
if processed_messages is not None:
|
||||
set_request_reasoning_end_token_ids(
|
||||
sampling_params,
|
||||
processed_messages.reasoning_end_token_ids,
|
||||
)
|
||||
# _process_messages set skip_special_tokens on a chat_request
|
||||
# we then discard, so re-apply it to the engine sampling dict.
|
||||
if processed_messages is not None and (
|
||||
@@ -591,6 +599,15 @@ class OpenAIServingResponses(OpenAIServingChat):
|
||||
|
||||
is_multimodal = self.tokenizer_manager.model_config.is_multimodal
|
||||
processed_messages = self._process_messages(chat_request, is_multimodal)
|
||||
# ``_process_messages`` merges server defaults into the temporary Chat
|
||||
# request before rendering. Response parsing happens later from the
|
||||
# original request, so carry over the exact template kwargs that selected
|
||||
# the wire-format delimiters.
|
||||
request.chat_template_kwargs = (
|
||||
dict(chat_request.chat_template_kwargs)
|
||||
if chat_request.chat_template_kwargs is not None
|
||||
else None
|
||||
)
|
||||
|
||||
if is_multimodal:
|
||||
request_prompts = [processed_messages.prompt]
|
||||
|
||||
@@ -29,6 +29,7 @@ from sglang.srt.function_call.hermes_detector import HermesDetector
|
||||
from sglang.srt.function_call.hunyuan_detector import HunyuanDetector
|
||||
from sglang.srt.function_call.inkling_detector import InklingDetector
|
||||
from sglang.srt.function_call.internlm_detector import InternlmDetector
|
||||
from sglang.srt.function_call.k2_v3_detector import K2V3Detector
|
||||
from sglang.srt.function_call.kimik2_detector import KimiK2Detector
|
||||
from sglang.srt.function_call.kimik3_detector import KimiK3Detector
|
||||
from sglang.srt.function_call.lfm2_detector import Lfm2Detector
|
||||
@@ -76,6 +77,7 @@ class FunctionCallParser:
|
||||
"glm45": Glm4MoeDetector,
|
||||
"glm47": Glm47MoeDetector,
|
||||
"gpt-oss": GptOssDetector,
|
||||
"k2_horizon": K2V3Detector,
|
||||
"kimi_k2": KimiK2Detector,
|
||||
"kimi_k3": KimiK3Detector,
|
||||
"lfm2": Lfm2Detector,
|
||||
|
||||
@@ -0,0 +1,593 @@
|
||||
"""Tool-call parsing for the canonical K2 Horizon (K2-v3) IFM format."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import Tool
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.function_call.base_format_detector import BaseFormatDetector
|
||||
from sglang.srt.function_call.core_types import (
|
||||
StreamingParseResult,
|
||||
ToolCallItem,
|
||||
_GetInfoFunc,
|
||||
)
|
||||
from sglang.srt.function_call.utils import (
|
||||
get_schema_properties,
|
||||
infer_type_from_json_schema,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class K2V3Detector(BaseFormatDetector):
|
||||
"""Parse K2 Horizon's canonical ``<ifm|...>`` tool-call format.
|
||||
|
||||
The wire content is self-describing: XML uses key/value tags, typed XML
|
||||
adds argument-type tags, and JSON begins with an object or array.
|
||||
|
||||
Streaming keeps a partial block atomic and emits it as soon as its closing
|
||||
tag arrives. This deliberately favors correctness across arbitrarily split
|
||||
tags over emitting incomplete JSON argument fragments.
|
||||
"""
|
||||
|
||||
_TOOL_CALLS_START = "<ifm|tool_calls>"
|
||||
_TOOL_CALLS_END = "</ifm|tool_calls>"
|
||||
_TOOL_CALL_START = "<ifm|tool_call>"
|
||||
_TOOL_CALL_END = "</ifm|tool_call>"
|
||||
|
||||
_ARG_KEY_START = "<ifm|arg_key>"
|
||||
|
||||
_THINK_PAIRS = (
|
||||
("<ifm|think>", "</ifm|think>"),
|
||||
("<ifm|think_fast>", "</ifm|think_fast>"),
|
||||
("<ifm|think_faster>", "</ifm|think_faster>"),
|
||||
)
|
||||
_GROUP_TOKENS = (_TOOL_CALLS_START, _TOOL_CALLS_END)
|
||||
_STREAM_MARKERS = (
|
||||
_TOOL_CALL_START,
|
||||
_TOOL_CALLS_START,
|
||||
_TOOL_CALLS_END,
|
||||
"<ifm|think>",
|
||||
"</ifm|think>",
|
||||
"<ifm|think_fast>",
|
||||
"</ifm|think_fast>",
|
||||
"<ifm|think_faster>",
|
||||
"</ifm|think_faster>",
|
||||
)
|
||||
|
||||
_ARG_PATTERN = re.compile(
|
||||
r"<ifm\|arg_key>(.*?)</ifm\|arg_key>\s*"
|
||||
r"(?:<ifm\|arg_type>(.*?)</ifm\|arg_type>\s*)?"
|
||||
r"<ifm\|arg_value>(.*?)</ifm\|arg_value>",
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.bot_token = self._TOOL_CALL_START
|
||||
self.eot_token = self._TOOL_CALL_END
|
||||
self._next_tool_index = 0
|
||||
self._at_stream_start = True
|
||||
self._inside_tool_group = False
|
||||
|
||||
def has_tool_call(self, text: str) -> bool:
|
||||
return self._TOOL_CALL_START in text
|
||||
|
||||
def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
|
||||
if not self.has_tool_call(text):
|
||||
return StreamingParseResult(normal_text=text)
|
||||
|
||||
# Reasoning extraction belongs to ReasoningParser. Keep a leading IFM
|
||||
# reasoning block out of tool parsing, but pass it through unchanged so
|
||||
# separate_reasoning=False returns the original assistant output.
|
||||
reasoning_prefix, working = self._split_leading_reasoning(text)
|
||||
|
||||
# Canonical template output wraps all calls in one group. Whitespace
|
||||
# inside that wrapper is formatting, while the prefix immediately
|
||||
# before it (notably the newline after </ifm|think>) is user-visible.
|
||||
group_start = working.find(self._TOOL_CALLS_START)
|
||||
if group_start != -1:
|
||||
group_body_start = group_start + len(self._TOOL_CALLS_START)
|
||||
group_end = working.find(self._TOOL_CALLS_END, group_body_start)
|
||||
if group_end != -1:
|
||||
parsed_group = self._parse_complete_region(
|
||||
working[group_body_start:group_end], tools
|
||||
)
|
||||
if parsed_group is None:
|
||||
# The release template emits one canonical group. Avoid a
|
||||
# partial recovery that silently drops its wrapper or only
|
||||
# some malformed calls: surface the wire text intact.
|
||||
return StreamingParseResult(normal_text=reasoning_prefix + working)
|
||||
normal_text = (
|
||||
working[:group_start]
|
||||
+ working[group_end + len(self._TOOL_CALLS_END) :]
|
||||
)
|
||||
return StreamingParseResult(
|
||||
normal_text=reasoning_prefix + normal_text,
|
||||
calls=self._items_from_calls(parsed_group, tools, 0),
|
||||
)
|
||||
|
||||
calls: list[ToolCallItem] = []
|
||||
normal_parts: list[str] = []
|
||||
cursor = 0
|
||||
|
||||
while True:
|
||||
start = working.find(self._TOOL_CALL_START, cursor)
|
||||
if start == -1:
|
||||
normal_parts.append(working[cursor:])
|
||||
break
|
||||
|
||||
normal_parts.append(working[cursor:start])
|
||||
body_start = start + len(self._TOOL_CALL_START)
|
||||
end = working.find(self._TOOL_CALL_END, body_start)
|
||||
if end == -1:
|
||||
# Preserve an unterminated block verbatim instead of silently
|
||||
# dropping the model's output.
|
||||
normal_parts.append(working[start:])
|
||||
break
|
||||
|
||||
raw_block = working[start : end + len(self._TOOL_CALL_END)]
|
||||
parsed = self._parse_block(working[body_start:end], tools)
|
||||
if parsed is None:
|
||||
normal_parts.append(raw_block)
|
||||
else:
|
||||
calls.extend(self._items_from_calls(parsed, tools, len(calls)))
|
||||
cursor = end + len(self._TOOL_CALL_END)
|
||||
|
||||
normal_text = reasoning_prefix + self._strip_group_tokens("".join(normal_parts))
|
||||
return StreamingParseResult(normal_text=normal_text, calls=calls)
|
||||
|
||||
def parse_streaming_increment(
|
||||
self, new_text: str, tools: List[Tool]
|
||||
) -> StreamingParseResult:
|
||||
self._buffer += new_text
|
||||
return self._drain_stream(tools, flush=False)
|
||||
|
||||
def finish(self, tools: List[Tool]) -> StreamingParseResult:
|
||||
return self._drain_stream(tools, flush=True)
|
||||
|
||||
def _drain_stream(self, tools: List[Tool], *, flush: bool) -> StreamingParseResult:
|
||||
normal_parts: list[str] = []
|
||||
calls: list[ToolCallItem] = []
|
||||
|
||||
while self._buffer:
|
||||
if self._at_stream_start:
|
||||
state, reasoning_prefix = self._consume_stream_reasoning_prefix(
|
||||
flush=flush
|
||||
)
|
||||
if state == "hold":
|
||||
break
|
||||
if state == "consumed":
|
||||
normal_parts.append(reasoning_prefix)
|
||||
continue
|
||||
self._at_stream_start = False
|
||||
|
||||
start = self._buffer.find(self._TOOL_CALL_START)
|
||||
if self._inside_tool_group:
|
||||
# The line breaks surrounding calls are part of IFM framing,
|
||||
# not assistant content. Hold/drop them with the wrapper.
|
||||
self._buffer = self._buffer.lstrip()
|
||||
if self._buffer.startswith(self._TOOL_CALLS_END):
|
||||
self._buffer = self._buffer[len(self._TOOL_CALLS_END) :]
|
||||
self._inside_tool_group = False
|
||||
continue
|
||||
if (
|
||||
self._TOOL_CALLS_END.startswith(self._buffer)
|
||||
and self._buffer != self._TOOL_CALLS_END
|
||||
and not flush
|
||||
):
|
||||
break
|
||||
start = self._buffer.find(self._TOOL_CALL_START)
|
||||
else:
|
||||
group_start = self._buffer.find(self._TOOL_CALLS_START)
|
||||
if group_start != -1 and (start == -1 or group_start < start):
|
||||
normal_parts.append(self._buffer[:group_start])
|
||||
self._buffer = self._buffer[
|
||||
group_start + len(self._TOOL_CALLS_START) :
|
||||
]
|
||||
self._inside_tool_group = True
|
||||
continue
|
||||
|
||||
if start == -1:
|
||||
visible, hold = self._split_visible_stream_text(
|
||||
self._buffer, flush=flush
|
||||
)
|
||||
normal_parts.append(visible)
|
||||
self._buffer = hold
|
||||
break
|
||||
|
||||
prefix = self._buffer[:start]
|
||||
normal_parts.append(self._strip_group_tokens(prefix))
|
||||
self._buffer = self._buffer[start:]
|
||||
|
||||
body_start = len(self._TOOL_CALL_START)
|
||||
end = self._buffer.find(self._TOOL_CALL_END, body_start)
|
||||
if end == -1:
|
||||
if flush:
|
||||
normal_parts.append(self._buffer)
|
||||
self._buffer = ""
|
||||
break
|
||||
|
||||
raw_block = self._buffer[: end + len(self._TOOL_CALL_END)]
|
||||
parsed = self._parse_block(self._buffer[body_start:end], tools)
|
||||
if parsed is None:
|
||||
normal_parts.append(raw_block)
|
||||
else:
|
||||
new_items = self._items_from_calls(parsed, tools, self._next_tool_index)
|
||||
calls.extend(new_items)
|
||||
self._next_tool_index += len(new_items)
|
||||
self._buffer = self._buffer[end + len(self._TOOL_CALL_END) :]
|
||||
|
||||
return StreamingParseResult(normal_text="".join(normal_parts), calls=calls)
|
||||
|
||||
def _parse_complete_region(
|
||||
self, text: str, tools: List[Tool]
|
||||
) -> Optional[list[tuple[str, dict[str, Any]]]]:
|
||||
"""Parse a canonical tool-call group, requiring full framing.
|
||||
|
||||
Only whitespace may occur between calls. Any malformed or unterminated
|
||||
block fails the group atomically so its raw text can be surfaced.
|
||||
"""
|
||||
|
||||
parsed: list[tuple[str, dict[str, Any]]] = []
|
||||
cursor = 0
|
||||
while cursor < len(text):
|
||||
start = text.find(self._TOOL_CALL_START, cursor)
|
||||
if start == -1:
|
||||
return parsed if parsed and not text[cursor:].strip() else None
|
||||
if text[cursor:start].strip():
|
||||
return None
|
||||
body_start = start + len(self._TOOL_CALL_START)
|
||||
end = text.find(self._TOOL_CALL_END, body_start)
|
||||
if end == -1:
|
||||
return None
|
||||
block_calls = self._parse_block(text[body_start:end], tools)
|
||||
if block_calls is None:
|
||||
return None
|
||||
parsed.extend(block_calls)
|
||||
cursor = end + len(self._TOOL_CALL_END)
|
||||
return parsed or None
|
||||
|
||||
def _consume_stream_reasoning_prefix(self, *, flush: bool) -> tuple[str, str]:
|
||||
"""Separate a leading reasoning block from tool parsing without loss."""
|
||||
|
||||
stripped = self._buffer.lstrip()
|
||||
whitespace_len = len(self._buffer) - len(stripped)
|
||||
for start_token, end_token in self._THINK_PAIRS:
|
||||
if (
|
||||
stripped
|
||||
and start_token.startswith(stripped)
|
||||
and stripped != start_token
|
||||
):
|
||||
if not flush:
|
||||
return "hold", ""
|
||||
passthrough = self._buffer
|
||||
self._buffer = ""
|
||||
return "consumed", passthrough
|
||||
if not stripped.startswith(start_token):
|
||||
continue
|
||||
end = stripped.find(end_token, len(start_token))
|
||||
if end == -1:
|
||||
if not flush:
|
||||
return "hold", ""
|
||||
passthrough = self._buffer
|
||||
self._buffer = ""
|
||||
return "consumed", passthrough
|
||||
consumed_len = whitespace_len + end + len(end_token)
|
||||
passthrough = self._buffer[:consumed_len]
|
||||
self._buffer = self._buffer[consumed_len:]
|
||||
return "consumed", passthrough
|
||||
|
||||
# Whitespace at the beginning may precede a reasoning token in the
|
||||
# next chunk, so keep it until that distinction is observable.
|
||||
if not stripped and whitespace_len and not flush:
|
||||
return "hold", ""
|
||||
return "none", ""
|
||||
|
||||
def _split_visible_stream_text(self, text: str, *, flush: bool) -> tuple[str, str]:
|
||||
if flush:
|
||||
return self._strip_group_tokens(text), ""
|
||||
|
||||
partial_len = max(
|
||||
(
|
||||
self._partial_marker_suffix_len(text, marker)
|
||||
for marker in self._STREAM_MARKERS
|
||||
),
|
||||
default=0,
|
||||
)
|
||||
if partial_len:
|
||||
visible, hold = text[:-partial_len], text[-partial_len:]
|
||||
else:
|
||||
visible, hold = text, ""
|
||||
return self._strip_group_tokens(visible), hold
|
||||
|
||||
@staticmethod
|
||||
def _partial_marker_suffix_len(text: str, marker: str) -> int:
|
||||
for size in range(min(len(text), len(marker) - 1), 0, -1):
|
||||
if marker.startswith(text[-size:]):
|
||||
return size
|
||||
return 0
|
||||
|
||||
@classmethod
|
||||
def _strip_group_tokens(cls, text: str) -> str:
|
||||
for token in cls._GROUP_TOKENS:
|
||||
text = text.replace(token, "")
|
||||
return text
|
||||
|
||||
@classmethod
|
||||
def _split_leading_reasoning(cls, text: str) -> tuple[str, str]:
|
||||
"""Return an exact leading reasoning prefix and the parseable suffix."""
|
||||
|
||||
cursor = 0
|
||||
while cursor < len(text):
|
||||
remainder = text[cursor:]
|
||||
stripped = remainder.lstrip()
|
||||
whitespace_len = len(remainder) - len(stripped)
|
||||
matched_pair = next(
|
||||
(
|
||||
(start_token, end_token)
|
||||
for start_token, end_token in cls._THINK_PAIRS
|
||||
if stripped.startswith(start_token)
|
||||
),
|
||||
None,
|
||||
)
|
||||
if matched_pair is None:
|
||||
# A truncated reasoning opener is still passthrough content;
|
||||
# never reinterpret tool-looking text following it.
|
||||
if stripped and any(
|
||||
start_token.startswith(stripped)
|
||||
for start_token, _ in cls._THINK_PAIRS
|
||||
):
|
||||
return text, ""
|
||||
return text[:cursor], text[cursor:]
|
||||
|
||||
start_token, end_token = matched_pair
|
||||
end = stripped.find(end_token, len(start_token))
|
||||
if end == -1:
|
||||
return text, ""
|
||||
cursor += whitespace_len + end + len(end_token)
|
||||
|
||||
return text, ""
|
||||
|
||||
def _parse_block(
|
||||
self, block: str, tools: List[Tool]
|
||||
) -> Optional[list[tuple[str, dict[str, Any]]]]:
|
||||
# The wire framing is self-describing: JSON begins with an object or
|
||||
# array, while XML and typed XML share the same tag parser.
|
||||
stripped = block.strip()
|
||||
looks_json = stripped.startswith(("{", "["))
|
||||
try:
|
||||
if looks_json:
|
||||
return self._parse_json_block(stripped, tools)
|
||||
return self._parse_xml_block(block, tools)
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
logger.warning("Malformed K2-v3 tool-call block; forwarding it as text")
|
||||
return None
|
||||
|
||||
def _parse_json_block(
|
||||
self, block: str, tools: List[Tool]
|
||||
) -> list[tuple[str, dict[str, Any]]]:
|
||||
payload = json.loads(block)
|
||||
raw_calls = payload if isinstance(payload, list) else [payload]
|
||||
parsed: list[tuple[str, dict[str, Any]]] = []
|
||||
for raw_call in raw_calls:
|
||||
if not isinstance(raw_call, Mapping):
|
||||
raise ValueError("K2-v3 JSON tool call must be an object")
|
||||
function = raw_call.get("function", raw_call)
|
||||
if (
|
||||
not isinstance(function, Mapping)
|
||||
or not isinstance(function.get("name"), str)
|
||||
or not function["name"]
|
||||
):
|
||||
raise ValueError("K2-v3 JSON tool call is missing a function name")
|
||||
name = function["name"]
|
||||
arguments = function.get("arguments", function.get("parameters", {}))
|
||||
if isinstance(arguments, str):
|
||||
arguments = json.loads(arguments) if arguments.strip() else {}
|
||||
if arguments is None:
|
||||
arguments = {}
|
||||
if not isinstance(arguments, Mapping):
|
||||
raise ValueError("K2-v3 tool-call arguments must be an object")
|
||||
parsed.append(
|
||||
(
|
||||
name,
|
||||
self._coerce_arguments(
|
||||
name, dict(arguments), tools, from_text=False
|
||||
),
|
||||
)
|
||||
)
|
||||
return parsed
|
||||
|
||||
def _parse_xml_block(
|
||||
self, block: str, tools: List[Tool]
|
||||
) -> list[tuple[str, dict[str, Any]]]:
|
||||
first_arg = block.find(self._ARG_KEY_START)
|
||||
name = (block if first_arg == -1 else block[:first_arg]).strip()
|
||||
if not name:
|
||||
raise ValueError("K2-v3 XML tool call is missing a function name")
|
||||
if first_arg == -1:
|
||||
return [(name, {})]
|
||||
|
||||
argument_text = block[first_arg:]
|
||||
arguments: dict[str, Any] = {}
|
||||
cursor = 0
|
||||
for match in self._ARG_PATTERN.finditer(argument_text):
|
||||
if argument_text[cursor : match.start()].strip():
|
||||
raise ValueError("Malformed K2-v3 XML argument framing")
|
||||
key = match.group(1).strip()
|
||||
if not key:
|
||||
raise ValueError("K2-v3 XML argument is missing a key")
|
||||
inline_type = match.group(2).strip() if match.group(2) else None
|
||||
arguments[key] = self._coerce_value(
|
||||
match.group(3),
|
||||
name,
|
||||
key,
|
||||
tools,
|
||||
inline_type=inline_type,
|
||||
from_text=True,
|
||||
)
|
||||
cursor = match.end()
|
||||
if argument_text[cursor:].strip():
|
||||
raise ValueError("Malformed K2-v3 XML argument framing")
|
||||
return [(name, arguments)]
|
||||
|
||||
def _items_from_calls(
|
||||
self,
|
||||
parsed: list[tuple[str, dict[str, Any]]],
|
||||
tools: List[Tool],
|
||||
start_index: int,
|
||||
) -> list[ToolCallItem]:
|
||||
known_names = {tool.function.name for tool in tools}
|
||||
items: list[ToolCallItem] = []
|
||||
for name, arguments in parsed:
|
||||
if name not in known_names and not envs.SGLANG_FORWARD_UNKNOWN_TOOLS.get():
|
||||
logger.warning("Model attempted to call undefined function: %s", name)
|
||||
continue
|
||||
items.append(
|
||||
ToolCallItem(
|
||||
tool_index=start_index + len(items),
|
||||
name=name,
|
||||
parameters=json.dumps(arguments, ensure_ascii=False),
|
||||
)
|
||||
)
|
||||
return items
|
||||
|
||||
@staticmethod
|
||||
def _resolve_local_ref(schema: Any, root_schema: Any) -> Optional[dict[str, Any]]:
|
||||
"""Resolve the local references emitted by the K2 tool template.
|
||||
|
||||
The release templates inline ``#/$defs/...`` and ``#/definitions/...``
|
||||
before presenting a tool to the model. Parsing must inspect that same
|
||||
effective schema; otherwise a referenced string such as ``"00123"``
|
||||
is incorrectly coerced to an integer.
|
||||
"""
|
||||
|
||||
if not isinstance(schema, Mapping):
|
||||
return None
|
||||
resolved = dict(schema)
|
||||
seen: set[str] = set()
|
||||
while isinstance(resolved.get("$ref"), str):
|
||||
ref = resolved["$ref"]
|
||||
if ref in seen:
|
||||
return None
|
||||
seen.add(ref)
|
||||
if ref.startswith("#/$defs/"):
|
||||
definitions = (
|
||||
root_schema.get("$defs", {})
|
||||
if isinstance(root_schema, Mapping)
|
||||
else {}
|
||||
)
|
||||
name = ref.removeprefix("#/$defs/")
|
||||
elif ref.startswith("#/definitions/"):
|
||||
definitions = (
|
||||
root_schema.get("definitions", {})
|
||||
if isinstance(root_schema, Mapping)
|
||||
else {}
|
||||
)
|
||||
name = ref.removeprefix("#/definitions/")
|
||||
else:
|
||||
return resolved
|
||||
target = definitions.get(name) if isinstance(definitions, Mapping) else None
|
||||
if not isinstance(target, Mapping):
|
||||
return resolved
|
||||
# Match the template: annotations at the reference site override
|
||||
# fields from the referenced definition.
|
||||
resolved = {
|
||||
**target,
|
||||
**{key: value for key, value in resolved.items() if key != "$ref"},
|
||||
}
|
||||
return resolved
|
||||
|
||||
@classmethod
|
||||
def _argument_type(
|
||||
cls, tool_name: str, argument_name: str, tools: List[Tool]
|
||||
) -> Optional[str]:
|
||||
for tool in tools:
|
||||
if tool.function.name != tool_name:
|
||||
continue
|
||||
root_schema = tool.function.parameters
|
||||
schema = get_schema_properties(root_schema).get(argument_name)
|
||||
schema = cls._resolve_local_ref(schema, root_schema)
|
||||
if not isinstance(schema, dict):
|
||||
return None
|
||||
|
||||
# For unions the canonical xml_typed template emits the actual
|
||||
# argument value type. Do not guess one branch and override that
|
||||
# wire-level type hint. Untyped XML can still infer from its value,
|
||||
# and JSON values have already been decoded with their true type.
|
||||
schema_type = schema.get("type")
|
||||
if "anyOf" in schema or "oneOf" in schema:
|
||||
return None
|
||||
if (
|
||||
isinstance(schema_type, list)
|
||||
and len([item for item in schema_type if item != "null"]) > 1
|
||||
):
|
||||
return None
|
||||
return infer_type_from_json_schema(schema)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _coerce_value(
|
||||
cls,
|
||||
value: Any,
|
||||
tool_name: str,
|
||||
argument_name: str,
|
||||
tools: List[Tool],
|
||||
*,
|
||||
inline_type: Optional[str] = None,
|
||||
from_text: bool,
|
||||
) -> Any:
|
||||
target_type = cls._argument_type(tool_name, argument_name, tools) or inline_type
|
||||
if target_type == "any":
|
||||
return value
|
||||
if target_type == "string":
|
||||
return (
|
||||
value
|
||||
if isinstance(value, str)
|
||||
else json.dumps(value, ensure_ascii=False)
|
||||
)
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
if target_type is None:
|
||||
# JSON strings are already unambiguously strings after json.loads.
|
||||
# Only the XML dialect needs best-effort JSON decoding of an
|
||||
# untyped textual value.
|
||||
if not from_text:
|
||||
return value
|
||||
try:
|
||||
return json.loads(value.strip())
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
try:
|
||||
return json.loads(value.strip())
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def _coerce_arguments(
|
||||
cls,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
tools: List[Tool],
|
||||
*,
|
||||
from_text: bool,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
key: cls._coerce_value(value, tool_name, key, tools, from_text=from_text)
|
||||
for key, value in arguments.items()
|
||||
}
|
||||
|
||||
def supports_structural_tag(self) -> bool:
|
||||
# XML arguments are not a JSON-schema body. Required/named tool choice
|
||||
# therefore uses SGLang's standard JSON constraint and JsonArrayParser.
|
||||
return False
|
||||
|
||||
def structure_info(self) -> _GetInfoFunc:
|
||||
raise NotImplementedError(
|
||||
"K2-v3 native IFM does not use legacy structural tags"
|
||||
)
|
||||
@@ -0,0 +1,310 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# 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 primitives for mixture-of-value attention (MoVA)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from sglang.srt.utils import set_weight_attrs
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
_ROUTED_LINEAR_CHUNK_SIZE = 64 * 1024
|
||||
|
||||
|
||||
def _prepare_mova_moe_config(config: dict) -> dict:
|
||||
"""Copy a generic MoE config and remove options MoVA cannot consume."""
|
||||
|
||||
config = dict(config)
|
||||
# The ordinary two-GEMM MoE runner handles USE_TMA separately and removes
|
||||
# it before launching the Triton kernel. MoVA reuses only the first GEMM
|
||||
# and does not construct TMA descriptors, so forwarding this tuning-only
|
||||
# key as a kernel constexpr would fail at launch.
|
||||
config.pop("USE_TMA", None)
|
||||
return config
|
||||
|
||||
|
||||
def mova_router_topk(
|
||||
router_logits: torch.Tensor,
|
||||
router_bias: Optional[torch.Tensor],
|
||||
*,
|
||||
score_func: str,
|
||||
top_k: int,
|
||||
scaling_factor: float,
|
||||
renormalize: bool = True,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply xLLM's selection-only router-bias semantics.
|
||||
|
||||
Scores are computed in fp32. ``router_bias`` changes which value experts
|
||||
are selected, but the mixture coefficients are gathered from the unbiased
|
||||
scores. Scaling happens after optional top-k renormalization.
|
||||
"""
|
||||
|
||||
if top_k <= 0 or top_k > router_logits.shape[-1]:
|
||||
raise ValueError(
|
||||
f"top_k must be in [1, {router_logits.shape[-1]}], got {top_k}"
|
||||
)
|
||||
if router_logits.is_cuda and (score_func == "sigmoid" or router_bias is None):
|
||||
# Reuse SGLang's fused sigmoid/softmax top-k kernels. They implement
|
||||
# the same selection-only correction-bias contract and return fp32
|
||||
# mixture weights plus int32 expert ids.
|
||||
from sglang.srt.layers.moe.topk import fused_topk
|
||||
|
||||
weights, selected = fused_topk(
|
||||
hidden_states=router_logits,
|
||||
gating_output=router_logits,
|
||||
topk=top_k,
|
||||
# Native xLLM leaves a top-1 route at its raw probability.
|
||||
renormalize=renormalize and top_k > 1,
|
||||
correction_bias=router_bias,
|
||||
scoring_func=score_func,
|
||||
)
|
||||
return (weights * scaling_factor).to(router_logits.dtype), selected
|
||||
if score_func == "sigmoid":
|
||||
scores = torch.sigmoid(router_logits.float())
|
||||
elif score_func == "softmax":
|
||||
scores = F.softmax(router_logits, dim=-1, dtype=torch.float32)
|
||||
else:
|
||||
raise ValueError(f"Unsupported MoVA router score function: {score_func}")
|
||||
|
||||
selection_scores = scores
|
||||
if router_bias is not None:
|
||||
selection_scores = selection_scores + router_bias.to(selection_scores)
|
||||
|
||||
selected = torch.topk(selection_scores, top_k, dim=-1).indices
|
||||
weights = torch.gather(scores, dim=-1, index=selected)
|
||||
if renormalize and top_k > 1:
|
||||
weights = weights / weights.sum(dim=-1, keepdim=True)
|
||||
weights = weights * scaling_factor
|
||||
return weights.to(router_logits.dtype), selected.to(torch.int32)
|
||||
|
||||
|
||||
def routed_linear_reference(
|
||||
hidden_states: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Straightforward MoVA value projection used as the correctness oracle."""
|
||||
|
||||
if hidden_states.ndim != 2:
|
||||
raise ValueError("MoVA routed linear expects [tokens, hidden] inputs")
|
||||
if expert_weights.ndim != 3:
|
||||
raise ValueError("MoVA expert weights must be [experts, output, hidden]")
|
||||
if routing_weights.shape != selected_experts.shape:
|
||||
raise ValueError("MoVA routing weights and expert ids must have equal shape")
|
||||
if hidden_states.shape[0] != selected_experts.shape[0]:
|
||||
raise ValueError("MoVA route count must match the token count")
|
||||
if hidden_states.shape[1] != expert_weights.shape[2]:
|
||||
raise ValueError("MoVA input and expert hidden dimensions differ")
|
||||
if hidden_states.shape[0] == 0:
|
||||
return hidden_states.new_empty((0, expert_weights.shape[1]))
|
||||
|
||||
# This deliberately favors clarity over memory use. Production CUDA paths
|
||||
# use ``routed_linear`` below and never materialize selected expert weights.
|
||||
selected_weights = expert_weights[selected_experts.to(torch.long)]
|
||||
projected = torch.einsum("mknh,mh->mkn", selected_weights, hidden_states)
|
||||
projected = F.silu(projected)
|
||||
return (projected * routing_weights.to(projected).unsqueeze(-1)).sum(dim=1)
|
||||
|
||||
|
||||
def _fake_mova_routed_linear_cuda(
|
||||
hidden_states: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return hidden_states.new_empty((hidden_states.shape[0], expert_weights.shape[1]))
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
op_name="mova_routed_linear_cuda",
|
||||
fake_impl=_fake_mova_routed_linear_cuda,
|
||||
)
|
||||
def _routed_linear_cuda_chunk(
|
||||
hidden_states: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# Keep these imports local: CPU model inspection and mapping tests should
|
||||
# not initialize Triton or require the CUDA extension.
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.kernels.ops.moe.fused_moe_triton_kernels import invoke_fused_moe_kernel
|
||||
from sglang.srt.layers.moe.fused_moe_triton import (
|
||||
moe_align_block_size,
|
||||
try_get_optimal_moe_config,
|
||||
)
|
||||
|
||||
num_tokens = hidden_states.shape[0]
|
||||
top_k = selected_experts.shape[1]
|
||||
num_experts, output_size, input_size = expert_weights.shape
|
||||
# ``try_get_optimal_moe_config`` uses the last dimension of its synthetic
|
||||
# second-GEMM shape as N. MoVA has no second GEMM, so describe the desired
|
||||
# routed projection output explicitly.
|
||||
config = _prepare_mova_moe_config(
|
||||
try_get_optimal_moe_config(
|
||||
expert_weights.shape,
|
||||
(num_experts, input_size, output_size),
|
||||
top_k,
|
||||
None,
|
||||
num_tokens,
|
||||
)
|
||||
)
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
|
||||
selected_experts, config["BLOCK_SIZE_M"], num_experts
|
||||
)
|
||||
projected = torch.empty(
|
||||
(num_tokens * top_k, output_size),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
|
||||
invoke_fused_moe_kernel(
|
||||
hidden_states,
|
||||
expert_weights,
|
||||
None,
|
||||
projected,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
routing_weights,
|
||||
selected_experts,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
False, # Routing weights are applied after SiLU.
|
||||
top_k,
|
||||
config,
|
||||
compute_type=compute_type,
|
||||
use_fp8_w8a8=False,
|
||||
use_int8_w8a8=False,
|
||||
use_int8_w8a16=False,
|
||||
use_int4_w4a16=False,
|
||||
per_channel_quant=False,
|
||||
filter_expert=False,
|
||||
)
|
||||
projected = F.silu(projected.view(num_tokens, top_k, output_size))
|
||||
return (projected * routing_weights.to(projected).unsqueeze(-1)).sum(dim=1)
|
||||
|
||||
|
||||
def routed_linear(
|
||||
hidden_states: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Run routed value projections using SGLang's fused-MoE first GEMM."""
|
||||
|
||||
if not hidden_states.is_cuda:
|
||||
return routed_linear_reference(
|
||||
hidden_states, expert_weights, routing_weights, selected_experts
|
||||
)
|
||||
if hidden_states.dtype not in (torch.float16, torch.bfloat16):
|
||||
raise ValueError("Fused MoVA routed linear supports fp16 and bf16 only")
|
||||
if hidden_states.shape[0] == 0:
|
||||
return hidden_states.new_empty((0, expert_weights.shape[1]))
|
||||
if not hidden_states.is_contiguous() or not expert_weights.is_contiguous():
|
||||
raise ValueError("Fused MoVA inputs and expert weights must be contiguous")
|
||||
|
||||
outputs = []
|
||||
for begin in range(0, hidden_states.shape[0], _ROUTED_LINEAR_CHUNK_SIZE):
|
||||
end = min(begin + _ROUTED_LINEAR_CHUNK_SIZE, hidden_states.shape[0])
|
||||
outputs.append(
|
||||
_routed_linear_cuda_chunk(
|
||||
hidden_states[begin:end],
|
||||
expert_weights,
|
||||
routing_weights[begin:end],
|
||||
selected_experts[begin:end],
|
||||
)
|
||||
)
|
||||
return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0)
|
||||
|
||||
|
||||
class RoutedValueExperts(nn.Module):
|
||||
"""Persistent output-sharded MoVA value-expert weights."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_experts: int,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
*,
|
||||
tp_rank: int,
|
||||
tp_size: int,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if output_size % tp_size:
|
||||
raise ValueError(
|
||||
f"MoVA value width {output_size} is not divisible by TP={tp_size}"
|
||||
)
|
||||
self.num_experts = num_experts
|
||||
self.input_size = input_size
|
||||
self.output_size = output_size
|
||||
self.output_size_per_partition = output_size // tp_size
|
||||
self.tp_rank = tp_rank
|
||||
self.weight = nn.Parameter(
|
||||
torch.empty(num_experts, self.output_size_per_partition, input_size),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
|
||||
|
||||
def weight_loader(
|
||||
self,
|
||||
param: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: Optional[int] = None,
|
||||
) -> None:
|
||||
output_begin = self.tp_rank * self.output_size_per_partition
|
||||
if loaded_shard_id is None:
|
||||
expected = (self.num_experts, self.output_size, self.input_size)
|
||||
if tuple(loaded_weight.shape) != expected:
|
||||
raise ValueError(
|
||||
f"Packed MoVA value weight must be {expected}, got "
|
||||
f"{tuple(loaded_weight.shape)}"
|
||||
)
|
||||
local_weight = loaded_weight.narrow(
|
||||
1, output_begin, self.output_size_per_partition
|
||||
)
|
||||
param.data.copy_(local_weight)
|
||||
return
|
||||
|
||||
if not 0 <= loaded_shard_id < self.num_experts:
|
||||
raise ValueError(f"Invalid MoVA value expert id: {loaded_shard_id}")
|
||||
expected = (self.output_size, self.input_size)
|
||||
if tuple(loaded_weight.shape) != expected:
|
||||
raise ValueError(
|
||||
f"MoVA value expert must be {expected}, got {tuple(loaded_weight.shape)}"
|
||||
)
|
||||
local_weight = loaded_weight.narrow(
|
||||
0, output_begin, self.output_size_per_partition
|
||||
)
|
||||
param.data[loaded_shard_id].copy_(local_weight)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
routing_weights: torch.Tensor,
|
||||
selected_experts: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return routed_linear(
|
||||
hidden_states,
|
||||
self.weight,
|
||||
routing_weights,
|
||||
selected_experts,
|
||||
)
|
||||
@@ -925,6 +925,30 @@ class Scheduler(
|
||||
reasoning_parser.detector.think_end_token,
|
||||
)
|
||||
|
||||
selectable_tokens = getattr(
|
||||
reasoning_parser.detector,
|
||||
"request_selectable_think_end_tokens",
|
||||
(),
|
||||
)
|
||||
if selectable_tokens:
|
||||
selectable_sequences = []
|
||||
for end_token in selectable_tokens:
|
||||
token_ids = self.tokenizer.encode(
|
||||
end_token, add_special_tokens=False
|
||||
)
|
||||
if not token_ids:
|
||||
raise ValueError(
|
||||
f"Request-selectable reasoning terminator {end_token!r} "
|
||||
"could not be encoded"
|
||||
)
|
||||
selectable_sequences.append(token_ids)
|
||||
self.model_config.request_selectable_think_end_id_sequences = [
|
||||
list(sequence)
|
||||
for sequence in dict.fromkeys(
|
||||
tuple(sequence) for sequence in selectable_sequences
|
||||
)
|
||||
]
|
||||
|
||||
def init_mamba_backend(self) -> None:
|
||||
if initialize_mamba_selective_state_update_backend is not None:
|
||||
initialize_mamba_selective_state_update_backend(self.server_args)
|
||||
|
||||
@@ -41,6 +41,9 @@ from sglang.srt.runtime_context import (
|
||||
max_speculative_num_draft_tokens,
|
||||
)
|
||||
from sglang.srt.sampling.sampling_observer import CommittedTokens
|
||||
from sglang.srt.sampling.sampling_params import (
|
||||
get_request_reasoning_end_token_ids,
|
||||
)
|
||||
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
|
||||
from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
|
||||
from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
|
||||
@@ -1206,9 +1209,23 @@ class SchedulerBatchResultProcessor:
|
||||
req: Req,
|
||||
next_token_id: Union[int, List[int]],
|
||||
):
|
||||
if not req.require_reasoning:
|
||||
return
|
||||
think_end_ids = self.model_config.think_end_ids
|
||||
if req.require_reasoning and think_end_ids:
|
||||
req.update_reasoning_tokens(next_token_id, think_end_ids)
|
||||
if req._think_end_matcher is None:
|
||||
request_think_end_ids = get_request_reasoning_end_token_ids(
|
||||
req.sampling_params.custom_params,
|
||||
allowed_sequences=getattr(
|
||||
self.model_config,
|
||||
"request_selectable_think_end_id_sequences",
|
||||
None,
|
||||
),
|
||||
)
|
||||
if request_think_end_ids is not None:
|
||||
think_end_ids = request_think_end_ids
|
||||
if not think_end_ids:
|
||||
return
|
||||
req.update_reasoning_tokens(next_token_id, think_end_ids)
|
||||
|
||||
def _mamba_prefix_cache_update(
|
||||
self,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -471,6 +471,63 @@ class KimiK2Detector(BaseReasoningFormatDetector):
|
||||
)
|
||||
|
||||
|
||||
class K2V3Detector(BaseReasoningFormatDetector):
|
||||
"""Reasoning detector for canonical K2 Horizon IFM tokens.
|
||||
|
||||
K2 Horizon's template prefills the opening token, so generated text starts
|
||||
inside reasoning and this parser is always forced on. ``reasoning_effort``
|
||||
selects the matching IFM token pair.
|
||||
"""
|
||||
|
||||
_EFFORT_TOKENS = {
|
||||
"high": ("<ifm|think>", "</ifm|think>"),
|
||||
"medium": ("<ifm|think_fast>", "</ifm|think_fast>"),
|
||||
"low": ("<ifm|think_faster>", "</ifm|think_faster>"),
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream_reasoning: bool = True,
|
||||
force_reasoning: bool = True,
|
||||
continue_final_message: bool = False,
|
||||
previous_content: str = "",
|
||||
force_nonempty_content: bool = False,
|
||||
reasoning_effort: object = "high",
|
||||
):
|
||||
if not force_reasoning:
|
||||
raise ValueError("K2-v3 reasoning parser requires force_reasoning=True")
|
||||
|
||||
# Five release templates reject unsupported levels. The 0.9B template's
|
||||
# fallback emits the medium token, so medium is the only possible wire
|
||||
# format for an unsupported value that reaches generation.
|
||||
effort = (
|
||||
reasoning_effort
|
||||
if isinstance(reasoning_effort, str)
|
||||
and reasoning_effort in self._EFFORT_TOKENS
|
||||
else "medium"
|
||||
)
|
||||
start_token, end_token = self._EFFORT_TOKENS[effort]
|
||||
super().__init__(
|
||||
start_token,
|
||||
end_token,
|
||||
force_reasoning=True,
|
||||
stream_reasoning=stream_reasoning,
|
||||
# Common prefix of the singular and plural tool-call open tags.
|
||||
# This also closes reasoning for a malformed turn that omits its
|
||||
# explicit </ifm|think...> token.
|
||||
tool_start_token="<ifm|tool_call",
|
||||
continue_final_message=continue_final_message,
|
||||
previous_content=previous_content,
|
||||
reasoning_default="always",
|
||||
force_nonempty_content=force_nonempty_content,
|
||||
)
|
||||
# Catalog only: scheduler-side request validation encodes these, but
|
||||
# the active matcher must use just the delimiter selected above.
|
||||
self.request_selectable_think_end_tokens = tuple(
|
||||
tokens[1] for tokens in self._EFFORT_TOKENS.values()
|
||||
)
|
||||
|
||||
|
||||
class KimiK3Detector(BaseReasoningFormatDetector):
|
||||
"""Detector for the Kimi K3 XTML think channel.
|
||||
|
||||
@@ -1938,6 +1995,7 @@ class ReasoningParser:
|
||||
"ling3": Ling3Detector,
|
||||
"hunyuan": HunyuanDetector,
|
||||
"gpt-oss": GptOssDetector,
|
||||
"k2_horizon": K2V3Detector,
|
||||
"kimi": KimiDetector,
|
||||
"kimi_k2": KimiK2Detector,
|
||||
"kimi_k3": KimiK3Detector,
|
||||
@@ -2007,6 +2065,22 @@ class ReasoningParser:
|
||||
if chat_template_kwargs.get("force_nonempty_content") is True:
|
||||
kwargs["force_nonempty_content"] = True
|
||||
|
||||
if model_type.lower() == "k2_horizon":
|
||||
# Template kwargs are the final values passed to Jinja and therefore
|
||||
# take precedence over the convenience fields on API requests.
|
||||
effort = chat_template_kwargs.get("reasoning_effort")
|
||||
if effort is None:
|
||||
effort = getattr(request, "reasoning_effort", None)
|
||||
if effort is None:
|
||||
# The Responses API carries the same value in its standard
|
||||
# nested shape (``reasoning.effort``). Prompt rendering already
|
||||
# mirrors it into a ChatCompletionRequest; parsing must select
|
||||
# the matching IFM delimiter as well.
|
||||
reasoning = getattr(request, "reasoning", None)
|
||||
effort = getattr(reasoning, "effort", None)
|
||||
if effort is not None:
|
||||
kwargs["reasoning_effort"] = effort
|
||||
|
||||
if tokenizer is not None:
|
||||
sig = inspect.signature(detector_class)
|
||||
if "tokenizer" in sig.parameters:
|
||||
|
||||
@@ -145,6 +145,15 @@ def _has_toggle_default_assignment(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
REASONING_MODE_RULES = (
|
||||
DetectionRule(
|
||||
name="k2_v3_reasoning_effort",
|
||||
value=ReasoningToggleConfig(special_case="always"),
|
||||
predicate=lambda ctx: (
|
||||
ctx.has_text("<ifm|think>")
|
||||
and ctx.has_text("<ifm|think_fast>")
|
||||
and ctx.has_text("reasoning_effort")
|
||||
),
|
||||
),
|
||||
DetectionRule(
|
||||
name="gpt_oss_channel_markers",
|
||||
value=ReasoningToggleConfig(special_case="always"),
|
||||
@@ -288,6 +297,14 @@ def _is_kimi_k2(ctx):
|
||||
return ctx.has_vocab("<|tool_calls_section_begin|>")
|
||||
|
||||
|
||||
def _is_k2_v3(ctx):
|
||||
return (
|
||||
ctx.has_text("<ifm|think>")
|
||||
and ctx.has_text("<ifm|tool_calls>")
|
||||
and ctx.has_text("<ifm|tool_call>")
|
||||
)
|
||||
|
||||
|
||||
def _is_nemotron_3(ctx):
|
||||
return ctx.has_text("truncate_history_thinking") and (
|
||||
ctx.reasoning_config is not None
|
||||
@@ -443,6 +460,7 @@ def _is_deepseek_r1_think_tags(ctx):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
REASONING_PARSER_RULES = (
|
||||
DetectionRule(name="k2_horizon", value="k2_horizon", predicate=_is_k2_v3),
|
||||
DetectionRule(name="apertus2509", value="apertus2509", predicate=_is_apertus2509),
|
||||
DetectionRule(name="gemma4", value="gemma4", predicate=_is_gemma4),
|
||||
DetectionRule(name="kimi", value="kimi", predicate=_is_kimi),
|
||||
@@ -477,6 +495,7 @@ REASONING_PARSER_RULES = (
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TOOL_CALL_PARSER_RULES = (
|
||||
DetectionRule(name="k2_horizon", value="k2_horizon", predicate=_is_k2_v3),
|
||||
DetectionRule(name="apertus2509", value="apertus2509", predicate=_is_apertus2509),
|
||||
DetectionRule(name="gemma4", value="gemma4", predicate=_is_gemma4),
|
||||
DetectionRule(name="gpt_oss", value="gpt-oss", predicate=_is_gpt_oss),
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
|
||||
import logging
|
||||
import math
|
||||
from typing import Dict, List, Optional, Set, Union
|
||||
from typing import Dict, List, Optional, Sequence, Set, Union
|
||||
|
||||
import msgspec
|
||||
|
||||
@@ -45,6 +45,72 @@ MAX_STOP_REGEX_COUNT = 32
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Private transport from the OpenAI request renderer to scheduler-side
|
||||
# reasoning grammar/accounting. It lives in custom_params so the selected
|
||||
# delimiter follows every existing tokenizer/session/disaggregation path
|
||||
# without changing the public SamplingParams wire layout.
|
||||
REQUEST_REASONING_END_TOKEN_IDS_KEY = "__sglang_reasoning_end_token_ids"
|
||||
MAX_REQUEST_REASONING_END_TOKEN_IDS = 32
|
||||
|
||||
|
||||
def get_request_reasoning_end_token_ids(
|
||||
custom_params: Optional[Dict[str, CustomParamValue]],
|
||||
*,
|
||||
allowed_sequences: Optional[Sequence[Sequence[int]]] = None,
|
||||
vocab_size: Optional[int] = None,
|
||||
strict: bool = False,
|
||||
) -> Optional[List[int]]:
|
||||
"""Return a validated request-selected reasoning terminator, if present."""
|
||||
if not isinstance(custom_params, dict):
|
||||
return None
|
||||
if REQUEST_REASONING_END_TOKEN_IDS_KEY not in custom_params:
|
||||
return None
|
||||
token_ids = custom_params.get(REQUEST_REASONING_END_TOKEN_IDS_KEY)
|
||||
invalid = (
|
||||
not isinstance(token_ids, list)
|
||||
or not token_ids
|
||||
or len(token_ids) > MAX_REQUEST_REASONING_END_TOKEN_IDS
|
||||
or any(type(token_id) is not int or token_id < 0 for token_id in token_ids)
|
||||
or (
|
||||
vocab_size is not None
|
||||
and isinstance(token_ids, list)
|
||||
and any(
|
||||
type(token_id) is int and token_id >= vocab_size
|
||||
for token_id in token_ids
|
||||
)
|
||||
)
|
||||
)
|
||||
if invalid:
|
||||
if strict:
|
||||
raise ValueError(
|
||||
"request reasoning end token IDs must be a non-empty, "
|
||||
f"vocabulary-bounded list of at most "
|
||||
f"{MAX_REQUEST_REASONING_END_TOKEN_IDS} integers"
|
||||
)
|
||||
return None
|
||||
if allowed_sequences is None or tuple(token_ids) not in {
|
||||
tuple(sequence) for sequence in allowed_sequences
|
||||
}:
|
||||
return None
|
||||
return list(token_ids)
|
||||
|
||||
|
||||
def set_request_reasoning_end_token_ids(
|
||||
sampling_params: Dict,
|
||||
token_ids: Optional[List[int]],
|
||||
) -> None:
|
||||
"""Attach the renderer-selected reasoning terminator to a request."""
|
||||
if token_ids is None:
|
||||
return
|
||||
if not token_ids or any(
|
||||
type(token_id) is not int or token_id < 0 for token_id in token_ids
|
||||
):
|
||||
raise ValueError("reasoning end token IDs must be non-empty integers")
|
||||
custom_params = dict(sampling_params.get("custom_params") or {})
|
||||
custom_params[REQUEST_REASONING_END_TOKEN_IDS_KEY] = list(token_ids)
|
||||
sampling_params["custom_params"] = custom_params
|
||||
|
||||
|
||||
class SamplingParams(msgspec.Struct, kw_only=True, array_like=True):
|
||||
"""
|
||||
The sampling parameters.
|
||||
@@ -205,6 +271,12 @@ class SamplingParams(msgspec.Struct, kw_only=True, array_like=True):
|
||||
f"{token_id}."
|
||||
)
|
||||
|
||||
get_request_reasoning_end_token_ids(
|
||||
self.custom_params,
|
||||
vocab_size=vocab_size,
|
||||
strict=True,
|
||||
)
|
||||
|
||||
grammars = [
|
||||
self.json_schema,
|
||||
self.regex,
|
||||
|
||||
@@ -42,6 +42,7 @@ from sglang.srt.configs import (
|
||||
InternS2PreviewConfig,
|
||||
JetNemotronConfig,
|
||||
JetVLMConfig,
|
||||
K2HorizonConfig,
|
||||
KimiK3Config,
|
||||
KimiK25Config,
|
||||
KimiLinearConfig,
|
||||
@@ -70,6 +71,7 @@ from sglang.srt.configs import (
|
||||
Step3p5Config,
|
||||
Step3p7Config,
|
||||
Step3VLConfig,
|
||||
XllmConfig,
|
||||
)
|
||||
from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
|
||||
from sglang.srt.configs.internvl import InternVLChatConfig
|
||||
@@ -100,6 +102,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
|
||||
DeepseekVL2Config,
|
||||
MultiModalityConfig,
|
||||
KimiVLConfig,
|
||||
K2HorizonConfig,
|
||||
LocateAnythingConfig,
|
||||
InternVLChatConfig,
|
||||
LagunaConfig,
|
||||
@@ -142,6 +145,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
|
||||
InklingVisionConfig,
|
||||
InklingMMConfig,
|
||||
MiniMaxM3VLConfig,
|
||||
XllmConfig,
|
||||
]
|
||||
}
|
||||
|
||||
@@ -654,9 +658,15 @@ def get_tokenizer_from_processor(processor):
|
||||
|
||||
|
||||
# Turn-final markers that some checkpoints ship without EOS metadata:
|
||||
# <|eom_id|> (Llama-3 tool use) and <|content_model_end_sampling|> (Inkling,
|
||||
# whose bundled tokenizer config leaves eos_token unset).
|
||||
_ADDITIONAL_STOP_TOKEN_TEXTS = ("<|eom_id|>", "<|content_model_end_sampling|>")
|
||||
# <|eom_id|> (Llama-3 tool use), <|content_model_end_sampling|> (Inkling,
|
||||
# whose bundled tokenizer config leaves eos_token unset), and
|
||||
# <|ifm|im_end|> (some K2 Horizon checkpoints, notably 0.9B, name only
|
||||
# <|endoftext|> as EOS).
|
||||
_ADDITIONAL_STOP_TOKEN_TEXTS = (
|
||||
"<|eom_id|>",
|
||||
"<|content_model_end_sampling|>",
|
||||
"<|ifm|im_end|>",
|
||||
)
|
||||
|
||||
|
||||
def attach_additional_stop_token_ids(tokenizer):
|
||||
|
||||
Reference in New Issue
Block a user