diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 492760909..bc6211b5e 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -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", diff --git a/python/sglang/srt/configs/k2_horizon.py b/python/sglang/srt/configs/k2_horizon.py new file mode 100644 index 000000000..1096b11bc --- /dev/null +++ b/python/sglang/srt/configs/k2_horizon.py @@ -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" diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 2a5a4f459..2a87a7e0b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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( diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index c218956a4..bdd7b7c45 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -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) diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py index f9b21f56a..bb9403b78 100644 --- a/python/sglang/srt/constrained/reasoner_grammar_backend.py +++ b/python/sglang/srt/constrained/reasoner_grammar_backend.py @@ -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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 9a63da8ac..c03483ec6 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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 diff --git a/python/sglang/srt/entrypoints/openai/protocol.py b/python/sglang/srt/entrypoints/openai/protocol.py index 82a384a09..346c7c02e 100644 --- a/python/sglang/srt/entrypoints/openai/protocol.py +++ b/python/sglang/srt/entrypoints/openai/protocol.py @@ -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): diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 790f84a28..8aaa07e34 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -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( diff --git a/python/sglang/srt/entrypoints/openai/serving_responses.py b/python/sglang/srt/entrypoints/openai/serving_responses.py index bec853e81..a82321487 100644 --- a/python/sglang/srt/entrypoints/openai/serving_responses.py +++ b/python/sglang/srt/entrypoints/openai/serving_responses.py @@ -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] diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 090898b41..12dce448b 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -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, diff --git a/python/sglang/srt/function_call/k2_v3_detector.py b/python/sglang/srt/function_call/k2_v3_detector.py new file mode 100644 index 000000000..222e2002a --- /dev/null +++ b/python/sglang/srt/function_call/k2_v3_detector.py @@ -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 ```` 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 = "" + _TOOL_CALLS_END = "" + _TOOL_CALL_START = "" + _TOOL_CALL_END = "" + + _ARG_KEY_START = "" + + _THINK_PAIRS = ( + ("", ""), + ("", ""), + ("", ""), + ) + _GROUP_TOKENS = (_TOOL_CALLS_START, _TOOL_CALLS_END) + _STREAM_MARKERS = ( + _TOOL_CALL_START, + _TOOL_CALLS_START, + _TOOL_CALLS_END, + "", + "", + "", + "", + "", + "", + ) + + _ARG_PATTERN = re.compile( + r"(.*?)\s*" + r"(?:(.*?)\s*)?" + r"(.*?)", + 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 ) 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" + ) diff --git a/python/sglang/srt/layers/mova.py b/python/sglang/srt/layers/mova.py new file mode 100644 index 000000000..dc5ddb5c5 --- /dev/null +++ b/python/sglang/srt/layers/mova.py @@ -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, + ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b002d9d87..cf0f12b17 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 7e400e5f4..5aeec8a2f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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, diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py new file mode 100644 index 000000000..4fa46971a --- /dev/null +++ b/python/sglang/srt/models/xllm.py @@ -0,0 +1,1888 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2023-2024 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. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/qwen2_moe.py +# for the xLLM K2MoE architecture. +# Key differences from Qwen2Moe: +# - Sigmoid routing (not softmax) +# - Gate bias used for expert selection only (correction_bias pattern) +# - Router scaling factor applied after renormalization +# - No shared_expert_gate (shared expert output added directly) +# - Dense layers specified via mlp_only_layers config +# - Partial RoPE (rope_head_dim < head_dim) +"""Inference-only xLLM K2MoE and MoVA models compatible with HF weights.""" + +import math +from contextlib import nullcontext +from typing import Any, Dict, Iterable, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group, tensor_model_parallel_all_reduce +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_skip_post_experts_all_reduce, +) +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat +from sglang.srt.layers.moe.utils import ( + RoutingMethodType, + filter_moe_weight_param_global_expert, +) +from sglang.srt.layers.mova import RoutedValueExperts, mova_router_topk +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.cuda_graph_config import ( + Backend, + Phase, + check_cuda_graph_backend, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.runtime_context import get_exec, get_parallel +from sglang.srt.utils import add_prefix, make_layers + +_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY = "xllm_source_router_gemm_partitions" +_XLLM_SOURCE_ROUTER_PARTITIONS_MISSING = object() +_XLLM_CHECKPOINT_FORMAT_CONFIG_KEY = "_sglang_xllm_checkpoint_format" +_K2_HORIZON_HF_CHECKPOINT_FORMAT = "k2_horizon_hf" +_CONFIG_ATTR_MISSING = object() + + +class XllmGroupRMSNorm(nn.Module): + """Reference grouped RMSNorm used by the xLLM model family.""" + + def __init__( + self, + hidden_size: int, + n_groups: int = 1, + eps: float = 1e-6, + zero_centered: bool = False, + ): + super().__init__() + self.n_groups = n_groups + self.hidden_size = hidden_size + if n_groups <= 0 or hidden_size % n_groups: + raise ValueError( + f"hidden_size={hidden_size} must be divisible by n_groups={n_groups}" + ) + self.variance_epsilon = eps + self.zero_centered = zero_centered + self.weight = nn.Parameter( + torch.zeros(hidden_size) if zero_centered else torch.ones(hidden_size) + ) + + def forward(self, hidden_states, residual=None, post_residual_addition=None): + if residual is not None: + hidden_states = hidden_states + residual + residual = hidden_states + if post_residual_addition is not None: + hidden_states = hidden_states + post_residual_addition + residual = hidden_states + orig_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + hidden_states = hidden_states.reshape( + *hidden_states.shape[:-1], self.n_groups, -1 + ) + hidden_states = hidden_states * torch.rsqrt( + hidden_states.pow(2).mean(-1, keepdim=True) + self.variance_epsilon + ) + hidden_states = hidden_states.reshape(*hidden_states.shape[:-2], -1) + weight = self.weight + 1.0 if self.zero_centered else self.weight + hidden_states = (weight * hidden_states).to(orig_dtype) + if residual is not None: + return hidden_states, residual + return hidden_states + + +def _is_k2_horizon_hf_checkpoint(config: PretrainedConfig) -> bool: + return ( + getattr(config, _XLLM_CHECKPOINT_FORMAT_CONFIG_KEY, None) + == _K2_HORIZON_HF_CHECKPOINT_FORMAT + ) + + +def _set_k2_horizon_alias( + config: PretrainedConfig, + *, + source_name: str, + target_name: str, + value: Any, +) -> None: + """Set one explicit K2Horizon schema alias, rejecting contradictions.""" + + current = getattr(config, target_name, _CONFIG_ATTR_MISSING) + if current is not _CONFIG_ATTR_MISSING and current is not None: + if current != value: + raise ValueError( + f"K2Horizon config has conflicting {source_name}={value!r} " + f"and {target_name}={current!r}" + ) + return + setattr(config, target_name, value) + + +def _normalize_k2_horizon_config(config: PretrainedConfig) -> None: + """Translate the canonical K2Horizon HF schema to the native xLLM path. + + This adapter intentionally maps only fields that K2Horizon spells + differently. In particular, source router GEMM topology is provenance, + not an architecture property, so it must be supplied explicitly. + """ + + # Dense K2Horizon artifacts may omit the MoVA fields entirely, while the + # remote config class supplies zero defaults. Treat both representations + # identically, but keep malformed/negative values distinct from dense. + mova_num_experts = getattr(config, "mova_num_experts", 0) + if isinstance(mova_num_experts, bool) or not isinstance(mova_num_experts, int): + raise ValueError( + "K2Horizon mova_num_experts must be a non-negative integer, " + f"got {mova_num_experts!r}." + ) + if mova_num_experts < 0: + raise ValueError( + "K2Horizon mova_num_experts must be a non-negative integer, " + f"got {mova_num_experts!r}." + ) + is_mova = mova_num_experts > 0 + + if is_mova: + if _get_xllm_source_router_gemm_partitions(config) is None: + raise ValueError( + "K2Horizon MoVA requires explicit source router GEMM provenance; " + "SGLang will not infer it from runtime tensor parallelism. After " + "confirming the training contract, pass " + "--json-model-override-args " + "'{\"xllm_source_router_gemm_partitions\": 2}' (use 1 only for " + "a confirmed MP1 source checkpoint)." + ) + + _set_k2_horizon_alias( + config, + source_name="mova_num_experts", + target_name="num_values", + value=mova_num_experts, + ) + mova_num_experts_per_tok = getattr( + config, "mova_num_experts_per_tok", _CONFIG_ATTR_MISSING + ) + if ( + isinstance(mova_num_experts_per_tok, bool) + or not isinstance(mova_num_experts_per_tok, int) + or not 0 < mova_num_experts_per_tok <= mova_num_experts + ): + raise ValueError( + "K2Horizon mova_num_experts_per_tok must be a positive integer no " + f"larger than mova_num_experts, got {mova_num_experts_per_tok!r}" + ) + _set_k2_horizon_alias( + config, + source_name="mova_num_experts_per_tok", + target_name="num_values_per_tok", + value=mova_num_experts_per_tok, + ) + else: + mova_num_experts_per_tok = getattr(config, "mova_num_experts_per_tok", 0) + if ( + isinstance(mova_num_experts_per_tok, bool) + or not isinstance(mova_num_experts_per_tok, int) + or mova_num_experts_per_tok != 0 + ): + raise ValueError( + "Dense K2Horizon requires mova_num_experts_per_tok=0, got " + f"{mova_num_experts_per_tok!r}" + ) + _set_k2_horizon_alias( + config, + source_name="mova_num_experts", + target_name="num_values", + value=0, + ) + _set_k2_horizon_alias( + config, + source_name="mova_num_experts_per_tok", + target_name="num_values_per_tok", + value=0, + ) + for field in ("num_experts", "num_experts_per_tok", "num_shared_experts"): + value = getattr(config, field, 0) + if isinstance(value, bool) or not isinstance(value, int) or value != 0: + raise ValueError(f"Dense K2Horizon requires {field}=0, got {value!r}") + # Some dense exports omit the MoE-only fields. Downstream model + # construction reads them directly, so materialize the validated + # dense defaults instead of relying on getattr fallbacks forever. + setattr(config, field, 0) + if getattr(config, "query_key_norm", False): + raise ValueError( + "Dense K2Horizon native loading does not support query/key " + "normalization" + ) + if getattr(config, "sliding_window", None) is not None or getattr( + config, "use_sliding_window", False + ): + raise ValueError( + "Dense K2Horizon native loading supports full causal attention only" + ) + attention_gate_func = getattr( + config, "attention_gate_func", _CONFIG_ATTR_MISSING + ) + native_gate_func = getattr(config, "attn_gate_func", _CONFIG_ATTR_MISSING) + if ( + attention_gate_func not in (_CONFIG_ATTR_MISSING, None) + or native_gate_func not in (_CONFIG_ATTR_MISSING, None) + or getattr(config, "apply_attn_gate", False) + ): + raise ValueError( + "Dense K2Horizon native loading does not support gated attention" + ) + + attention_gate_func = getattr(config, "attention_gate_func", _CONFIG_ATTR_MISSING) + if attention_gate_func is not _CONFIG_ATTR_MISSING: + _set_k2_horizon_alias( + config, + source_name="attention_gate_func", + target_name="attn_gate_func", + value=attention_gate_func, + ) + _set_k2_horizon_alias( + config, + source_name="attention_gate_func", + target_name="apply_attn_gate", + value=attention_gate_func is not None, + ) + + rope_parameters = getattr(config, "rope_parameters", _CONFIG_ATTR_MISSING) + if rope_parameters is not _CONFIG_ATTR_MISSING and rope_parameters is not None: + if not isinstance(rope_parameters, dict): + raise ValueError( + "K2Horizon rope_parameters must be a dictionary, got " + f"{type(rope_parameters).__name__}" + ) + rope_type = rope_parameters.get( + "rope_type", rope_parameters.get("type", _CONFIG_ATTR_MISSING) + ) + if is_mova and rope_type != "default": + raise ValueError( + "K2Horizon direct loading supports only explicit default " + f"rope_parameters, got rope_type={rope_type!r}" + ) + if not is_mova and rope_type not in ("default", "yarn"): + raise ValueError( + "Dense K2Horizon direct loading supports only explicit default " + f"or yarn rope_parameters, got rope_type={rope_type!r}" + ) + if ( + "rope_type" in rope_parameters + and "type" in rope_parameters + and rope_parameters["rope_type"] != rope_parameters["type"] + ): + raise ValueError( + "K2Horizon rope_parameters has conflicting rope_type and type" + ) + rope_theta = rope_parameters.get("rope_theta", _CONFIG_ATTR_MISSING) + if rope_theta is _CONFIG_ATTR_MISSING: + if is_mova: + raise ValueError( + "K2Horizon default rope_parameters must explicitly provide " + "rope_theta" + ) + # Dense K2Horizon YaRN artifacts generated during the TF5 config + # transition persisted theta at the legacy top level only. + rope_theta = getattr(config, "rope_theta", _CONFIG_ATTR_MISSING) + if ( + rope_theta is _CONFIG_ATTR_MISSING + or isinstance(rope_theta, bool) + or not isinstance(rope_theta, (int, float)) + or not math.isfinite(rope_theta) + or rope_theta <= 0 + ): + raise ValueError( + "K2Horizon rope_theta must be a positive finite number, got " + f"{rope_theta!r}" + ) + _set_k2_horizon_alias( + config, + source_name="rope_parameters.rope_theta", + target_name="rope_theta", + value=rope_theta, + ) + default_rope_scaling = dict(rope_parameters) + default_rope_scaling.pop("type", None) + default_rope_scaling["rope_theta"] = rope_theta + default_rope_scaling["rope_type"] = rope_type + if rope_type == "yarn": + supported_yarn_keys = { + "attention_factor", + "beta_fast", + "beta_slow", + "factor", + "original_max_position_embeddings", + "rope_theta", + "rope_type", + "truncate", + "type", + } + unknown_yarn_keys = set(rope_parameters) - supported_yarn_keys + if unknown_yarn_keys: + raise ValueError( + "Dense K2Horizon YaRN has unsupported rope_parameters keys: " + f"{sorted(unknown_yarn_keys)}" + ) + factor = rope_parameters.get("factor", _CONFIG_ATTR_MISSING) + if ( + isinstance(factor, bool) + or not isinstance(factor, (int, float)) + or not math.isfinite(factor) + or factor <= 0 + ): + raise ValueError( + "Dense K2Horizon YaRN factor must be positive and finite, " + f"got {factor!r}" + ) + original_max_position_embeddings = rope_parameters.get( + "original_max_position_embeddings", _CONFIG_ATTR_MISSING + ) + if ( + isinstance(original_max_position_embeddings, bool) + or not isinstance(original_max_position_embeddings, int) + or original_max_position_embeddings <= 0 + ): + raise ValueError( + "Dense K2Horizon YaRN original_max_position_embeddings must " + f"be a positive integer, got {original_max_position_embeddings!r}" + ) + _set_k2_horizon_alias( + config, + source_name="rope_parameters.original_max_position_embeddings", + target_name="original_max_position_embeddings", + value=original_max_position_embeddings, + ) + max_position_embeddings = getattr( + config, "max_position_embeddings", _CONFIG_ATTR_MISSING + ) + expected_max_position_embeddings = factor * original_max_position_embeddings + if ( + isinstance(max_position_embeddings, bool) + or not isinstance(max_position_embeddings, int) + or not math.isclose( + max_position_embeddings, + expected_max_position_embeddings, + rel_tol=0.0, + abs_tol=1e-9, + ) + ): + raise ValueError( + "Dense K2Horizon YaRN requires max_position_embeddings == " + "factor * original_max_position_embeddings; got " + f"{max_position_embeddings!r} != " + f"{expected_max_position_embeddings!r}" + ) + for field, default in (("beta_fast", 32), ("beta_slow", 1)): + value = rope_parameters.get(field, default) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError( + f"Dense K2Horizon YaRN {field} must be positive and " + f"finite, got {value!r}" + ) + truncate = rope_parameters.get("truncate", True) + if not isinstance(truncate, bool): + raise ValueError( + f"Dense K2Horizon YaRN truncate must be a bool, got {truncate!r}" + ) + attention_factor = default_rope_scaling.pop("attention_factor", None) + if attention_factor is not None: + if ( + isinstance(attention_factor, bool) + or not isinstance(attention_factor, (int, float)) + or not math.isfinite(attention_factor) + or attention_factor <= 0 + ): + raise ValueError( + "Dense K2Horizon YaRN attention_factor must be positive " + f"and finite, got {attention_factor!r}" + ) + # HF's attention_factor is the final multiplier applied to + # cos/sin. SGLang's attn_factor multiplies its own standard + # YaRN mscale, so translate between those two conventions. + default_attention_factor = ( + 1.0 if factor <= 1 else 0.1 * math.log(factor) + 1.0 + ) + default_rope_scaling["attn_factor"] = ( + attention_factor / default_attention_factor + ) + current_rope_scaling = getattr(config, "rope_scaling", _CONFIG_ATTR_MISSING) + # Transformers 5 exposes rope_scaling as a property alias for the + # original rope_parameters dictionary. That is the same source field, + # not a second independently specified value. + if current_rope_scaling == rope_parameters: + setattr(config, "rope_scaling", default_rope_scaling) + else: + _set_k2_horizon_alias( + config, + source_name="rope_parameters", + target_name="rope_scaling", + value=default_rope_scaling, + ) + elif not is_mova: + raise ValueError( + "Dense K2Horizon native loading requires explicit rope_parameters" + ) + + mlp_only_layers = getattr(config, "mlp_only_layers", _CONFIG_ATTR_MISSING) + if mlp_only_layers is not _CONFIG_ATTR_MISSING: + if not isinstance(mlp_only_layers, (list, tuple)) or any( + isinstance(layer_id, bool) or not isinstance(layer_id, int) + for layer_id in mlp_only_layers + ): + raise ValueError("K2Horizon mlp_only_layers must be a list of integers") + expected_prefix = list(range(len(mlp_only_layers))) + if list(mlp_only_layers) != expected_prefix: + raise ValueError( + "K2Horizon MoVA requires mlp_only_layers to be a contiguous " + f"prefix starting at zero, got {list(mlp_only_layers)}" + ) + if not is_mova and list(mlp_only_layers) != list( + range(config.num_hidden_layers) + ): + raise ValueError( + "Dense K2Horizon native loading requires every layer in mlp_only_layers" + ) + _set_k2_horizon_alias( + config, + source_name="mlp_only_layers", + target_name="num_dense_layers", + value=len(mlp_only_layers), + ) + elif not is_mova: + raise ValueError( + "Dense K2Horizon native loading requires explicit mlp_only_layers" + ) + + current_format = getattr( + config, _XLLM_CHECKPOINT_FORMAT_CONFIG_KEY, _CONFIG_ATTR_MISSING + ) + if current_format not in ( + _CONFIG_ATTR_MISSING, + _K2_HORIZON_HF_CHECKPOINT_FORMAT, + ): + raise ValueError( + "K2Horizon native adapter requires checkpoint format " + f"{_K2_HORIZON_HF_CHECKPOINT_FORMAT!r}, got {current_format!r}" + ) + setattr( + config, + _XLLM_CHECKPOINT_FORMAT_CONFIG_KEY, + _K2_HORIZON_HF_CHECKPOINT_FORMAT, + ) + + +def _make_norm(config): + """Create the appropriate RMSNorm for this config.""" + n_groups = getattr(config, "layernorm_num_groups", 1) + is_mova = getattr(config, "num_values", 0) > 0 + if (n_groups is None or n_groups <= 1) and not is_mova: + return RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + return XllmGroupRMSNorm( + config.hidden_size, + n_groups=n_groups or 1, + eps=config.rms_norm_eps, + # Converted xLLM MoVA stores zero-centered norm deltas. Canonical + # K2Horizon HF stores ordinary one-centered RMSNorm weights directly. + zero_centered=is_mova and not _is_k2_horizon_hf_checkpoint(config), + ) + + +def _get_xllm_source_router_gemm_partitions( + config: PretrainedConfig, +) -> Optional[int]: + """Read optional source-router provenance without inferring it from TP.""" + + partitions = getattr( + config, + _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY, + _XLLM_SOURCE_ROUTER_PARTITIONS_MISSING, + ) + if partitions is _XLLM_SOURCE_ROUTER_PARTITIONS_MISSING: + return None + if ( + isinstance(partitions, bool) + or not isinstance(partitions, int) + or partitions not in (1, 2) + ): + raise ValueError( + f"{_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY}={partitions!r} is " + f"invalid (type={type(partitions).__name__}); when present it must " + "be the integer 1 or 2. Omit the key to preserve legacy router " + "GEMM behavior." + ) + if config.hidden_size % partitions: + raise ValueError( + f"explicit {_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY}={partitions} " + f"requires hidden_size divisible by {partitions}; got " + f"hidden_size={config.hidden_size}" + ) + return partitions + + +def _xllm_router_gemm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + source_partitions: Optional[int], +) -> torch.Tensor: + """Reproduce the source xLLM router GEMM's partition rounding contract.""" + + # Old xLLM artifacts have no source-topology provenance. Preserve their + # exact pre-contract behavior instead of guessing how the router was run. + if source_partitions is None: + return F.linear(hidden_states, weight) + if isinstance(source_partitions, bool) or not isinstance(source_partitions, int): + raise ValueError( + "explicit xLLM router source partitions must be the integer 1 or " + f"2; got {source_partitions!r} " + f"(type={type(source_partitions).__name__})" + ) + if source_partitions not in (1, 2): + raise ValueError( + "explicit xLLM router source partitions must be 1 or 2, got " + f"{source_partitions}" + ) + if hidden_states.ndim < 1 or weight.ndim != 2: + raise ValueError( + "xLLM router GEMM expects input [..., hidden] and weight " + f"[routes, hidden]; got input={tuple(hidden_states.shape)}, " + f"weight={tuple(weight.shape)}" + ) + if hidden_states.shape[-1] != weight.shape[-1]: + raise ValueError( + "xLLM router input and weight hidden dimensions differ; got " + f"input={tuple(hidden_states.shape)}, weight={tuple(weight.shape)}" + ) + if hidden_states.shape[-1] % source_partitions: + raise ValueError( + f"explicit xLLM router partitions={source_partitions} requires " + f"hidden size divisible by {source_partitions}; got hidden_size=" + f"{hidden_states.shape[-1]}" + ) + if hidden_states.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise ValueError( + "Explicit xLLM source router GEMM provenance requires BF16 input " + f"and weight; got input={hidden_states.dtype}, weight={weight.dtype}" + ) + + if source_partitions == 1: + return F.linear(hidden_states, weight).float() + + # Native xLLM row-shards each router across MP2. Each rank performs a BF16 + # partial GEMM, rounds that result to BF16, casts it to FP32, and then the + # FP32 all-reduce adds the two partials. Emulate that ordering locally. + input_parts = hidden_states.chunk(source_partitions, dim=-1) + weight_parts = weight.chunk(source_partitions, dim=-1) + first = F.linear(input_parts[0].contiguous(), weight_parts[0].contiguous()) + second = F.linear(input_parts[1].contiguous(), weight_parts[1].contiguous()) + return first.float() + second.float() + + +def _validate_mova_config( + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], +) -> None: + """Fail early for native runtime combinations that cannot be served exactly.""" + + if getattr(config, "model_type", None) in ("xllm", "k2_horizon"): + if torch.get_default_dtype() != torch.bfloat16: + raise ValueError( + "Native xLLM/K2 Horizon serving requires --dtype bfloat16: " + "the released checkpoints persist float32 dtype metadata but " + "their weights and validated runtime contract are BF16." + ) + if quant_config is not None: + raise ValueError( + "Native xLLM/K2 Horizon serving does not support quantized " + "model weights" + ) + + runtime = get_exec() + if runtime.overlap.enable_two_batch_overlap: + raise ValueError( + "Native xLLM/K2 Horizon serving does not yet support " + "--enable-two-batch-overlap" + ) + + moe_runtime = runtime.moe + unsupported_expert_remap = ( + moe_runtime.enable_eplb + or moe_runtime.init_expert_location != "trivial" + or moe_runtime.ep_num_redundant_experts > 0 + ) + if unsupported_expert_remap: + raise ValueError( + "Native xLLM/K2 Horizon serving does not yet support EPLB, " + "non-trivial initial expert placement, or redundant experts; " + "these modes require logical-to-physical expert remapping on " + "every MoE backend." + ) + + if getattr(config, "num_values", 0) <= 0: + # Legacy K2 checkpoints use model_type="xllm" for the ordinary + # attention path. XllmAttention implements partial RoPE and biased + # QKV projections, but it does not implement the original xLLM + # query/key normalization, sliding-window attention, or attention + # gating variants. Reject those layouts here instead of silently + # loading them with different attention math. + if getattr(config, "query_key_norm", False): + raise ValueError( + "Native dense xLLM attention does not support query/key normalization" + ) + if getattr(config, "sliding_window", None) is not None or getattr( + config, "use_sliding_window", False + ): + raise ValueError( + "Native dense xLLM attention supports full causal attention only" + ) + if getattr(config, "apply_attn_gate", False): + raise ValueError( + "Native dense xLLM attention does not support gated attention" + ) + return + _get_xllm_source_router_gemm_partitions(config) + if getattr(config, "attention_bias", False): + raise ValueError("K2 Horizon MoVA requires bias-free Q/K/V/O projections") + if getattr(config, "query_key_norm", False): + raise ValueError("K2 Horizon MoVA does not support query/key normalization") + if not getattr(config, "apply_attn_gate", False): + raise ValueError("K2 Horizon MoVA requires the xLLM attention gate") + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + if head_dim % 2: + raise ValueError(f"MoVA requires an even RoPE head dimension, got {head_dim}") + if config.num_attention_heads % config.num_key_value_heads: + raise ValueError("MoVA requires query heads to be divisible by KV heads") + if getattr(config, "rope_head_dim", head_dim) != head_dim: + raise ValueError("K2 Horizon MoVA requires full-head interleaved RoPE") + rope_scaling = getattr(config, "rope_scaling", None) + # Transformers 5 normalizes a JSON ``rope_scaling: null`` into an + # explicit default-RoPE dictionary. That representation does not change + # the rotary math and must not be confused with linear/dynamic scaling. + if rope_scaling is not None and not ( + isinstance(rope_scaling, dict) + and rope_scaling.get("rope_type", rope_scaling.get("type")) == "default" + ): + raise ValueError("K2 Horizon MoVA does not support non-default RoPE scaling") + if getattr(config, "sliding_window", None) is not None or getattr( + config, "use_sliding_window", False + ): + raise ValueError("K2 Horizon MoVA uses full causal RadixAttention only") + if getattr(config, "attn_gate_func", "silu") not in ("silu", "softplus"): + raise ValueError("MoVA supports only silu and softplus attention gates") + if getattr(config, "router_score_func", "sigmoid") not in ("sigmoid", "softmax"): + raise ValueError("MoVA supports only sigmoid and softmax value routing") + router_scale = getattr(config, "router_scaling_factor", 1.0) + if router_scale is None or not math.isfinite(router_scale) or router_scale <= 0: + raise ValueError( + f"MoVA requires a positive finite router scaling factor, got {router_scale}" + ) + num_dense_layers = getattr(config, "num_dense_layers", None) + if ( + num_dense_layers is None + or not 0 <= num_dense_layers <= config.num_hidden_layers + ): + raise ValueError( + "MoVA requires num_dense_layers in [0, num_hidden_layers], got " + f"{num_dense_layers}" + ) + expected_dense_layers = list(range(num_dense_layers)) + if list(getattr(config, "mlp_only_layers", [])) != expected_dense_layers: + raise ValueError( + "K2 Horizon MoVA requires dense attention and dense FFN prefix layers to " + f"match exactly; expected mlp_only_layers={expected_dense_layers}" + ) + if getattr(config, "decoder_sparse_step", 1) != 1: + raise ValueError("K2 Horizon MoVA requires decoder_sparse_step=1") + num_values = config.num_values + top_k = getattr(config, "num_values_per_tok", 0) + if not 0 < top_k <= num_values: + raise ValueError( + f"num_values_per_tok must be in [1, {num_values}], got {top_k}" + ) + if getattr(config, "num_experts", 0) <= 0: + raise ValueError("MoVA requires sparse MoE feed-forward layers") + n_groups = getattr(config, "layernorm_num_groups", 1) or 1 + if config.hidden_size % n_groups: + raise ValueError( + f"hidden size {config.hidden_size} is not divisible by {n_groups} norm groups" + ) + attn_tp_size = get_parallel().attn_tp_size + if config.num_attention_heads % attn_tp_size: + raise ValueError(f"MoVA query heads must be divisible by TP={attn_tp_size}") + if config.num_key_value_heads % attn_tp_size: + raise ValueError( + "K2 Horizon MoVA requires TP <= KV heads and KV heads divisible by TP; " + f"got TP={attn_tp_size}, KV heads={config.num_key_value_heads}" + ) + + +def _xllm_stacked_params_mapping(config: PretrainedConfig): + if getattr(config, "num_values", 0) <= 0: + return [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + mapping = [ + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + mapping.extend( + (".v_experts.weight", f".v_experts.{expert_id}.weight", expert_id) + for expert_id in range(config.num_values) + ) + return mapping + + +def permute_to_xllm(x): + """Interleave first half and second half: [0,1,...,63,64,...,127] -> [0,64,1,65,...,63,127]""" + return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1) + + +def permute_to_hf(x): + """Inverse of permute_to_xllm: [0,64,1,65,...,63,127] -> [0,1,...,63,64,...,127]""" + return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1) + + +class XllmMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward( + self, + x, + use_reduce_scatter: bool = False, + ): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x, skip_all_reduce=use_reduce_scatter) + return x + + +class XllmMoEGate(nn.Module): + """Router gate for xllm. + + Stores weight and bias separately. The bias is used as correction_bias + for expert selection (added to sigmoid scores) but not in the linear + computation of router logits. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.source_router_gemm_partitions = _get_xllm_source_router_gemm_partitions( + config + ) + self.weight = nn.Parameter( + torch.empty((config.num_experts, config.hidden_size)) + ) + if getattr(config, "moe_gate_bias", False): + # topk_sigmoid kernel requires correction_bias in float32 + self.bias = nn.Parameter( + torch.empty(config.num_experts, dtype=torch.float32) + ) + else: + self.bias = None + + def forward(self, hidden_states: torch.Tensor): + # The reference router applies sigmoid/softmax in FP32 after the BF16 + # GEMM (or after the explicit source-partition reduction). + return _xllm_router_gemm( + hidden_states, self.weight, self.source_router_gemm_partitions + ).float() + + +class XllmSparseMoeBlock(nn.Module): + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_parallel().tp_size + self.layer_id = layer_id + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + self.router_scaling_factor = getattr(config, "router_scaling_factor", 1.0) + + self.gate = XllmMoEGate(config) + + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + layer_id=layer_id, + scoring_func=getattr(config, "router_score_func", "sigmoid"), + correction_bias=self.gate.bias, + # xLLM needs explicit ids and weights so correction-bias routing, + # EPLB remapping, and post-renormalization scaling keep identical + # semantics on every MoE runner backend. + output_format=TopKOutputFormat.STANDARD, + ) + + self.experts = get_moe_impl_class(quant_config)( + layer_id=self.layer_id, + top_k=config.num_experts_per_tok, + num_experts=config.num_experts + get_exec().moe.ep_num_redundant_experts, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + prefix=add_prefix("experts", prefix), + routing_method_type=RoutingMethodType.RenormalizeNaive, + ) + + # Shared expert (no gating — output added directly) + num_shared_experts = getattr(config, "num_shared_experts", 0) + if num_shared_experts > 0: + shared_intermediate_size = config.moe_intermediate_size * num_shared_experts + self.shared_experts = XllmMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + **( + dict(tp_rank=0, tp_size=1) + if ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mori() + or get_moe_a2a_backend().is_flashinfer() + ) + else {} + ), + ) + else: + self.shared_experts = None + + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori(): + self.ep_size = get_parallel().moe_ep_size + self.num_experts = ( + config.num_experts + get_exec().moe.ep_num_redundant_experts + ) + self.top_k = config.num_experts_per_tok + + def get_moe_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + and filter_moe_weight_param_global_expert( + name, x, self.experts.num_local_experts + ) + ] + + def _forward_shared_experts(self, hidden_states: torch.Tensor): + if self.shared_experts is not None: + return self.shared_experts(hidden_states) + return None + + def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + shared_output = None + if hidden_states.shape[0] > 0: + router_logits = self.gate(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) + # DeepEP/EPLB requires the current dispatched TopK path so logical + # expert ids can be remapped and padded rows can be masked. + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=( + ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) + ), + ) + # Apply router scaling factor after renormalization + if self.router_scaling_factor != 1.0: + scaled_weights = topk_output.topk_weights * self.router_scaling_factor + if hasattr(topk_output, "_replace"): + topk_output = topk_output._replace(topk_weights=scaled_weights) + else: + topk_output.topk_weights = scaled_weights + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, layer_id=self.layer_id + ) + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + ) + + if shared_output is not None: + final_hidden_states.add_(shared_output) + + return final_hidden_states + + def _forward_router_experts(self, hidden_states: torch.Tensor): + router_logits = self.gate(hidden_states) + topk_output = self.topk.forward_native(hidden_states, router_logits) + # Apply router scaling factor after renormalization + # TopK output is a NamedTuple (immutable), so we must replace it + if self.router_scaling_factor != 1.0: + topk_output = topk_output._replace( + topk_weights=topk_output.topk_weights * self.router_scaling_factor + ) + return self.experts(hidden_states, topk_output) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori(): + return self._forward_deepep(hidden_states, forward_batch) + + if hidden_states.shape[0] == 0: + shared_output = None + topk_output = self.topk.empty_topk_output( + hidden_states.device, layer_id=self.layer_id + ) + final_hidden_states = self.experts(hidden_states, topk_output) + else: + shared_output = self._forward_shared_experts(hidden_states) + final_hidden_states = self._forward_router_experts(hidden_states) + + if shared_output is not None: + final_hidden_states += shared_output + if ( + self.tp_size > 1 + and not use_reduce_scatter + and not should_skip_post_experts_all_reduce(is_tp_path=True) + and not get_moe_a2a_backend().is_flashinfer() + ): + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class XllmAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rope_head_dim: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + qkv_bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + + self.total_num_heads = num_heads + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= attn_tp_size: + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=qkv_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + reduce_results=False, + prefix=add_prefix("o_proj", prefix), + ) + + # Partial RoPE: xLLM/HF stores each head in neox ordering, where the + # rotary dimensions are not contiguous when rope_head_dim < head_dim. + # Mirror HF exactly: permute to interleaved, split rope/nope, apply RoPE + # on the rope slice, then recombine and permute back. + self.rope_head_dim = rope_head_dim + self.use_xllm_partial_rope = rope_head_dim < head_dim + self.rotary_emb = get_rope( + self.rope_head_dim if self.use_xllm_partial_rope else self.head_dim, + rotary_dim=rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def _apply_partial_rope( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + q_heads = q.reshape(-1, self.num_heads, self.head_dim) + k_heads = k.reshape(-1, self.num_kv_heads, self.head_dim) + + q_interleaved = permute_to_xllm(q_heads) + k_interleaved = permute_to_xllm(k_heads) + + nope_dim = self.head_dim - self.rope_head_dim + q_rope, q_nope = q_interleaved.split([self.rope_head_dim, nope_dim], dim=-1) + k_rope, k_nope = k_interleaved.split([self.rope_head_dim, nope_dim], dim=-1) + + q_rope_flat = permute_to_hf(q_rope).reshape( + -1, self.num_heads * self.rope_head_dim + ) + k_rope_flat = permute_to_hf(k_rope).reshape( + -1, self.num_kv_heads * self.rope_head_dim + ) + q_rope_flat, k_rope_flat = self.rotary_emb(positions, q_rope_flat, k_rope_flat) + + q_rope = permute_to_xllm( + q_rope_flat.reshape(-1, self.num_heads, self.rope_head_dim) + ) + k_rope = permute_to_xllm( + k_rope_flat.reshape(-1, self.num_kv_heads, self.rope_head_dim) + ) + + q = permute_to_hf(torch.cat([q_rope, q_nope], dim=-1)).reshape( + -1, self.num_heads * self.head_dim + ) + k = permute_to_hf(torch.cat([k_rope, k_nope], dim=-1)).reshape( + -1, self.num_kv_heads * self.head_dim + ) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.use_xllm_partial_rope: + q, k = self._apply_partial_rope(positions, q, k) + else: + q, k = self.rotary_emb(positions, q, k) + + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class _XllmMoVAAttentionBase(nn.Module): + """Shared gated-GQA path for dense and routed-value MoVA layers.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig], + prefix: str, + ) -> None: + super().__init__() + if quant_config is not None: + raise ValueError( + "K2 Horizon MoVA supports unquantized bf16/fp16 weights only" + ) + + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.rope_head_dim = getattr(config, "rope_head_dim", self.head_dim) + self.apply_attn_gate = getattr(config, "apply_attn_gate", False) + self.attn_gate_func = getattr(config, "attn_gate_func", "silu") + self.scaling = self.head_dim**-0.5 + + self.tp_rank = get_parallel().attn_tp_rank + self.tp_size = get_parallel().attn_tp_size + if self.total_num_heads % self.tp_size: + raise ValueError( + f"Attention heads {self.total_num_heads} are not divisible by TP={self.tp_size}" + ) + if self.total_num_kv_heads % self.tp_size: + raise ValueError( + "K2 Horizon MoVA requires TP <= KV heads and KV heads divisible by TP; " + f"got TP={self.tp_size}, KV heads={self.total_num_kv_heads}" + ) + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = self.total_num_kv_heads // self.tp_size + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + # Keep the public correctness path as ordinary checkpoint-shaped + # projections. Packing Q/K/gate is a performance optimization and is + # deliberately outside the initial K2 Horizon integration. + self.q_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("q_proj", prefix), + ) + self.k_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("k_proj", prefix), + ) + self.gate_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("gate_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + reduce_results=False, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rope_head_dim, + max_position=getattr(config, "max_position_embeddings", 8192), + base=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=None, + prefix=add_prefix("attn", prefix), + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def _activate_gate(self, gate: torch.Tensor) -> torch.Tensor: + if self.attn_gate_func == "silu": + return F.silu(gate) + if self.attn_gate_func == "softplus": + return F.softplus(gate, beta=math.log(2)) + raise ValueError( + f"Unsupported xLLM attention gate function: {self.attn_gate_func}" + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + q, _ = self.q_proj(hidden_states) + k, _ = self.k_proj(hidden_states) + gate, _ = self.gate_proj(hidden_states) + value = self._project_value(hidden_states) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, value, forward_batch) + if self.apply_attn_gate: + attn_output = attn_output * self._activate_gate(gate) + output, _ = self.o_proj(attn_output) + return output + + +class XllmGatedAttention(_XllmMoVAAttentionBase): + """Dense GQA used by the prefix layers of a MoVA checkpoint.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + self.v_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("v_proj", prefix), + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + value, _ = self.v_proj(hidden_states) + return value + + +class XllmMoVAAttention(_XllmMoVAAttentionBase): + """Sparse MoVA attention with output-sharded routed value experts.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + self.num_values = config.num_values + self.num_values_per_tok = config.num_values_per_tok + self.router_score_func = getattr(config, "router_score_func", "sigmoid") + self.router_scaling_factor = getattr(config, "router_scaling_factor", 1.0) + self.renormalize = getattr(config, "norm_topk_prob", True) + self.source_router_gemm_partitions = _get_xllm_source_router_gemm_partitions( + config + ) + self.v_router = ReplicatedLinear( + config.hidden_size, + self.num_values, + bias=False, + quant_config=None, + prefix=add_prefix("v_router", prefix), + ) + if getattr(config, "moe_gate_bias", False): + # SGLang's fused sigmoid top-k requires correction bias in fp32. + # It remains a loadable parameter for Miles weight updates, but is + # never included in the router logits matmul. + self.v_router.bias = nn.Parameter( + torch.empty(self.num_values, dtype=torch.float32), + requires_grad=False, + ) + self.v_experts = RoutedValueExperts( + self.num_values, + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Router bias is deliberately omitted from the logits matmul. It only + # changes route selection inside ``mova_router_topk``. + router_logits = _xllm_router_gemm( + hidden_states, + self.v_router.weight, + self.source_router_gemm_partitions, + ).float() + routing_weights, selected_values = mova_router_topk( + router_logits, + self.v_router.bias, + score_func=self.router_score_func, + top_k=self.num_values_per_tok, + scaling_factor=self.router_scaling_factor, + renormalize=self.renormalize, + ) + return self.v_experts(hidden_states, routing_weights, selected_values) + + +class XllmDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + qkv_bias = getattr(config, "attention_bias", False) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + rope_head_dim = getattr(config, "rope_head_dim", head_dim) + + self.layer_id = layer_id + + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank + + # Determine if this layer is sparse (MoE) or dense + mlp_only_layers = getattr(config, "mlp_only_layers", []) + decoder_sparse_step = getattr(config, "decoder_sparse_step", 1) + if (layer_id not in mlp_only_layers) and ( + config.num_experts > 0 and (layer_id + 1) % decoder_sparse_step == 0 + ): + self.is_layer_sparse = True + else: + self.is_layer_sparse = False + + is_mova_config = getattr(config, "num_values", 0) > 0 + is_mova_attention = is_mova_config and layer_id >= config.num_dense_layers + if is_mova_attention: + self.self_attn = XllmMoVAAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + elif is_mova_config: + self.self_attn = XllmGatedAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + else: + self.self_attn = XllmAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + qkv_bias=qkv_bias, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + # Check neighbors for scatter modes + def _is_sparse(lid): + if lid < 0 or lid >= config.num_hidden_layers: + return False + return (lid not in mlp_only_layers) and ( + config.num_experts > 0 and (lid + 1) % decoder_sparse_step == 0 + ) + + is_previous_layer_sparse = _is_sparse(layer_id - 1) + is_next_layer_sparse = _is_sparse(layer_id + 1) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + + if self.is_layer_sparse: + self.mlp = XllmSparseMoeBlock( + layer_id=layer_id, + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = XllmMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.input_layernorm = _make_norm(config) + self.post_attention_layernorm = _make_norm(config) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=(self.layer_id == config.num_hidden_layers - 1), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, + residual, + forward_batch, + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + if isinstance(self.mlp, XllmMLP): + hidden_states = self.mlp( + hidden_states, use_reduce_scatter=use_reduce_scatter + ) + else: + hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) + + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + +class XllmModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + self.vocab_size = config.vocab_size + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + use_attn_tp_group=is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: XllmDecoderLayer( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = _make_norm(config) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + for i in range(self.start_layer, self.end_layer): + ctx = ( + nullcontext() + if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) + else get_global_expert_distribution_recorder().with_current_layer(i) + ) + with ctx: + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + else: + if hidden_states.shape[0] != 0: + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states + + +class XllmForCausalLM(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + _validate_mova_config(config, quant_config) + self.model = XllmModel( + config, + quant_config, + prefix=add_prefix("model", prefix), + ) + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config) + # Value experts are shards of one attention-TP parameter, not FFN/EP + # experts. ParameterMapper therefore stages all 64 canonical HF shards + # before writing the persistent packed tensor during live updates. + self.stacked_params_mapping = _xllm_stacked_params_mapping(config) + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + logits_output = self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + return logits_output + else: + return hidden_states + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = self.stacked_params_mapping + expert_params_mapping = self.expert_params_mapping + strict_checkpoint = getattr(self.config, "model_type", None) in ( + "xllm", + "k2_horizon", + ) + + def is_pipeline_missing_weight(name: str) -> bool: + if not strict_checkpoint: + return False + pp_group = getattr(self, "pp_group", None) + if pp_group is None: + return False + return ( + name == "model.embed_tokens.weight" and not pp_group.is_first_rank + ) or ( + name in ("model.norm.weight", "lm_head.weight") + and not pp_group.is_last_rank + ) + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + checkpoint_name = name + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if "rotary_emb.inv_freq" in name: + continue + if name == "model.embed_tokens.weight" and self.config.tie_word_embeddings: + # With PP>1, the final stage has no embedding table to alias, + # so initialize its separate head from the checkpoint embedding. + if self.pp_group.is_last_rank and "lm_head.weight" in params_dict: + param = params_dict["lm_head.weight"] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + if name == "lm_head.weight" and self.config.tie_word_embeddings: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if weight_name in (".gate_proj", ".up_proj") and ".mlp." not in name: + continue + # Skip experts (handled below in expert_params_mapping) + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + if strict_checkpoint: + raise RuntimeError( + "xLLM-family checkpoint weight did not resolve to " + "a native model parameter: " + f"checkpoint={checkpoint_name!r}, mapped={name!r}" + ) + continue + if name not in params_dict: + if strict_checkpoint: + raise RuntimeError( + "xLLM-family checkpoint weight did not resolve to " + "a native model parameter: " + f"checkpoint={checkpoint_name!r}, mapped={name!r}" + ) + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if is_pipeline_missing_weight(name): + continue + if name.endswith(".bias") and name not in params_dict: + if strict_checkpoint: + raise RuntimeError( + "xLLM-family checkpoint weight did not resolve " + "to a native model parameter: " + f"checkpoint={checkpoint_name!r}, mapped={name!r}" + ) + continue + if name not in params_dict: + if strict_checkpoint: + raise RuntimeError( + "xLLM-family checkpoint weight did not resolve " + "to a native model parameter: " + f"checkpoint={checkpoint_name!r}, mapped={name!r}" + ) + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + @classmethod + def get_model_config_for_expert_location(cls, config): + if getattr(config, "num_experts", 0) <= 0: + return None + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None, + ) + + +class K2HorizonForCausalLM(XllmForCausalLM): + """Load canonical K2Horizon HF checkpoints through the native xLLM path.""" + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + _normalize_k2_horizon_config(config) + super().__init__(config, quant_config=quant_config, prefix=prefix) + + +EntryClass = [XllmForCausalLM, K2HorizonForCausalLM] diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 82488eb75..fd82f7616 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -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": ("", ""), + "medium": ("", ""), + "low": ("", ""), + } + + 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 token. + tool_start_token="") + and ctx.has_text("") + 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("") + and ctx.has_text("") + and ctx.has_text("") + ) + + 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), diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index fdc8d8377..24a807ba5 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -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, diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 5ffa8530c..fd101dabb 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -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): diff --git a/rust/sglang-server/src/message/sampling.rs b/rust/sglang-server/src/message/sampling.rs index ae460695d..a75c3c265 100644 --- a/rust/sglang-server/src/message/sampling.rs +++ b/rust/sglang-server/src/message/sampling.rs @@ -27,6 +27,8 @@ const MAX_STOP_REGEX_LEN: usize = 256; /// Most `stop_regex` patterns accepted per request. Python's `re` cache holds 512 /// (`re._MAXCACHE`), so past that every pattern recompiles on every decode step. const MAX_STOP_REGEX_COUNT: usize = 32; +const REQUEST_REASONING_END_TOKEN_IDS_KEY: &str = "__sglang_reasoning_end_token_ids"; +const MAX_REQUEST_REASONING_END_TOKEN_IDS: usize = 32; /// JSON values accepted by Python's `CustomParamValue`: a scalar, a list of /// scalars, or a string-keyed object whose values are scalars. @@ -584,6 +586,38 @@ impl SamplingParams { } } } + if let Some(value) = self + .custom_params + .as_ref() + .and_then(|params| params.get(REQUEST_REASONING_END_TOKEN_IDS_KEY)) + { + let CustomParamValue::List(token_ids) = value else { + return Err(bad( + "request reasoning end token IDs must be a list of integers".into(), + )); + }; + if token_ids.is_empty() || token_ids.len() > MAX_REQUEST_REASONING_END_TOKEN_IDS { + return Err(bad(format!( + "request reasoning end token IDs must contain 1 to \ + {MAX_REQUEST_REASONING_END_TOKEN_IDS} integers" + ))); + } + for token_id in token_ids { + let in_vocab = match token_id { + JsonScalar::Signed(token_id) => { + *token_id >= 0 && (*token_id as u64) < vocab_size + } + JsonScalar::Unsigned(token_id) => *token_id < vocab_size, + _ => false, + }; + if !in_vocab { + return Err(bad(format!( + "request reasoning end token IDs must be integers in [0, {})", + vocab_size + ))); + } + } + } // Grammars are mutually exclusive. let grammars = [ &self.json_schema, @@ -1057,6 +1091,27 @@ mod tests { } } + #[test] + fn request_reasoning_end_token_ids_are_bounded_integers() { + let valid = norm(r#"{"custom_params":{"__sglang_reasoning_end_token_ids":[17,18]}}"#); + assert!(valid.custom_params.is_some()); + + for body in [ + r#"{"custom_params":{"__sglang_reasoning_end_token_ids":[]}}"#, + r#"{"custom_params":{"__sglang_reasoning_end_token_ids":[-1]}}"#, + r#"{"custom_params":{"__sglang_reasoning_end_token_ids":[true]}}"#, + r#"{"custom_params":{"__sglang_reasoning_end_token_ids":[32000]}}"#, + r#"{"custom_params":{"__sglang_reasoning_end_token_ids":"17"}}"#, + ] { + assert!( + serde_json::from_str::(body) + .unwrap() + .normalize(false, 32_000) + .is_err() + ); + } + } + /// `skip_tokenizer_init` has no tokenizer, so the text-matching stop features /// and `min_new_tokens` (needs eos_token_id) are 400s, not silent no-ops. /// Mirrors Python `raise_if_tokenizer_required`. diff --git a/test/registered/kernels/ops/moe/test_mova_routed_linear.py b/test/registered/kernels/ops/moe/test_mova_routed_linear.py new file mode 100644 index 000000000..e3e149358 --- /dev/null +++ b/test/registered/kernels/ops/moe/test_mova_routed_linear.py @@ -0,0 +1,57 @@ +"""CUDA correctness coverage for MoVA's production routed projection.""" + +import pytest +import torch + +from sglang.srt.layers.mova import routed_linear, routed_linear_reference +from sglang.srt.runtime_context import get_context +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA only") +def test_mova_routed_linear_matches_reference_implementation(): + """Exercise the fused-MoE CUDA path with multiple routes and experts.""" + + torch.manual_seed(0) + num_tokens, num_experts = 7, 5 + input_size, output_size, top_k = 128, 96, 2 + + hidden_states = torch.randn( + num_tokens, input_size, device="cuda", dtype=torch.bfloat16 + ) + expert_weights = ( + torch.randn( + num_experts, + output_size, + input_size, + device="cuda", + dtype=torch.bfloat16, + ) + * input_size**-0.5 + ).contiguous() + selected_experts = torch.tensor( + [[0, 3], [1, 4], [2, 0], [3, 1], [4, 2], [0, 4], [2, 3]], + device="cuda", + dtype=torch.int32, + ) + routing_weights = torch.rand(num_tokens, top_k, device="cuda") + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + + with get_context().override_server_args(): + actual = routed_linear( + hidden_states, expert_weights, routing_weights, selected_experts + ) + expected = routed_linear_reference( + hidden_states, expert_weights, routing_weights, selected_experts + ) + + assert actual.dtype == torch.bfloat16 + torch.testing.assert_close(actual.float(), expected.float(), rtol=2e-2, atol=2e-2) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/constrained/test_base_grammar_backend.py b/test/registered/unit/constrained/test_base_grammar_backend.py index 94dd29451..0c49ce598 100644 --- a/test/registered/unit/constrained/test_base_grammar_backend.py +++ b/test/registered/unit/constrained/test_base_grammar_backend.py @@ -399,6 +399,7 @@ class TestCreateGrammarBackend(unittest.TestCase): result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=[42]) self.assertIsInstance(result, ReasonerGrammarBackend) self.assertIs(result.grammar_backend, mock_backend) + self.assertEqual(result.think_end_ids, [42]) @patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend") def test_no_reasoner_wrapping_without_think_end_ids(self, mock_outlines_cls): diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index 95c57c639..c1031e2b0 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -26,6 +26,9 @@ from sglang.srt.constrained.base_grammar_backend import ( from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.distributed.communication_tags import P2PTag +from sglang.srt.sampling.sampling_params import ( + REQUEST_REASONING_END_TOKEN_IDS_KEY, +) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(2.0, "base-a-test-cpu") @@ -42,6 +45,7 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): scheduler.server_args.reasoning_parser = None scheduler.server_args.constrained_json_whitespace_pattern = None scheduler.server_args.constrained_json_disable_any_whitespace = False + scheduler.model_config.request_selectable_think_end_id_sequences = None # Distributed group mocks scheduler.dp_tp_cpu_group = MagicMock() @@ -303,6 +307,39 @@ class TestProcessReqWithGrammar(unittest.TestCase): self.assertEqual(req.grammar.max_think_tokens, 7) + def test_cache_hit_applies_only_request_selected_terminator(self): + mgr = self._make_mgr() + mgr.scheduler.model_config.request_selectable_think_end_id_sequences = [ + [2, 3], + [8, 9], + ] + grammar_obj = ReasonerGrammarObject( + grammar=None, + think_end_ids=[2, 3], + ) + grammar_obj.maybe_init_reasoning(True) + mgr.grammar_backend.get_cached_or_future_value.return_value = ( + grammar_obj, + True, + ) + + req = _make_req( + json_schema="schema", + custom_params={REQUEST_REASONING_END_TOKEN_IDS_KEY: [8, 9]}, + ) + req.require_reasoning = True + mgr.process_req_with_grammar(req) + + self.assertEqual(req.grammar.think_end_ids, (8, 9)) + + for token_id in (2, 3): + req.grammar.accept_token(token_id) + self.assertTrue(req.grammar._is_thinking()) + + for token_id in (8, 9): + req.grammar.accept_token(token_id) + self.assertTrue(req.grammar._is_generation()) + def test_strict_reasoning_grammar_applies_request_thinking_budget(self): mgr = self._make_mgr() mgr._enable_strict_thinking = True diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 8ea5bc466..969b44e1e 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -43,6 +43,9 @@ from sglang.srt.parser.jinja_template_utils import ( jinja_template_may_reorder_tool_results, ) from sglang.srt.parser.template_detection import ReasoningToggleConfig +from sglang.srt.sampling.sampling_params import ( + REQUEST_REASONING_END_TOKEN_IDS_KEY, +) from sglang.srt.utils import get_or_create_event_loop from sglang.test.ci.ci_register import register_cpu_ci @@ -794,6 +797,33 @@ class ServingChatTestCase(unittest.TestCase): self.assertEqual(req.reasoning_effort, "high") + def test_k2_selected_terminator_reaches_sampling_params(self): + self.tm._config_overrides["reasoning_parser"] = "k2_horizon" + self.chat = OpenAIServingChat(self.tm, self.template_manager) + self.template_manager.chat_template_name = None + self.template_manager.jinja_template_content_format = "string" + self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3] + self.tm.tokenizer.encode.side_effect = lambda text, **_: ( + [8, 9] if text == "" else [1, 2, 3] + ) + req = ChatCompletionRequest( + model="IFM/K2-Horizon-7B", + messages=[{"role": "user", "content": "hi"}], + chat_template_kwargs={"reasoning_effort": "medium"}, + ) + + processed = self.chat._process_messages(req, is_multimodal=False) + self.assertEqual(processed.reasoning_end_token_ids, [8, 9]) + + with patch.object(self.chat, "_process_messages", return_value=processed): + adapted, _ = self.chat._convert_to_internal_request(req) + self.assertEqual( + adapted.sampling_params["custom_params"][ + REQUEST_REASONING_END_TOKEN_IDS_KEY + ], + [8, 9], + ) + def test_kimi_tool_call_keeps_template_default_thinking(self): self.template_manager.chat_template_name = None self.template_manager.jinja_template_content_format = "string" diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses.py b/test/registered/unit/entrypoints/openai/test_serving_responses.py index 458128764..f255739d3 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses.py @@ -23,6 +23,9 @@ from sglang.srt.entrypoints.openai.serving_responses import ( ) from sglang.srt.function_call.core_types import ToolCallItem from sglang.srt.parser.template_detection import ReasoningToggleConfig +from sglang.srt.sampling.sampling_params import ( + REQUEST_REASONING_END_TOKEN_IDS_KEY, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -246,6 +249,38 @@ class ChatToolForwardingTestCase(CustomTestCase): self.assertEqual(request_prompts, [[4, 5, 6]]) self.assertEqual(engine_prompts, [[4, 5, 6]]) + def test_k2_output_parser_reuses_effective_template_default(self): + serving = make_serving() + serving.reasoning_parser = "k2_horizon" + serving.default_chat_template_kwargs = {"reasoning_effort": "low"} + serving.template_manager.chat_template_name = None + serving.tokenizer_manager.tokenizer.apply_chat_template.return_value = [4, 5, 6] + request = ResponsesRequest( + model="IFM/K2-Horizon-7B", + input="hi", + # Template kwargs are the final render inputs, so the server default + # below takes precedence over this API convenience field. + reasoning={"effort": "medium"}, + store=False, + ) + + asyncio.run( + serving._make_request(request, None, serving.tokenizer_manager.tokenizer) + ) + + render_call = serving.tokenizer_manager.tokenizer.apply_chat_template.call_args + self.assertEqual(render_call.kwargs["reasoning_effort"], "low") + self.assertEqual(request.chat_template_kwargs["reasoning_effort"], "low") + + output_items = serving._make_response_output_items( + request, + "work\nanswer", + tokenizer=Mock(), + require_reasoning=True, + ) + self.assertEqual(output_items[0].content[0].text, "work") + self.assertEqual(output_items[1].content[0].text, "\nanswer") + class ReasoningRequestForwardingTestCase(unittest.TestCase): def test_create_responses_uses_processed_reasoning_state(self): @@ -263,6 +298,7 @@ class ReasoningRequestForwardingTestCase(unittest.TestCase): video_data=None, modalities=[], stop=[], + reasoning_end_token_ids=[41, 42], ) captured = {} @@ -308,6 +344,12 @@ class ReasoningRequestForwardingTestCase(unittest.TestCase): self.assertEqual(response.status, "completed") self.assertFalse(captured["adapted_request"].require_reasoning) + self.assertEqual( + captured["adapted_request"].sampling_params["custom_params"][ + REQUEST_REASONING_END_TOKEN_IDS_KEY + ], + [41, 42], + ) self.assertFalse(parser_cls.call_args.kwargs["force_reasoning"]) diff --git a/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py b/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py index c199a7aa6..f0059540c 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py +++ b/test/registered/unit/entrypoints/openai/test_serving_responses_stream.py @@ -32,6 +32,37 @@ class NonHarmonyStreamTestCase(CustomTestCase): self.assertTrue(parser_cls.call_args.kwargs["force_reasoning"]) + def test_k2_nested_effort_selects_streaming_reasoning_delimiter(self): + serving = make_serving() + serving.reasoning_parser = "k2_horizon" + serving.tool_call_parser = None + request = ResponsesRequest( + model="IFM/K2-Horizon-7B", + input="hi", + reasoning={"effort": "medium"}, + stream=True, + store=False, + ) + + events = StreamFixture(serving, request, require_reasoning=True).run( + [engine_chunk("work\nanswer", 4, finish=True)] + ) + types = event_types(events) + payloads = event_payloads(events) + reasoning = "".join( + payload["delta"] + for event_type, payload in zip(types, payloads) + if event_type == "response.reasoning_text.delta" + ) + answer = "".join( + payload["delta"] + for event_type, payload in zip(types, payloads) + if event_type == "response.output_text.delta" + ) + + self.assertEqual(reasoning, "work") + self.assertEqual(answer, "\nanswer") + def test_emits_typed_sse_events_in_order(self): serving = make_serving() serving.reasoning_parser = None diff --git a/test/registered/unit/function_call/test_k2_v3_detector.py b/test/registered/unit/function_call/test_k2_v3_detector.py new file mode 100644 index 000000000..613e6d80b --- /dev/null +++ b/test/registered/unit/function_call/test_k2_v3_detector.py @@ -0,0 +1,325 @@ +import json +import unittest + +from sglang.srt.entrypoints.openai.protocol import Function, Tool +from sglang.srt.environ import envs +from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.function_call.k2_v3_detector import K2V3Detector +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def make_weather_tool() -> Tool: + return Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + "options": {"type": "object"}, + }, + }, + ), + ) + + +def make_reasoning_with_tool_marker(tag: str) -> str: + return ( + f"Consider " + "get_weather" + "city" + "Boston" + " as hypothetical text." + f"\n" + ) + + +class TestK2V3Detector(CustomTestCase): + def setUp(self): + self.tools = [make_weather_tool()] + + def test_xml_and_typed_values(self): + text = ( + "get_weather\n" + "city\n" + "string\n" + " Boston \n" + "days\n" + "integer\n" + "3\n" + "options\n" + '{"units": "metric"}\n' + "" + ) + result = K2V3Detector().detect_and_parse(text, self.tools) + self.assertEqual(result.normal_text, "") + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual( + json.loads(result.calls[0].parameters), + { + "city": " Boston ", + "days": 3, + "options": {"units": "metric"}, + }, + ) + + def test_json_content_is_detected_from_wire(self): + wire = ( + "\n" + '{"name":"get_weather","arguments":{"city":"Tokyo","days":2}}' + "\n" + ) + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual( + json.loads(result.calls[0].parameters), + {"city": "Tokyo", "days": 2}, + ) + + def test_json_strings_are_not_reinterpreted_without_a_schema(self): + wire = ( + "" + '{"name":"get_weather","arguments":' + '{"numeric":"123","boolean":"true","object":"{}"}}' + "" + ) + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual( + json.loads(result.calls[0].parameters), + {"numeric": "123", "boolean": "true", "object": "{}"}, + ) + + def test_local_schema_ref_preserves_xml_string_value(self): + tool = Tool( + type="function", + function=Function( + name="lookup", + parameters={ + "type": "object", + "$defs": {"Identifier": {"type": "string"}}, + "properties": {"id": {"$ref": "#/$defs/Identifier"}}, + }, + ), + ) + wire = ( + "lookup" + "id" + "123" + "" + ) + result = K2V3Detector().detect_and_parse(wire, [tool]) + self.assertEqual(json.loads(result.calls[0].parameters), {"id": "123"}) + + def test_xml_typed_union_uses_wire_value_type(self): + tool = Tool( + type="function", + function=Function( + name="lookup", + parameters={ + "type": "object", + "properties": { + "id": {"anyOf": [{"type": "string"}, {"type": "integer"}]} + }, + }, + ), + ) + wire = ( + "lookup" + "id" + "integer" + "123" + "" + ) + result = K2V3Detector().detect_and_parse(wire, [tool]) + self.assertEqual(json.loads(result.calls[0].parameters), {"id": 123}) + + def test_non_string_json_function_name_is_forwarded_as_text(self): + wire = '{"name":123,"arguments":{}}' + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(result.calls, []) + self.assertEqual(result.normal_text, wire) + + def test_parallel_calls_keep_wire_order_and_indices(self): + wire = ( + "\n" + "get_weather" + "city" + "Tokyo" + "\n" + "get_weather" + "city" + "Boston" + "\n" + "" + ) + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(result.normal_text, "") + self.assertEqual([call.tool_index for call in result.calls], [0, 1]) + self.assertEqual( + [json.loads(call.parameters)["city"] for call in result.calls], + ["Tokyo", "Boston"], + ) + + def test_non_streaming_preserves_reasoning_for_ordinary_answer(self): + wire = ( + " \n" + make_reasoning_with_tool_marker("think_fast") + "The answer is 42." + ) + parser = FunctionCallParser(self.tools, "k2_horizon") + normal, calls = parser.parse_non_stream(wire) + self.assertEqual(normal, wire) + self.assertEqual(calls, []) + + def test_streaming_preserves_reasoning_for_ordinary_answer(self): + wire = ( + " \n" + + make_reasoning_with_tool_marker("think_faster") + + "The answer is 42." + ) + parser = FunctionCallParser(self.tools, "k2_horizon") + normal = "" + calls = [] + for char in wire: + new_normal, new_calls = parser.parse_stream_chunk(char) + normal += new_normal + calls.extend(new_calls) + end_normal, end_calls = parser.parse_stream_end() + normal += end_normal + calls.extend(end_calls) + + self.assertEqual(normal, wire) + self.assertEqual(calls, []) + + def test_non_streaming_tool_call_preserves_reasoning_prefix(self): + reasoning = make_reasoning_with_tool_marker("think") + wire = reasoning + ( + "\n" + "get_weather" + "city" + "Tokyo" + "\n" + "" + ) + parser = FunctionCallParser(self.tools, "k2_horizon") + normal, calls = parser.parse_non_stream(wire) + self.assertEqual(normal, reasoning) + self.assertEqual(len(calls), 1) + self.assertEqual(json.loads(calls[0].parameters), {"city": "Tokyo"}) + + def test_forced_reasoning_without_opening_tag_is_preserved(self): + reasoning = "work\n" + wire = reasoning + ( + "get_weather" + "city" + "Tokyo" + "" + ) + normal, calls = FunctionCallParser(self.tools, "k2_horizon").parse_non_stream( + wire + ) + self.assertEqual(normal, reasoning) + self.assertEqual(len(calls), 1) + self.assertEqual(json.loads(calls[0].parameters), {"city": "Tokyo"}) + + def test_streaming_tool_call_preserves_reasoning_at_every_boundary(self): + reasoning = make_reasoning_with_tool_marker("think") + wire = reasoning + ( + "\n" + "get_weather" + "city" + "東京" + "\n" + "" + ) + parser = FunctionCallParser(self.tools, "k2_horizon") + normal = "" + calls = [] + for char in wire: + new_normal, new_calls = parser.parse_stream_chunk(char) + normal += new_normal + calls.extend(new_calls) + end_normal, end_calls = parser.parse_stream_end() + normal += end_normal + calls.extend(end_calls) + + self.assertEqual(normal, reasoning) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0].name, "get_weather") + self.assertEqual(json.loads(calls[0].parameters), {"city": "東京"}) + + def test_malformed_complete_block_is_forwarded_as_text(self): + wire = ( + "prefixget_weather" + "city" + "suffix" + ) + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(result.calls, []) + self.assertEqual(result.normal_text, wire) + + def test_malformed_canonical_group_is_forwarded_intact(self): + wire = ( + "prefix\n" + "get_weather" + "city" + "\n" + "suffix" + ) + result = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(result.calls, []) + self.assertEqual(result.normal_text, wire) + + def test_unknown_tool_respects_forwarding_policy(self): + wire = "unknown" + with envs.SGLANG_FORWARD_UNKNOWN_TOOLS.override(False): + dropped = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(dropped.calls, []) + + with envs.SGLANG_FORWARD_UNKNOWN_TOOLS.override(True): + forwarded = K2V3Detector().detect_and_parse(wire, self.tools) + self.assertEqual(len(forwarded.calls), 1) + self.assertEqual(forwarded.calls[0].name, "unknown") + + def test_any_schema_preserves_json_value(self): + tool = make_weather_tool() + tool.function.parameters["properties"]["options"] = {"type": "any"} + wire = ( + "" + '{"name":"get_weather","arguments":{"options":{"units":"metric"}}}' + "" + ) + result = K2V3Detector().detect_and_parse(wire, [tool]) + self.assertEqual( + json.loads(result.calls[0].parameters), + {"options": {"units": "metric"}}, + ) + + def test_unterminated_stream_is_released_on_finish(self): + wire = "prefixget_weathercity" + detector = K2V3Detector() + streamed = detector.parse_streaming_increment(wire, self.tools) + self.assertEqual(streamed.normal_text, "prefix") + self.assertEqual(streamed.calls, []) + finished = detector.finish(self.tools) + self.assertEqual( + finished.normal_text, + "get_weathercity", + ) + self.assertEqual(finished.calls, []) + + def test_function_call_parser_registry(self): + parser = FunctionCallParser( + tools=self.tools, + tool_call_parser="k2_horizon", + ) + self.assertIsInstance(parser.detector, K2V3Detector) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_mova.py b/test/registered/unit/layers/test_mova.py new file mode 100644 index 000000000..6d69a9d2a --- /dev/null +++ b/test/registered/unit/layers/test_mova.py @@ -0,0 +1,136 @@ +# 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. + +import torch + +from sglang.srt.layers.mova import ( + RoutedValueExperts, + _prepare_mova_moe_config, + mova_router_topk, + routed_linear, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def test_router_bias_changes_selection_but_not_mixture_weights(): + logits = torch.tensor([[2.0, 1.0, 0.0]], dtype=torch.float32) + bias = torch.tensor([0.0, 0.0, 10.0], dtype=torch.float32) + + weights, selected = mova_router_topk( + logits, + bias, + score_func="sigmoid", + top_k=1, + scaling_factor=2.5, + ) + + assert selected.tolist() == [[2]] + torch.testing.assert_close(weights, torch.sigmoid(logits[:, 2:3]) * 2.5) + + +def test_router_renormalizes_before_scaling(): + logits = torch.tensor([[2.0, 1.0, -2.0]], dtype=torch.float32) + weights, selected = mova_router_topk( + logits, + None, + score_func="sigmoid", + top_k=2, + scaling_factor=2.5, + ) + + scores = torch.sigmoid(logits) + expected_ids = torch.topk(scores, 2, dim=-1).indices + expected = torch.gather(scores, -1, expected_ids) + expected = expected / expected.sum(-1, keepdim=True) * 2.5 + torch.testing.assert_close(selected, expected_ids.to(torch.int32)) + torch.testing.assert_close(weights, expected) + + +def test_softmax_router_bias_is_selection_only(): + logits = torch.tensor([[3.0, 2.0, -4.0]], dtype=torch.float32) + bias = torch.tensor([0.0, 0.0, 20.0], dtype=torch.float32) + + weights, selected = mova_router_topk( + logits, + bias, + score_func="softmax", + top_k=1, + scaling_factor=1.0, + ) + + unbiased_scores = torch.softmax(logits, dim=-1) + assert selected.tolist() == [[2]] + torch.testing.assert_close(weights, unbiased_scores[:, 2:3]) + + +def test_routed_linear_cpu_matches_independent_expected_math(): + hidden = torch.tensor([[1.0, -2.0], [0.5, 3.0]]) + expert_weights = torch.arange(3 * 4 * 2, dtype=torch.float32).reshape(3, 4, 2) + selected = torch.tensor([[0, 2], [1, 0]], dtype=torch.int32) + weights = torch.tensor([[0.25, 0.75], [0.6, 0.4]]) + + expected_rows = [] + for token, expert_ids, mixture in zip(hidden, selected, weights): + projections = torch.stack( + [ + torch.nn.functional.silu(expert_weights[expert_id] @ token) + for expert_id in expert_ids.tolist() + ] + ) + expected_rows.append((projections * mixture[:, None]).sum(dim=0)) + + expected = torch.stack(expected_rows) + actual = routed_linear(hidden, expert_weights, weights, selected) + torch.testing.assert_close(actual, expected) + + +def test_value_expert_loader_shards_output_dimension(): + experts = RoutedValueExperts( + num_experts=2, + input_size=3, + output_size=4, + tp_rank=1, + tp_size=2, + ) + packed = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3) + experts.weight_loader(experts.weight, packed) + torch.testing.assert_close(experts.weight, packed[:, 2:]) + + single = torch.arange(4 * 3, dtype=torch.float32).reshape(4, 3) + 100 + experts.weight_loader(experts.weight, single, loaded_shard_id=0) + torch.testing.assert_close(experts.weight[0], single[2:]) + + +def test_mova_moe_config_drops_unsupported_tma_without_mutating_source(): + source = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "USE_TMA": True, + } + + prepared = _prepare_mova_moe_config(source) + + assert "USE_TMA" not in prepared + assert source["USE_TMA"] is True + + +if __name__ == "__main__": + import sys + + import pytest + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index a61365397..f8547049e 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -9,11 +9,15 @@ from types import SimpleNamespace import torch +from sglang.srt.disaggregation.decode import DecodeRequest, DecodeTransferQueue from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, ) -from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.sampling.sampling_params import ( + REQUEST_REASONING_END_TOKEN_IDS_KEY, + SamplingParams, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -110,6 +114,48 @@ def _make_result(num_draft_tokens, accept_lens, flat_tokens): ) +def _commit_disagg_handoff( + req: Req, + processor: SchedulerBatchResultProcessor, + token_id: int, + *, + replayed_boundary: bool = False, +) -> None: + queue = DecodeTransferQueue.__new__(DecodeTransferQueue) + queue.scheduler = SimpleNamespace(batch_result_processor=processor) + queue.spec_algorithm = SimpleNamespace(is_none=lambda: True) + queue.metadata_buffers = SimpleNamespace( + get_buf=lambda _: ( + torch.tensor([token_id], dtype=torch.long), + torch.zeros(7, dtype=torch.long), + torch.zeros(1), + torch.zeros(1, dtype=torch.long), + torch.zeros(1), + torch.zeros(1, dtype=torch.long), + None, + None, + None, + None, + None, + None, + None, + torch.tensor([1], dtype=torch.long), + ) + ) + req.bootstrap_host = "127.0.0.1" + req.bootstrap_room = 1 + if replayed_boundary: + req.pd_rebootstrap_forced_output_id = token_id + decode_req = DecodeRequest( + req=req, + kv_receiver=SimpleNamespace(clear=lambda: None), + metadata_buffer_index=0, + is_rebootstrap=replayed_boundary, + ) + + queue._commit_transfer_to_req(decode_req) + + class TestSpecV2GrammarTruncation(CustomTestCase): def test_resolve_truncates_after_grammar_completion(self): req = _make_req(terminate_after=2) @@ -147,6 +193,66 @@ class TestReasoningTokenAccounting(CustomTestCase): self.assertEqual(req.reasoning_tokens, 3) self.assertTrue(req._is_reasoning_over) + def test_request_selected_end_ignores_other_closer(self): + req = _make_req(terminate_after=99) + req.require_reasoning = True + req.sampling_params.custom_params = { + REQUEST_REASONING_END_TOKEN_IDS_KEY: [17, 18] + } + processor = _make_processor() + processor.model_config.think_end_ids = [7, 8] + processor.model_config.request_selectable_think_end_id_sequences = [ + [7, 8], + [17, 18], + ] + + # The global/default closer must not end a medium request. + processor._maybe_update_reasoning_tokens(req, [10, 7]) + processor._maybe_update_reasoning_tokens(req, [8, 11]) + self.assertFalse(req._is_reasoning_over) + + processor._maybe_update_reasoning_tokens(req, [10, 17]) + processor._maybe_update_reasoning_tokens(req, [18, 11]) + + self.assertEqual(req.reasoning_tokens, 7) + self.assertTrue(req._is_reasoning_over) + + def test_disagg_handoff_can_start_multi_token_selected_end(self): + req = _make_req(terminate_after=99) + req.require_reasoning = True + req.sampling_params.custom_params = { + REQUEST_REASONING_END_TOKEN_IDS_KEY: [17, 18] + } + processor = _make_processor() + processor.model_config.request_selectable_think_end_id_sequences = [ + [7, 8], + [17, 18], + ] + + _commit_disagg_handoff(req, processor, 17) + self.assertEqual(req.reasoning_tokens, 1) + self.assertFalse(req._is_reasoning_over) + + processor._maybe_update_reasoning_tokens(req, 18) + + self.assertEqual(req.reasoning_tokens, 2) + self.assertTrue(req._is_reasoning_over) + + def test_disagg_rebootstrap_does_not_recount_boundary(self): + req = _make_req(terminate_after=99) + req.require_reasoning = True + req.sampling_params.custom_params = {REQUEST_REASONING_END_TOKEN_IDS_KEY: [17]} + processor = _make_processor() + processor.model_config.request_selectable_think_end_id_sequences = [ + [7], + [17], + ] + + _commit_disagg_handoff(req, processor, 17, replayed_boundary=True) + + self.assertEqual(req.reasoning_tokens, 0) + self.assertFalse(req._is_reasoning_over) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_xllm.py b/test/registered/unit/models/test_xllm.py new file mode 100644 index 000000000..b7607dae3 --- /dev/null +++ b/test/registered/unit/models/test_xllm.py @@ -0,0 +1,349 @@ +# 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. + +import math +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.configs.k2_horizon import K2HorizonConfig, XllmConfig +from sglang.srt.models.xllm import ( + EntryClass, + K2HorizonForCausalLM, + XllmAttention, + XllmForCausalLM, + XllmGroupRMSNorm, + _normalize_k2_horizon_config, + _validate_mova_config, + _xllm_router_gemm, + _xllm_stacked_params_mapping, + permute_to_hf, + permute_to_xllm, +) +from sglang.srt.runtime_context import get_context +from sglang.srt.utils.hf_transformers.common import _CONFIG_REGISTRY +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class _IdentityRotary: + def __call__(self, positions, q, k): + self.positions_shape = tuple(positions.shape) + self.q_shape = tuple(q.shape) + self.k_shape = tuple(k.shape) + return q, k + + +def test_native_config_and_model_registration(): + assert _CONFIG_REGISTRY["xllm"] is XllmConfig + assert _CONFIG_REGISTRY["k2_horizon"] is K2HorizonConfig + assert XllmForCausalLM in EntryClass + assert K2HorizonForCausalLM in EntryClass + + +def test_dense_horizon_yarn_schema_normalization(): + config = K2HorizonConfig.from_dict( + { + "architectures": ["K2HorizonForCausalLM"], + "model_type": "k2_horizon", + "hidden_size": 16, + "num_hidden_layers": 2, + "mlp_only_layers": [0, 1], + "rope_theta": 1_000_000.0, + "max_position_embeddings": 32, + "rope_parameters": { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 8, + "attention_factor": 1.1, + "beta_fast": 32.0, + "beta_slow": 1.0, + "truncate": True, + }, + } + ) + + _normalize_k2_horizon_config(config) + + assert config.num_values == 0 + assert config.num_values_per_tok == 0 + assert config.num_experts == 0 + assert config.num_experts_per_tok == 0 + assert config.num_shared_experts == 0 + assert config.num_dense_layers == 2 + assert config.rope_scaling["rope_type"] == "yarn" + assert config.rope_scaling["original_max_position_embeddings"] == 8 + assert config.rope_scaling["attn_factor"] == pytest.approx( + 1.1 / (0.1 * math.log(4.0) + 1.0) + ) + assert config._sglang_xllm_checkpoint_format == "k2_horizon_hf" + + +def test_mova_horizon_schema_requires_and_applies_source_router_contract(): + config = K2HorizonConfig.from_dict( + { + "architectures": ["K2HorizonForCausalLM"], + "model_type": "k2_horizon", + "hidden_size": 16, + "num_hidden_layers": 4, + "num_experts": 8, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "mova_num_experts": 4, + "mova_num_experts_per_tok": 2, + "mlp_only_layers": [0], + "attention_gate_func": "softplus", + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10_000_000.0, + }, + "xllm_source_router_gemm_partitions": 2, + } + ) + + _normalize_k2_horizon_config(config) + + assert config.num_values == 4 + assert config.num_values_per_tok == 2 + assert config.num_dense_layers == 1 + assert config.apply_attn_gate is True + assert config.attn_gate_func == "softplus" + assert config.rope_theta == 10_000_000.0 + + +def test_mova_horizon_schema_rejects_missing_source_router_contract(): + config = K2HorizonConfig.from_dict( + { + "hidden_size": 16, + "num_hidden_layers": 4, + "mova_num_experts": 4, + "mova_num_experts_per_tok": 2, + } + ) + + with pytest.raises(ValueError, match="source router GEMM provenance"): + _normalize_k2_horizon_config(config) + + +def test_group_rms_norm_matches_groupwise_reference_and_residual_contract(): + norm = XllmGroupRMSNorm(hidden_size=4, n_groups=2, eps=0.0) + with torch.no_grad(): + norm.weight.copy_(torch.tensor([1.0, 2.0, 3.0, 4.0])) + + hidden = torch.tensor([[3.0, 4.0, 0.0, 5.0]]) + residual = torch.ones_like(hidden) + combined = hidden + residual + grouped = combined.reshape(1, 2, 2) + expected = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True)) + expected = expected.reshape(1, 4) * norm.weight + + output, returned_residual = norm(hidden, residual=residual) + torch.testing.assert_close(output, expected) + torch.testing.assert_close(returned_residual, combined) + + +def test_partial_rope_round_trips_non_rotary_dimensions(): + attention = object.__new__(XllmAttention) + torch.nn.Module.__init__(attention) + attention.num_heads = 2 + attention.num_kv_heads = 1 + attention.head_dim = 8 + attention.rope_head_dim = 4 + attention.rotary_emb = _IdentityRotary() + + positions = torch.arange(3) + q = torch.randn(3, attention.num_heads * attention.head_dim) + k = torch.randn(3, attention.num_kv_heads * attention.head_dim) + q_out, k_out = XllmAttention._apply_partial_rope(attention, positions, q, k) + + torch.testing.assert_close(q_out, q) + torch.testing.assert_close(k_out, k) + assert attention.rotary_emb.positions_shape == (3,) + assert attention.rotary_emb.q_shape == (3, 8) + assert attention.rotary_emb.k_shape == (3, 4) + + +def test_permutation_helpers_use_expected_interleave_order(): + hf_value = torch.arange(8, dtype=torch.float32).reshape(1, 1, 8) + xllm_value = torch.tensor([[[0.0, 4.0, 1.0, 5.0, 2.0, 6.0, 3.0, 7.0]]]) + + torch.testing.assert_close(permute_to_xllm(hf_value), xllm_value) + torch.testing.assert_close(permute_to_hf(xllm_value), hf_value) + + +def test_mp2_router_gemm_preserves_source_rounding_order(): + hidden = torch.tensor([[1.25, -0.75, 0.5, 2.0]], dtype=torch.bfloat16) + weight = torch.tensor( + [[0.25, 1.5, -1.0, 0.75], [2.0, -0.5, 1.25, 0.5]], + dtype=torch.bfloat16, + ) + hidden_parts = hidden.chunk(2, dim=-1) + weight_parts = weight.chunk(2, dim=-1) + expected = ( + torch.nn.functional.linear( + hidden_parts[0].contiguous(), weight_parts[0].contiguous() + ).float() + + torch.nn.functional.linear( + hidden_parts[1].contiguous(), weight_parts[1].contiguous() + ).float() + ) + + torch.testing.assert_close(_xllm_router_gemm(hidden, weight, 2), expected) + with pytest.raises(ValueError, match="requires BF16"): + _xllm_router_gemm(hidden.float(), weight.float(), 2) + + +def test_mova_weight_mapping_uses_checkpoint_shaped_attention_projections(): + config = SimpleNamespace(num_values=4) + mapping = _xllm_stacked_params_mapping(config) + + assert mapping == [ + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + (".v_experts.weight", ".v_experts.0.weight", 0), + (".v_experts.weight", ".v_experts.1.weight", 1), + (".v_experts.weight", ".v_experts.2.weight", 2), + (".v_experts.weight", ".v_experts.3.weight", 3), + ] + + +def test_dense_weight_mapping_packs_qkv_and_gate_up(): + config = SimpleNamespace(num_values=0) + + assert _xllm_stacked_params_mapping(config) == [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + +def test_strict_loader_rejects_unknown_checkpoint_weight(): + model = object.__new__(XllmForCausalLM) + torch.nn.Module.__init__(model) + model.config = SimpleNamespace(model_type="xllm", tie_word_embeddings=False) + model.model = SimpleNamespace(start_layer=0, end_layer=1) + model.pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + model.stacked_params_mapping = [] + model.expert_params_mapping = [] + + with pytest.raises(RuntimeError, match="did not resolve"): + model.load_weights([("unexpected.weight", torch.ones(1))]) + + +def _native_runtime_config(*, enable_two_batch_overlap=False, **overrides): + values = { + "enable_eplb": False, + "init_expert_location": "trivial", + "ep_num_redundant_experts": 0, + "enable_two_batch_overlap": enable_two_batch_overlap, + } + values.update(overrides) + return values + + +def test_native_xllm_requires_bfloat16(monkeypatch): + config = XllmConfig(num_values=0, num_experts=192) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.float16) + + with ( + get_context().override_server_args(**_native_runtime_config()), + pytest.raises(ValueError, match="requires --dtype bfloat16"), + ): + _validate_mova_config(config, quant_config=None) + + +def test_native_xllm_rejects_quantized_weights(monkeypatch): + config = XllmConfig(num_values=0, num_experts=0) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) + + with ( + get_context().override_server_args(**_native_runtime_config()), + pytest.raises(ValueError, match="does not support quantized"), + ): + _validate_mova_config(config, quant_config=object()) + + +def test_native_xllm_accepts_bfloat16_without_expert_remapping(monkeypatch): + config = XllmConfig(num_values=0, num_experts=192) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) + + with get_context().override_server_args(**_native_runtime_config()): + _validate_mova_config(config, quant_config=None) + + +@pytest.mark.parametrize( + ("config_override", "error"), + [ + ({"query_key_norm": True}, "query/key normalization"), + ({"sliding_window": 4096}, "full causal attention only"), + ({"use_sliding_window": True}, "full causal attention only"), + ({"apply_attn_gate": True}, "gated attention"), + ], + ids=["qk-norm", "sliding-window", "use-sliding-window", "attention-gate"], +) +def test_native_dense_xllm_rejects_unimplemented_attention_features( + monkeypatch, config_override, error +): + config = XllmConfig(num_values=0, num_experts=192, **config_override) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) + + with ( + get_context().override_server_args(**_native_runtime_config()), + pytest.raises(ValueError, match=error), + ): + _validate_mova_config(config, quant_config=None) + + +def test_native_xllm_rejects_two_batch_overlap(monkeypatch): + config = XllmConfig(num_values=0, num_experts=192) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) + + with ( + get_context().override_server_args( + **_native_runtime_config(enable_two_batch_overlap=True) + ), + pytest.raises(ValueError, match="does not yet support.*two-batch-overlap"), + ): + _validate_mova_config(config, quant_config=None) + + +@pytest.mark.parametrize( + "runtime_override", + [ + {"enable_eplb": True}, + {"init_expert_location": "random"}, + {"ep_num_redundant_experts": 1}, + ], + ids=["eplb", "initial-placement", "redundant-expert"], +) +def test_native_xllm_rejects_unmapped_expert_modes(monkeypatch, runtime_override): + config = XllmConfig(num_values=0, num_experts=192) + monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) + + with ( + get_context().override_server_args( + **_native_runtime_config(**runtime_override) + ), + pytest.raises(ValueError, match="does not yet support EPLB"), + ): + _validate_mova_config(config, quant_config=None) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/parser/test_k2_v3_reasoning_parser.py b/test/registered/unit/parser/test_k2_v3_reasoning_parser.py new file mode 100644 index 000000000..c0db21b13 --- /dev/null +++ b/test/registered/unit/parser/test_k2_v3_reasoning_parser.py @@ -0,0 +1,116 @@ +import unittest + +from sglang.srt.entrypoints.openai.protocol import ( + ChatCompletionRequest, + ResponsesRequest, +) +from sglang.srt.parser.reasoning_parser import K2V3Detector, ReasoningParser +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestK2V3ReasoningParser(CustomTestCase): + def test_reasoning_effort_selects_ifm_pair(self): + expected = { + "high": ("", ""), + "medium": ("", ""), + "low": ("", ""), + } + for effort, tokens in expected.items(): + with self.subTest(effort=effort): + detector = K2V3Detector(reasoning_effort=effort) + self.assertEqual( + (detector.think_start_token, detector.think_end_token), tokens + ) + self.assertEqual( + set(detector.request_selectable_think_end_tokens), + { + "", + "", + "", + }, + ) + + def test_release_template_fallback_uses_medium_pair(self): + # K2-Horizon-0.9B maps unsupported levels to . + detector = K2V3Detector(reasoning_effort="none") + self.assertEqual(detector.think_start_token, "") + + def test_end_only_output_preserves_newlines(self): + result = K2V3Detector().detect_and_parse("\nreasoning\n\nanswer") + self.assertEqual(result.reasoning_text, "\nreasoning\n") + self.assertEqual(result.normal_text, "\nanswer") + + def test_tool_group_implicitly_ends_malformed_reasoning(self): + result = K2V3Detector().detect_and_parse( + "reasoning\nx" + ) + self.assertEqual(result.reasoning_text, "reasoning\n") + self.assertTrue(result.normal_text.startswith("")) + + def test_streaming_partial_tags(self): + detector = K2V3Detector(reasoning_effort="medium") + reasoning = "" + normal = "" + wire = "work\nanswer" + for char in wire: + result = detector.parse_streaming_increment(char) + reasoning += result.reasoning_text + normal += result.normal_text + end = detector.finish() + reasoning += end.reasoning_text + normal += end.normal_text + self.assertEqual(reasoning, "work") + self.assertEqual(normal, "\nanswer") + + def test_reasoning_parser_reads_request_effort(self): + request = ChatCompletionRequest( + model="IFM/K2-Horizon-7B", + messages=[{"role": "user", "content": "hi"}], + reasoning_effort="low", + ) + parser = ReasoningParser("k2_horizon", request=request) + self.assertIsInstance(parser.detector, K2V3Detector) + self.assertEqual(parser.detector.think_end_token, "") + + def test_template_kwarg_effort_has_rendering_precedence(self): + request = ChatCompletionRequest( + model="IFM/K2-Horizon-7B", + messages=[{"role": "user", "content": "hi"}], + reasoning_effort="high", + chat_template_kwargs={"reasoning_effort": "low"}, + ) + parser = ReasoningParser("k2_horizon", request=request) + self.assertEqual(parser.detector.think_end_token, "") + + def test_responses_request_effort_selects_non_stream_delimiter(self): + end_tokens = { + "high": "", + "medium": "", + "low": "", + } + for effort, end_token in end_tokens.items(): + with self.subTest(effort=effort): + request = ResponsesRequest( + model="IFM/K2-Horizon-7B", + input="hi", + reasoning={"effort": effort}, + store=False, + ) + parser = ReasoningParser( + "k2_horizon", stream_reasoning=False, request=request + ) + self.assertEqual( + parser.parse_non_stream(f"work{end_token}\nanswer"), + ("work", "\nanswer"), + ) + + def test_force_reasoning_cannot_be_disabled(self): + with self.assertRaisesRegex(ValueError, "requires force_reasoning=True"): + ReasoningParser("k2_horizon", force_reasoning=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/parser/test_template_manager.py b/test/registered/unit/parser/test_template_manager.py index de3f26f26..f1b887e2d 100644 --- a/test/registered/unit/parser/test_template_manager.py +++ b/test/registered/unit/parser/test_template_manager.py @@ -483,6 +483,28 @@ class TestTemplateDetectionRuleMatrix(unittest.TestCase): self.assertEqual(config.toggle_param, "enable_thinking") self.assertTrue(config.default_enabled) + def test_k2_horizon_detects_always_on_reasoning_and_tool_parsers(self): + template = """ + {% set tool_call_fmt = tool_call_format | default('xml') %} + {% set effort = reasoning_effort | default('high') %} + + + """ + force, config = detect_reasoning_pattern(template) + tokenizer = _DummyTokenizer([]) + + self.assertTrue(force) + self.assertIsNotNone(config) + self.assertEqual(config.special_case, "always") + self.assertEqual( + detect_reasoning_parser(template, tokenizer, config, force), + "k2_horizon", + ) + self.assertEqual( + detect_tool_call_parser(template, tokenizer, config, force), + "k2_horizon", + ) + class TestToolCallParserDetection(unittest.TestCase): """Tests for detect_tool_call_parser() using real model tokenizers.""" diff --git a/test/registered/unit/sampling/test_sampling_params.py b/test/registered/unit/sampling/test_sampling_params.py index 8f897000d..cd1fe3cfc 100644 --- a/test/registered/unit/sampling/test_sampling_params.py +++ b/test/registered/unit/sampling/test_sampling_params.py @@ -16,9 +16,11 @@ import msgspec from sglang.srt.sampling.sampling_params import ( MAX_LEN, + MAX_REQUEST_REASONING_END_TOKEN_IDS, MAX_STOP_COUNT, MAX_STOP_REGEX_COUNT, MAX_STOP_REGEX_LEN, + REQUEST_REASONING_END_TOKEN_IDS_KEY, TOP_K_ALL, SamplingParams, get_max_seq_length, @@ -108,6 +110,28 @@ class TestSamplingParamsVerify(CustomTestCase): sp = self._make() sp.verify(self.VOCAB_SIZE) + def test_request_reasoning_end_token_ids_are_vocab_bounded_integers(self): + self._make( + custom_params={REQUEST_REASONING_END_TOKEN_IDS_KEY: [17, 18]} + ).verify(self.VOCAB_SIZE) + + invalid_values = [ + [], + [-1], + [True], + [self.VOCAB_SIZE], + "17", + list(range(MAX_REQUEST_REASONING_END_TOKEN_IDS + 1)), + ] + for value in invalid_values: + with ( + self.subTest(value=value), + self.assertRaisesRegex(ValueError, "request reasoning end token IDs"), + ): + self._make( + custom_params={REQUEST_REASONING_END_TOKEN_IDS_KEY: value} + ).verify(self.VOCAB_SIZE) + def test_negative_temperature_raises(self): """Test that verify() rejects negative temperature (must be >= 0).""" sp = self._make(temperature=-0.5) diff --git a/test/registered/unit/utils/test_hf_transformers.py b/test/registered/unit/utils/test_hf_transformers.py index 94f666e33..7f7f1e5ca 100644 --- a/test/registered/unit/utils/test_hf_transformers.py +++ b/test/registered/unit/utils/test_hf_transformers.py @@ -581,6 +581,11 @@ class TestAttachAdditionalStopTokenIds(unittest.TestCase): attach_additional_stop_token_ids(tok) self.assertEqual(tok.additional_stop_token_ids, {128008}) + def test_k2_horizon_im_end_registers_as_stop(self): + tok = self._tokenizer({"<|ifm|im_end|>": 64019}) + attach_additional_stop_token_ids(tok) + self.assertEqual(tok.additional_stop_token_ids, {64019}) + def test_no_known_marker_yields_none(self): tok = self._tokenizer({"<|other|>": 7}) attach_additional_stop_token_ids(tok)