[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:
Yash Akhauri
2026-09-03 16:39:43 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang Xinyuan Tong
parent 02d9b3060a
commit 3bac084d4e
34 changed files with 4487 additions and 16 deletions
+3
View File
@@ -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",
+30
View File
@@ -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"
)
+310
View File
@@ -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,
)
+24
View File
@@ -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),
+73 -1
View File
@@ -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):