minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Model] Add native IFM K2 Horizon serving support (#37654)

Browse Source 
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
Yash Akhauri
2026-09-03 16:39:43 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang Xinyuan Tong
parent 02d9b3060a
commit 3bac084d4e
34 changed files with 4487 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/constrained/test_base_grammar_backend.py
+1
View File
@@ -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):
test/registered/unit/constrained/test_grammar_manager.py
+37
View File
@@ -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
test/registered/unit/entrypoints/openai/test_serving_chat.py
+30
View File
@@ -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 == "</ifm|think_fast>" 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"
test/registered/unit/entrypoints/openai/test_serving_responses.py
+42
View File
@@ -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</ifm|think_faster>\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"])
test/registered/unit/entrypoints/openai/test_serving_responses_stream.py
+31
View File
@@ -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</ifm|think_fast>\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
test/registered/unit/function_call/test_k2_v3_detector.py
+325
View File
@@ -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"<ifm|{tag}>Consider "
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>Boston</ifm|arg_value>"
"</ifm|tool_call> as hypothetical text."
f"</ifm|{tag}>\n"
)
class TestK2V3Detector(CustomTestCase):
def setUp(self):
self.tools = [make_weather_tool()]
def test_xml_and_typed_values(self):
text = (
"<ifm|tool_call>get_weather\n"
"<ifm|arg_key>city</ifm|arg_key>\n"
"<ifm|arg_type>string</ifm|arg_type>\n"
"<ifm|arg_value> Boston </ifm|arg_value>\n"
"<ifm|arg_key>days</ifm|arg_key>\n"
"<ifm|arg_type>integer</ifm|arg_type>\n"
"<ifm|arg_value>3</ifm|arg_value>\n"
"<ifm|arg_key>options</ifm|arg_key>\n"
'<ifm|arg_value>{"units": "metric"}</ifm|arg_value>\n'
"</ifm|tool_call>"
)
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 = (
"<ifm|tool_calls>\n<ifm|tool_call>"
'{"name":"get_weather","arguments":{"city":"Tokyo","days":2}}'
"</ifm|tool_call>\n</ifm|tool_calls>"
)
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 = (
"<ifm|tool_call>"
'{"name":"get_weather","arguments":'
'{"numeric":"123","boolean":"true","object":"{}"}}'
"</ifm|tool_call>"
)
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 = (
"<ifm|tool_call>lookup"
"<ifm|arg_key>id</ifm|arg_key>"
"<ifm|arg_value>123</ifm|arg_value>"
"</ifm|tool_call>"
)
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 = (
"<ifm|tool_call>lookup"
"<ifm|arg_key>id</ifm|arg_key>"
"<ifm|arg_type>integer</ifm|arg_type>"
"<ifm|arg_value>123</ifm|arg_value>"
"</ifm|tool_call>"
)
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 = '<ifm|tool_call>{"name":123,"arguments":{}}</ifm|tool_call>'
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 = (
"<ifm|tool_calls>\n"
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>Tokyo</ifm|arg_value>"
"</ifm|tool_call>\n"
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>Boston</ifm|arg_value>"
"</ifm|tool_call>\n"
"</ifm|tool_calls>"
)
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 + (
"<ifm|tool_calls>\n"
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>Tokyo</ifm|arg_value>"
"</ifm|tool_call>\n"
"</ifm|tool_calls>"
)
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</ifm|think>\n"
wire = reasoning + (
"<ifm|tool_calls><ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>Tokyo</ifm|arg_value>"
"</ifm|tool_call></ifm|tool_calls>"
)
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 + (
"<ifm|tool_calls>\n"
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"<ifm|arg_value>東京</ifm|arg_value>"
"</ifm|tool_call>\n"
"</ifm|tool_calls>"
)
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 = (
"prefix<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"</ifm|tool_call>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<ifm|tool_calls>\n"
"<ifm|tool_call>get_weather"
"<ifm|arg_key>city</ifm|arg_key>"
"</ifm|tool_call>\n"
"</ifm|tool_calls>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 = "<ifm|tool_call>unknown</ifm|tool_call>"
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 = (
"<ifm|tool_call>"
'{"name":"get_weather","arguments":{"options":{"units":"metric"}}}'
"</ifm|tool_call>"
)
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 = "prefix<ifm|tool_call>get_weather<ifm|arg_key>city"
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,
"<ifm|tool_call>get_weather<ifm|arg_key>city",
)
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()
test/registered/unit/layers/test_mova.py
+136
View File
@@ -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"]))
test/registered/unit/managers/test_batch_result_processor_spec_grammar.py
+107 -1
View File
@@ -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()
test/registered/unit/models/test_xllm.py
+349
View File
@@ -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"]))
test/registered/unit/parser/test_k2_v3_reasoning_parser.py
+116
View File
@@ -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": ("<ifm|think>", "</ifm|think>"),
"medium": ("<ifm|think_fast>", "</ifm|think_fast>"),
"low": ("<ifm|think_faster>", "</ifm|think_faster>"),
}
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),
{
"</ifm|think>",
"</ifm|think_fast>",
"</ifm|think_faster>",
},
)
def test_release_template_fallback_uses_medium_pair(self):
# K2-Horizon-0.9B maps unsupported levels to <ifm|think_fast>.
detector = K2V3Detector(reasoning_effort="none")
self.assertEqual(detector.think_start_token, "<ifm|think_fast>")
def test_end_only_output_preserves_newlines(self):
result = K2V3Detector().detect_and_parse("\nreasoning\n</ifm|think>\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\n<ifm|tool_calls><ifm|tool_call>x</ifm|tool_call>"
)
self.assertEqual(result.reasoning_text, "reasoning\n")
self.assertTrue(result.normal_text.startswith("<ifm|tool_calls>"))
def test_streaming_partial_tags(self):
detector = K2V3Detector(reasoning_effort="medium")
reasoning = ""
normal = ""
wire = "work</ifm|think_fast>\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, "</ifm|think_faster>")
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, "</ifm|think_faster>")
def test_responses_request_effort_selects_non_stream_delimiter(self):
end_tokens = {
"high": "</ifm|think>",
"medium": "</ifm|think_fast>",
"low": "</ifm|think_faster>",
}
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()
test/registered/unit/parser/test_template_manager.py
+22
View File
@@ -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') %}
<ifm|think><ifm|think_fast><ifm|think_faster>
<ifm|tool_calls><ifm|tool_call>
"""
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."""
test/registered/unit/sampling/test_sampling_params.py
+24
View File
@@ -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)
test/registered/unit/utils/test_hf_transformers.py
+5
View File
@@ -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)