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

[misc] Remove unit test cases that fail the admission criteria (#30690)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-09 15:31:28 -07:00
committed by GitHub
parent 7e936f690e
commit c53559ba10
20 changed files with 59 additions and 4005 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.codespellrc
+1 -1
View File
@@ -1,3 +1,3 @@
[codespell]
ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, dout, IST
ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias, allready, inout, nd, fo, visibles, nothink, renderD, ond, tbe, CopyIn, notin, subtile, subtiles, dout, IST, kInf
skip = *.json, *.jsonl, *.patch, *.txt, *.lock
python/sglang/srt/parser/conversation.py
-2
View File
@@ -358,8 +358,6 @@ class Conversation:
ret = system_prompt + self.sep
for role, message in self.messages:
if message:
if type(message) is tuple:
message, _, _ = message
ret += role + message + self.sep
else:
ret += role
test/registered/unit/entrypoints/anthropic/test_serving.py
+3 -55
View File
@@ -134,16 +134,6 @@ async def _collect_anthropic_events(serving, anthropic_request):
class TestAnthropicServing(unittest.TestCase):
# System-first guard (Qwen-style): rejects non-first system → must merge.
QWEN_SYSTEM_FIRST_TEMPLATE = (
"{%- for message in messages %}"
"{%- if message.role == 'system' and not loop.first %}"
"{{- raise_exception('system must be first') }}"
"{%- endif %}"
"{{- message.role }}: {{ message.content }}\n"
"{%- endfor %}"
)
# Renders system at any position (GLM/Kimi/Qwen3) → can pass through.
INLINE_SYSTEM_TEMPLATE = (
"{%- for message in messages %}"
@@ -675,19 +665,6 @@ class TestAnthropicServing(unittest.TestCase):
self.assertEqual(anthropic_response.content[1].type, "text")
self.assertEqual(anthropic_response.content[1].text, "the answer is 4")
def test_request_thinking_enabled_invokes_apply_reasoning_enabled(self):
"""``thinking={"type":"enabled", "budget_tokens":N}`` flips reasoning on.
``budget_tokens`` is required by the SDK shape on ``enabled``; the
local backend does not enforce it but accepts the value.
"""
serving = self._serving()
request = self._anthropic_request(
thinking={"type": "enabled", "budget_tokens": 1024}, stream=False
)
serving._convert_to_chat_completion_request(request)
self.assertEqual(serving.openai_serving_chat.apply_reasoning_calls, [True])
def test_request_thinking_disabled_invokes_apply_reasoning_enabled(self):
"""``thinking={"type": "disabled"}`` must flip the reasoning toggle off."""
serving = self._serving()
@@ -828,34 +805,17 @@ class TestAnthropicServing(unittest.TestCase):
self.assertEqual(chat_request.max_tokens, 16)
self.assertTrue(any("task_budget" in r and "32768" in r for r in log.output))
def test_request_task_budget_with_remaining_is_accepted(self):
"""SDK's ``BetaTokenTaskBudgetParam`` has a ``remaining`` field
used for client-side compaction. Must round-trip cleanly."""
serving = self._serving()
request = self._anthropic_request(
output_config={
"task_budget": {"type": "tokens", "total": 32768, "remaining": 12000}
},
stream=False,
)
# Must not raise; pre-existing logging still works.
serving._convert_to_chat_completion_request(request)
self.assertEqual(request.output_config.task_budget.remaining, 12000)
def test_request_betas_is_accepted_and_logged(self):
"""The Anthropic SDK attaches ``betas`` to many requests; must not 400."""
"""``betas`` is accepted and logged; the local backend has no beta system."""
import logging
serving = self._serving()
request = self._anthropic_request(
betas=["thinking-2025-08-04", "computer-use-2025-01-24"],
stream=False,
)
request = self._anthropic_request(betas=["thinking-2025-08-04"], stream=False)
with self.assertLogs(
"sglang.srt.entrypoints.anthropic.serving", level=logging.INFO
) as log:
serving._convert_to_chat_completion_request(request)
self.assertTrue(any("betas" in r for r in log.output))
self.assertTrue(any("thinking-2025-08-04" in r for r in log.output))
def test_assistant_thinking_history_is_rewrapped_for_chat_template(self):
"""Past-turn thinking blocks get re-emitted via wrap_reasoning_history."""
@@ -1152,18 +1112,6 @@ class TestAnthropicServing(unittest.TestCase):
serving._convert_to_chat_completion_request(request)
self.assertIn("tool_choice", str(ctx.exception))
def test_server_tool_only_with_tool_choice_auto_is_allowed(self):
"""tool_choice=auto over server-only tools is a no-op (model decides)."""
serving = self._serving()
request = self._anthropic_request(
stream=False,
tools=[{"type": "web_search_20250305", "name": "web_search"}],
tool_choice={"type": "auto"},
)
# Must not raise; the request just runs with no client-side tools.
chat_request = serving._convert_to_chat_completion_request(request)
self.assertIsNone(chat_request.tools)
def test_tool_choice_named_custom_tool_is_resolved(self):
"""tool_choice={type:'tool', name:'X'} where X is a custom tool wires through."""
serving = self._serving()
test/registered/unit/entrypoints/openai/test_serving_chat.py
-169
View File
@@ -240,49 +240,6 @@ class ServingChatTestCase(unittest.TestCase):
self.assertTrue(adapted.require_reasoning)
def test_kimi_tool_call_keeps_explicit_reasoning(self):
self.template_manager.reasoning_config = ReasoningToggleConfig(
toggle_param="thinking", default_enabled=True
)
self.tm.server_args.reasoning_parser = "kimi_k2"
self.tm.server_args.tool_call_parser = "kimi_k2"
self.chat.reasoning_parser = "kimi_k2"
self.chat.tool_call_parser = "kimi_k2"
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "What is 2+2?"}],
tools=[
{
"type": "function",
"function": {
"name": "add",
"parameters": {
"type": "object",
"properties": {"a": {"type": "integer"}},
},
},
}
],
tool_choice="required",
chat_template_kwargs={"thinking": True},
)
with patch.object(self.chat, "_process_messages") as proc_mock:
proc_mock.return_value = MessageProcessingResult(
"",
[1, 2, 3],
None,
None,
[],
[],
None,
)
adapted, _ = self.chat._convert_to_internal_request(req)
self.assertTrue(adapted.require_reasoning)
def test_kimi_tool_call_respects_explicit_reasoning_disable(self):
self.template_manager.reasoning_config = ReasoningToggleConfig(
toggle_param="thinking", default_enabled=True
@@ -895,98 +852,6 @@ class ServingChatTestCase(unittest.TestCase):
)
# ------------- kimi_k2 tool_call_id formatting -------------
def test_kimi_k2_non_streaming_tool_call_id_format(self):
"""Ensure non-streaming tool_call.id matches functions.{name}:{index} for kimi_k2 parser."""
# Force kimi_k2 parser
self.chat.tool_call_parser = "kimi_k2"
# Mock FunctionCallParser.parse_non_stream to return one tool call
with patch(
"sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser"
) as ParserMock:
parser_instance = ParserMock.return_value
# Build a mock ToolCallItem-like object
call_info = Mock()
call_info.name = "get_weather"
call_info.parameters = '{"city":"Paris"}'
call_info.tool_index = 0
parser_instance.has_tool_call.return_value = True
parser_instance.parse_non_stream.return_value = ("", [call_info])
finish_reason = {"type": "stop", "matched": None}
tools = [
{"type": "function", "function": {"name": "get_weather"}},
]
tool_calls, remaining_text, finish_reason = self.chat._process_tool_calls(
text="<|tool_calls_section_begin|>...",
tools=tools,
finish_reason=finish_reason,
)
self.assertIsNotNone(tool_calls)
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0].id, "functions.get_weather:0")
self.assertEqual(tool_calls[0].function.name, "get_weather")
def test_kimi_k2_streaming_tool_call_id_format(self):
"""Ensure streaming first chunk tool_call.id matches functions.{name}:{index} for kimi_k2 parser."""
# Force kimi_k2 parser
self.chat.tool_call_parser = "kimi_k2"
# Prepare request with tools
req = ChatCompletionRequest(
model="x",
messages=[{"role": "user", "content": "Hi?"}],
tools=[{"type": "function", "function": {"name": "get_weather"}}],
stream=True,
)
# Patch FunctionCallParser used inside _process_tool_call_stream
with patch(
"sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser"
) as ParserMock:
parser_instance = ParserMock.return_value
# First call returns one ToolCallItem-like chunk (with name)
first_chunk_call = Mock()
first_chunk_call.tool_index = 0
first_chunk_call.name = "get_weather"
first_chunk_call.parameters = ""
parser_instance.parse_stream_chunk.side_effect = [
("", [first_chunk_call]),
("", []),
]
async def collect_first_tool_chunk():
gen = self.chat._process_tool_call_stream(
index=0,
delta="irrelevant",
parser_dict={},
content={"meta_info": {"id": "chatcmpl-test"}},
request=req,
has_tool_calls={},
)
# Get first yielded SSE line
line = None
async for emitted in gen:
line = emitted
break
return line
loop = get_or_create_event_loop()
line = loop.run_until_complete(collect_first_tool_chunk())
self.assertIsNotNone(line)
self.assertTrue(line.startswith("data: "))
payload = json.loads(line[len("data: ") :])
tool_calls = payload["choices"][0]["delta"]["tool_calls"]
self.assertEqual(tool_calls[0]["id"], "functions.get_weather:0")
def test_kimi_k2_non_streaming_tool_call_id_with_history(self):
"""Ensure non-streaming tool_call.id increase with tool calls history for kimi_k2 parser."""
@@ -1225,28 +1090,6 @@ class ServingChatTestCase(unittest.TestCase):
task="bogus",
)
def test_latest_reminder_role_accepted(self):
"""`latest_reminder` is a first-class message role on generic param."""
from sglang.srt.entrypoints.openai.protocol import (
ChatCompletionMessageGenericParam,
)
msg = ChatCompletionMessageGenericParam(
role="latest_reminder", content="Be terse."
)
self.assertEqual(msg.role, "latest_reminder")
# Full request with reminder before user parses cleanly.
req = ChatCompletionRequest(
model="x",
messages=[
{"role": "latest_reminder", "content": "Be terse."},
{"role": "user", "content": "Hi"},
],
)
self.assertEqual(req.messages[0].role, "latest_reminder")
self.assertEqual(req.messages[1].role, "user")
def test_attach_task_to_last_user_message(self):
"""Helper attaches task to the nearest user/developer message."""
from sglang.srt.entrypoints.openai import encoding_dsv4
@@ -1954,14 +1797,6 @@ class ServingChatTestCase(unittest.TestCase):
)
self.assertIsNone(result)
def test_extract_routed_dp_rank_from_header_with_header(self):
"""Test that header value is extracted correctly."""
self.fastapi_request.headers = {"x-data-parallel-rank": "2"}
result = self.chat.extract_routed_dp_rank_from_header(
self.fastapi_request, body_routed_dp_rank=None
)
self.assertEqual(result, 2)
def test_extract_routed_dp_rank_header_overrides_body(self):
"""Test that header value has higher priority than body."""
self.fastapi_request.headers = {"x-data-parallel-rank": "3"}
@@ -2404,10 +2239,6 @@ class TestProcessToolCallsWithRequiredToolChoice(unittest.TestCase):
class TestNormalizeToolContent(unittest.TestCase):
"""Unit tests for normalize_tool_content()."""
def test_openai_text_parts_flattened(self):
result = normalize_tool_content("tool", [{"type": "text", "text": "10525"}])
self.assertEqual(result, "10525")
def test_multiple_text_parts_joined(self):
result = normalize_tool_content(
"tool",
test/registered/unit/function_call/test_function_call_parser.py
+19 -908
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
-39
View File
@@ -35,7 +35,6 @@ if _HAS_MLX:
MlxAuxiliaryStateReqToTokenPool,
MlxModelCacheLayout,
find_attention_layers,
is_attention_module,
patch_model_attention,
)
from sglang.srt.hardware_backend.mlx.model_runner import (
@@ -156,9 +155,6 @@ class TestMlxAttentionPatching(unittest.TestCase):
self.assertFalse(isinstance(model.layers[0].linear_attn, MLXAttentionWrapper))
self.assertIsInstance(model.layers[1].self_attn, MLXAttentionWrapper)
def test_projection_only_mixer_is_not_attention(self):
self.assertFalse(is_attention_module(ProjectionOnlyMixer()))
def test_cache_layout_separates_attention_and_auxiliary_layers(self):
layout = MlxModelCacheLayout.from_attention_discovery(
[object(), object(), object(), object()],
@@ -1012,41 +1008,6 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertIsNone(req.mamba_last_track_seqlen)
self.assertEqual(pool.auxiliary_state_pool.available_size(), 2)
def test_auxiliary_state_component_keeps_new_live_slot_owned_by_radix(self):
pool = MlxAuxiliaryStateReqToTokenPool(
size=2,
max_context_len=8,
device="cpu",
enable_memory_saver=False,
auxiliary_state_size=4,
)
req = FakeRequest()
pool.alloc([req])
component = MlxAuxiliaryStateComponent(
SimpleNamespace(req_to_token_pool=pool),
SimpleNamespace(enable_mamba_extra_buffer=False),
)
insert_params = InsertParams()
cache_len = component.prepare_for_caching_req(
req=req,
insert_params=insert_params,
token_ids_len=7,
is_finished=True,
)
component.cleanup_after_caching_req(
req=req,
is_finished=True,
insert_result=InsertResult(prefix_len=0, mamba_exist=False),
insert_params=insert_params,
)
self.assertEqual(cache_len, 7)
self.assertFalse(getattr(insert_params, "mlx_auxiliary_state_uses_track_slot"))
self.assertEqual(insert_params.mamba_value.tolist(), [1])
self.assertIsNone(req.mamba_pool_idx)
self.assertEqual(pool.auxiliary_state_pool.available_size(), 3)
def test_auxiliary_state_component_frees_stale_track_slot_when_live_slot_inserted(
self,
):
test/registered/unit/mem_cache/test_hicache_nixl_storage.py
-173
View File
@@ -9,10 +9,8 @@ import shutil
import socket
import subprocess
import tempfile
import threading
import time
import unittest
from typing import List
import torch
@@ -513,177 +511,6 @@ class TestNixlUnified(CustomTestCase):
self.assertEqual(self.hicache.batch_exists(["key1", "key2"]), 1)
def _run_concurrent_stress(
self, is_zero_copy_mode: bool, hicache: HiCacheNixl = None
):
"""One getter thread + one setter thread share the same HiCacheNixl
for ``is_zero_copy_mode``. Defaults to ``self.hicache`` (FILE backend);
pass ``hicache`` to exercise a different backend (e.g. OBJ).
Phase 1 pre-seeds N preset pages and stores them under fixed keys.
Phase 2 runs the getter (reads the presets back and verifies content)
concurrently with the setter (writes a stream of fresh distinct keys
from a disjoint source region). The kv_buffer regions touched by the
two threads are disjoint so any data corruption observed is from the
backend's shared state (bounce buffers, devId maps, fd pool).
"""
if hicache is None:
hicache = self.hicache
# 8 preset pages, 8 getter dst pages, 8 setter src pages -> 24 in use.
mock_host = MockMemPoolHost(is_zero_copy_mode=is_zero_copy_mode, num_pages=32)
hicache.register_mem_pool_host(mock_host)
hicache.is_zero_copy = is_zero_copy_mode
page_size = mock_host.page_size
dtype = mock_host.dtype
num_pages = 8
# Disjoint per-thread regions in kv_buffer (indexed by token index).
preset_src = (0, num_pages)
getter_dst = (num_pages, 2 * num_pages)
setter_src = (2 * num_pages, 3 * num_pages)
# zero_copy=page_first uses dim 1 for the token axis; non-zero-copy=
# layer_first uses dim 2. All buffer accesses below go through this so
# the rest of the harness stays layout-agnostic.
def token_index(start_token: int, n_tokens: int):
s = slice(start_token, start_token + n_tokens)
if is_zero_copy_mode:
return (slice(None), s, slice(None), slice(None), slice(None))
return (slice(None), slice(None), s, slice(None), slice(None))
def page_index(start_page: int, n_pages: int):
return token_index(start_page * page_size, n_pages * page_size)
def fill_pages(start_page: int, n_pages: int, value_fn):
"""value_fn(i) -> scalar value for page i."""
for i in range(n_pages):
idx = page_index(start_page + i, 1)
shape = mock_host.kv_buffer[idx].shape
mock_host.kv_buffer[idx] = torch.full(
shape, float(value_fn(i)), dtype=dtype
)
# Phase 1: distinct value per preset page so a wrong-page result is
# detectable; setter source is constant (value irrelevant to the
# test, just needs to be valid).
fill_pages(preset_src[0], num_pages, lambda i: i + 1)
fill_pages(setter_src[0], num_pages, lambda i: -1.0)
preset_keys = [f"preset_{int(is_zero_copy_mode)}_{i}" for i in range(num_pages)]
preset_indices = torch.arange(
preset_src[0] * page_size,
preset_src[1] * page_size,
dtype=torch.int64,
)
self.assertTrue(
all(hicache.batch_set_v1(preset_keys, preset_indices)),
"phase 1: presetting keys failed",
)
# Expected per-page-i payload after a successful get into getter_dst.
expected_pages = [
mock_host.kv_buffer[page_index(preset_src[0] + i, 1)].clone()
for i in range(num_pages)
]
# Phase 2.
stop = threading.Event()
errors: List[str] = []
errors_lock = threading.Lock()
def record_error(msg: str):
with errors_lock:
errors.append(msg)
def getter_loop():
dst_indices = torch.arange(
getter_dst[0] * page_size,
getter_dst[1] * page_size,
dtype=torch.int64,
)
loops = 0
while not stop.is_set():
# Zero the dst pages so a no-op get is observable.
mock_host.kv_buffer[page_index(getter_dst[0], num_pages)] = 0.0
ok = hicache.batch_get_v1(preset_keys, dst_indices)
if not all(ok):
record_error(f"getter loop {loops}: batch_get_v1 returned {ok}")
return
for i in range(num_pages):
got = mock_host.kv_buffer[page_index(getter_dst[0] + i, 1)]
if not torch.equal(got, expected_pages[i]):
record_error(f"getter loop {loops}: preset page {i} corrupted")
return
loops += 1
def setter_loop():
src_indices = torch.arange(
setter_src[0] * page_size,
setter_src[1] * page_size,
dtype=torch.int64,
)
loops = 0
while not stop.is_set():
keys = [
f"setter_{int(is_zero_copy_mode)}_{loops}_{i}"
for i in range(num_pages)
]
ok = hicache.batch_set_v1(keys, src_indices)
if not all(ok):
record_error(f"setter loop {loops}: batch_set_v1 returned {ok}")
return
loops += 1
t_get = threading.Thread(target=getter_loop, daemon=True)
t_set = threading.Thread(target=setter_loop, daemon=True)
t_get.start()
t_set.start()
# Bounded run: long enough to interleave many ops under NIXL I/O
# GIL release, short enough for a unit test.
time.sleep(3.0)
stop.set()
t_get.join(timeout=10)
t_set.join(timeout=10)
self.assertFalse(
t_get.is_alive() or t_set.is_alive(),
"stress threads failed to stop",
)
self.assertEqual(errors, [], f"concurrency errors: {errors}")
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
def test_concurrent_getter_setter_file_zero_copy(self):
"""Stress: concurrent getter+setter, FILE backend, zero-copy."""
self._run_concurrent_stress(is_zero_copy_mode=True)
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
def test_concurrent_getter_setter_file_non_zero_copy(self):
"""Stress: concurrent getter+setter, FILE backend, non-zero-copy."""
self._run_concurrent_stress(is_zero_copy_mode=False)
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
@unittest.skipUnless(
MinioFixture.is_available(), "minio binary or boto3 not available"
)
def test_concurrent_getter_setter_obj_zero_copy(self):
"""Stress: concurrent getter+setter, OBJ backend (MinIO), zero-copy."""
self._run_concurrent_stress(
is_zero_copy_mode=True, hicache=self._make_obj_hicache()
)
@unittest.skipUnless(STRESS_ENABLED, "set SGLANG_RUN_NIXL_STRESS=1 to run")
@unittest.skipUnless(
MinioFixture.is_available(), "minio binary or boto3 not available"
)
def test_concurrent_getter_setter_obj_non_zero_copy(self):
"""Stress: concurrent getter+setter, OBJ backend (MinIO), non-zero-copy."""
self._run_concurrent_stress(
is_zero_copy_mode=False, hicache=self._make_obj_hicache()
)
@unittest.skipUnless(hasattr(os, "O_DIRECT"), "O_DIRECT not available on this platform")
class TestNixlDirectIO(CustomTestCase):
test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py
-89
View File
@@ -776,58 +776,6 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
self.assertEqual(captured["host_indices"].device.type, "cpu")
self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu")
def test_hybrid_write_moves_indices_without_page_first_layout(self):
captured = {}
class FakeHostGroup:
layout = "layer_first"
can_use_write_back_jit = True
def backup_from_device_all_layer(
self,
device_pool,
host_indices,
device_indices,
io_backend,
pool_transfers=None,
):
captured["host_indices"] = host_indices
captured["pool_transfers"] = pool_transfers
op = CacheOperation(
host_indices=_indices(0, 4),
device_indices=_indices(4, 8),
node_id=1,
pool_transfers=[
PoolTransfer(
name=PoolName.DEEPSEEK_V4_C4,
host_indices=_indices(0, 4),
device_indices=_indices(4, 8),
)
],
)
controller = HybridCacheController.__new__(HybridCacheController)
controller.write_queue = [op]
controller.io_backend = "kernel"
controller.mem_pool_host = FakeHostGroup()
controller.mem_pool_device = None
controller.has_draft = False
controller.write_stream = object()
controller.ack_write_queue = []
controller._record_transfer_indices_on_stream = lambda *args: None
controller.move_hybrid_indices = mock.Mock(
return_value=(op.host_indices, op.device_indices, op.pool_transfers)
)
with mock.patch.object(
hybrid_cache_controller, "device_module", _FakeDeviceModule
):
controller.start_writing()
controller.move_hybrid_indices.assert_called_once()
self.assertEqual(captured["host_indices"].device.type, "cpu")
self.assertEqual(captured["pool_transfers"][0].host_indices.device.type, "cpu")
def test_write_back_jit_cache_controller_keeps_host_indices_on_cpu(self):
captured = {}
@@ -905,43 +853,6 @@ class TestHiCacheStagedWriteBackDispatch(unittest.TestCase):
controller.move_indices.assert_called_once()
self.assertEqual(captured["host_indices"].device.type, "cpu")
def test_cache_controller_moves_indices_without_page_first_layout(self):
captured = {}
class FakeHostPool:
layout = "layer_first"
can_use_write_back_jit = True
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
captured["host_indices"] = host_indices
op = ManagerCacheOperation(
host_indices=_indices(0, 4),
device_indices=_indices(4, 8),
node_id=1,
)
controller = HiCacheController.__new__(HiCacheController)
controller.write_queue = [op]
controller.io_backend = "kernel"
controller.mem_pool_host = FakeHostPool()
controller.mem_pool_device = None
controller.has_draft = False
controller.write_stream = object()
controller.ack_write_queue = []
controller.move_indices = mock.Mock(
return_value=(op.host_indices, op.device_indices)
)
with mock.patch.object(
manager_cache_controller, "device_module", _FakeDeviceModule
):
controller.start_writing()
controller.move_indices.assert_called_once()
self.assertEqual(captured["host_indices"].device.type, "cpu")
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+5 -231
View File
@@ -408,16 +408,6 @@ class TestMultiEndedAllocator(unittest.TestCase):
expected = full_alloc.virtual_to_physical[v]
self.assertTrue(bool((buf == expected).all().item()))
def test_translate_kv_loc_without_out_returns_fresh_tensor(self):
"""REGRESSION: without `out=`, behavior returns a fresh tensor."""
_, full_alloc, _, full_kv, _ = self._build_pair()
v = self._alloc(full_alloc, full_kv, 5)
ret = full_alloc.translate_kv_loc(v)
# Fresh tensor: different storage from v2p table
self.assertNotEqual(ret.data_ptr(), full_alloc.virtual_to_physical.data_ptr())
expected = full_alloc.virtual_to_physical[v]
self.assertTrue(bool((ret == expected).all().item()))
def test_translate_kv_loc_out_matches_no_out(self):
"""REGRESSION: result of translate_kv_loc(v, out=buf) byte-equals
translate_kv_loc(v)."""
@@ -833,62 +823,6 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0)
self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0)
# 7. Joint byte-budget pre-check.
def test_swa_joint_byte_budget_pre_check(self):
# Pick sizes where the byte gap, not slot-index headroom, is the bind.
full_spec = MHASubPoolSpec(
name="full",
layer_num=2,
head_num=2,
head_dim=4,
store_dtype=torch.float16,
grow_direction="up",
)
swa_spec = MHASubPoolSpec(
name="swa",
layer_num=2,
head_num=2,
head_dim=4,
store_dtype=torch.float16,
grow_direction="down",
)
n_full, n_swa = 10, 10
total = n_full * full_spec.entry_bytes() + n_swa * swa_spec.entry_bytes()
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full_spec, swa_spec],
device=_DEV,
enable_memory_saver=False,
)
kvcache = _FakeUnifiedSWAKVPool(pool)
allocator = UnifiedSWATokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=kvcache,
device=_DEV,
full_max_total_num_tokens=n_full,
swa_max_total_num_tokens=n_swa,
need_sort=False,
forward_stream=None,
)
fa = allocator.full_attn_allocator
sa = allocator.swa_attn_allocator
# Compute the "naive min" against the joint budget — at idle, the
# joint budget is strictly less than min(full.available, swa.available)
# because the joint uses (entry_full + entry_swa) per slot.
naive = min(fa.available_size(), sa.available_size())
joint = allocator.available_size()
# The joint must be no greater than naive (typically strictly less).
self.assertLessEqual(joint, naive)
# And it must equal `gap_bytes // (entry_full + entry_swa)` clamped
# by slot-room.
gap = sa._byte_low_frontier() - fa._byte_high_frontier()
expected = min(
gap // (fa.entry_bytes + sa.entry_bytes),
fa.max_slots - fa.min_slot_index - fa.allocated_count(),
sa.max_slots - sa.min_slot_index - sa.allocated_count(),
)
self.assertEqual(joint, expected)
# 8. Watermark rollback on partial alloc failure.
def test_swa_alloc_swa_failure_is_fail_loud(self):
"""The SWA composite runs a tight JOINT pre-check before allocating, so
@@ -1194,6 +1128,11 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
self.assertNotEqual(int(full_alloc.virtual_to_physical[v_page].item()), -1)
# 5. free() recovers pages via unique(// page_size) — matches upstream.
# REGRESSION: `allocated_count()` MUST return
# TOKENS, not pages -- matching upstream's convention that all external
# capacity methods report tokens. At page_size > 1, returning pages
# here breaks the leak invariant
# (`available + evictable + ... == total`, with all terms in tokens).
def test_paged_free_unique_by_page(self):
_, full_alloc, _, full_kv, _ = self._build()
a = full_alloc.alloc(self.PAGE_SIZE * 2) # 2 pages = 2*PS tokens
@@ -1501,31 +1440,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
int(free_before.numel()),
)
# 12. translate_kv_loc preserves token-level identity end-to-end.
def test_paged_translate_kv_loc_token_round_trip(self):
_, full_alloc, _, _, _ = self._build()
v = full_alloc.alloc(self.PAGE_SIZE * 2)
# Build the composite-style translation manually: virt_page * ps + offset.
ps = self.PAGE_SIZE
virt_pages = v // ps
offsets = v % ps
phys_pages = full_alloc.virtual_to_physical[virt_pages]
phys_tokens = phys_pages * ps + offsets
# `phys_tokens` should be a coherent set of two contiguous PAGES.
phys_pages_unique = sorted(set(phys_pages.tolist()))
self.assertEqual(len(phys_pages_unique), 2)
# Within each page the tokens go through offsets 0..7 in order.
for p in phys_pages_unique:
page_phys = sorted(
int(t)
for i, t in enumerate(phys_tokens.tolist())
if int(phys_pages[i].item()) == p
)
self.assertEqual(
page_phys,
[p * ps + i for i in range(ps)],
)
# REGRESSION: `translate_kv_loc(virt, out=buf)` must work
# under page_size > 1 — the page-math branch writes via
# `index_select(out=out)` + in-place `mul_` / `add_` and must match the
@@ -1679,28 +1593,6 @@ class TestPagedMultiEndedAllocator(unittest.TestCase):
)
self.assertTrue(bool((buf[:ps] == 0).all().item()))
# 13. REGRESSION: `allocated_count()` MUST return
# TOKENS, not pages — matching upstream's convention that all external
# capacity methods report tokens. At page_size > 1, returning pages
# here breaks the leak invariant
# (`available + evictable + ... == total`, with all terms in tokens).
def test_paged_allocated_count_returns_tokens(self):
_, full_alloc, _, _, _ = self._build()
PS = self.PAGE_SIZE
# Idle → allocated_count == 0.
self.assertEqual(full_alloc.allocated_count(), 0)
# Alloc 2 pages = 2 * PS tokens.
v = full_alloc.alloc(2 * PS)
self.assertIsNotNone(v)
# allocated_count() must report TOKENS (= 2 * PS), not pages (= 2).
self.assertEqual(
full_alloc.allocated_count(),
2 * PS,
"REGRESSION: allocated_count() must return TOKENS at page_size > 1",
)
# _allocated_pages() is the page-granular internal helper.
self.assertEqual(full_alloc._allocated_pages(), 2)
# 14. REGRESSION: the leak-invariant terms used by the
# scheduler runtime checker must all be in TOKENS. Specifically
# `full_available_size() + allocated_tokens == static_cap` must hold for
@@ -2032,24 +1924,6 @@ class TestLazyCompaction(unittest.TestCase):
p = int(alloc.virtual_to_physical[v].item())
kv.buf[p] = int(v)
def test_lazy_state_initialized(self):
"""Lazy allocator initializes the new state cleanly."""
_pool, fa, _kv = self._make_full(lazy=True)
self.assertTrue(fa.lazy_compaction)
self.assertEqual(len(fa._free_phys_pages), 0)
self.assertEqual(fa._pending_reuse, {})
self.assertEqual(fa.live_page_count, 0)
# Watermark + free virtual list start equivalent to eager.
self.assertEqual(fa.watermark_physical, fa.min_page_index)
def test_lazy_alloc_increments_live_page_count(self):
_pool, fa, _kv = self._make_full(lazy=True)
tokens = fa.alloc(8)
self.assertIsNotNone(tokens)
self.assertEqual(int(tokens.numel()), 8)
self.assertEqual(fa.live_page_count, 8)
self.assertEqual(len(fa._free_phys_pages), 0)
def test_lazy_free_boundary_shortcut(self):
"""Boundary absorption is DEFERRED to `_flush` (the hot
path `_free_lazy` does only a `torch.cat`, no watermark mutation).
@@ -2075,21 +1949,6 @@ class TestLazyCompaction(unittest.TestCase):
self.assertEqual(len(fa._free_phys_pages), 0)
self.assertEqual(fa.live_page_count, 2)
def test_lazy_free_non_boundary_pushes_hole(self):
"""Freeing a non-boundary page enters _free_phys_pages, watermark
stays put.
"""
_pool, fa, _kv = self._make_full(lazy=True)
a = fa.alloc(5)
wm_before = fa.watermark_physical
# Free a middle id (NOT the topmost), boundary-shortcut should
# NOT fire.
mid = a[2:3].clone()
fa.free(mid)
self.assertEqual(fa.watermark_physical, wm_before)
self.assertEqual(len(fa._free_phys_pages), 1)
self.assertEqual(fa.live_page_count, 4)
def test_lazy_free_inward_walk(self):
"""The inward walk (multiple contiguous holes absorbed
into the watermark in one pass) is DEFERRED to `_flush`. After
@@ -2248,29 +2107,6 @@ class TestLazyCompaction(unittest.TestCase):
)
self.assertEqual(lazy_data[v], lazy_stamps[v], f"lazy: KV[v={v}] != stamp")
def test_lazy_hole_set_directional_pop(self):
"""The _HoleSet pops smallest-first for grow-up; alloc must drain
the deepest hole first (the greedy clustering rule keeps near-
boundary holes available for cheap absorption by compaction).
"""
_pool, fa, _kv = self._make_full(lazy=True)
a = fa.alloc(6)
# Free middle and lower middles so the holes are NOT at boundary.
fa.free(a[1:2].clone()) # frees physical at index v2p[a[1]]
fa.free(a[3:4].clone())
# Capture which physical pages are now in the hole set.
# `_free_phys_pages` is a torch.Tensor; `.tolist()` returns
# Python ints so `sorted` produces ints (not 0-dim tensors).
holes_before = sorted(fa._free_phys_pages.tolist())
self.assertEqual(len(holes_before), 2)
# Alloc 1 — should drain a hole (grow-up).
# With sort-after-merge OFF (default), the drain order
# is FIFO over the free-list tensor — NOT "smallest first".
# We only assert that the bound physical is ONE OF the holes.
a2 = fa.alloc(1)
bound_phys = int(fa.virtual_to_physical[int(a2.item())].item())
self.assertIn(bound_phys, holes_before)
def test_lazy_non_urgent_stops_at_write_set_blocker(self):
"""Write-race case: when the topmost survivor IS in an
in-flight batch's write-set, non-urgent `_flush` STOPS the
@@ -2358,57 +2194,6 @@ class TestLazyCompaction(unittest.TestCase):
self.assertEqual(len(fa._pending_reuse), 0)
self.assertEqual(len(fa._pending_reuse_pages_cpu), 0)
def test_lazy_pending_reuse_urgent_wait(self):
"""Under urgent drain, an unfired event triggers wait_event; we
simulate this by checking that the drain ALSO releases unfired
entries (with a fake event whose `query` is False — `wait_event` is
a no-op in CPU mode since there's no current stream's wait_event for
a FakeEvent, so we test the release path)."""
_pool, fa, _kv = self._make_full(lazy=True)
a = fa.alloc(4)
class _FakeEvent:
def __init__(self):
self.waited = False
def query(self):
return False # never fires
# Inject ONE batch entry into _pending_reuse keyed by
# Event. Value is `(cpu_list, gpu_tensor)`. The parallel CPU
# set must also be updated.
# (Simulates a prior compaction whose event hasn't fired.)
p = int(fa.virtual_to_physical[int(a[2].item())].item())
# Clear v2p/p2v so post-drain reuse is safe.
fa.virtual_to_physical[int(a[2].item())] = -1
fa.physical_to_virtual[p] = -1
ev = _FakeEvent()
gpu_t = torch.tensor([p], dtype=torch.int64, device=fa.device)
fa._pending_reuse[ev] = ([p], gpu_t)
fa._pending_reuse_pages_cpu.add(p)
# Urgent drain — should release p despite event.query()=False.
# (CPU shim: torch.cuda.current_stream() may not exist; wrap try.)
try:
fa._drain_pending_reuse(urgent=True)
except Exception:
# CPU: wait_event may not work; this test is GPU-only.
self.skipTest("wait_event requires CUDA")
self.assertEqual(len(fa._pending_reuse), 0)
self.assertEqual(len(fa._pending_reuse_pages_cpu), 0)
def test_lazy_flush_opportunistic_hook(self):
"""The public flush_opportunistic method runs the non-urgent path
and is safe to call when no holes exist."""
_pool, fa, _kv = self._make_full(lazy=True)
# No holes → returns 0 moves, no-op.
self.assertEqual(fa.flush_opportunistic(), 0)
# Create a hole then call flush_opportunistic; latest_event=None
# means src releases immediately.
a = fa.alloc(3)
fa.free(a[0:1].clone())
moves = fa.flush_opportunistic()
self.assertGreaterEqual(moves, 1)
class TestO3FusedAllocBind(unittest.TestCase):
"""Fused take_physical_pages + bind_pages.
@@ -2465,17 +2250,6 @@ class TestO3FusedAllocBind(unittest.TestCase):
ma.bind_peer(fa)
return pool, fa, full_kv
def test_helper_exists_and_returns_tensor(self):
"""The helper `_alloc_bind_fast_or_slow` is wired and returns a
tensor on success."""
_pool, fa, _kv = self._make_full(lazy=True)
v_pages = torch.tensor([10, 11, 12], dtype=torch.int64, device="cuda")
phys = fa._alloc_bind_fast_or_slow(v_pages, 3)
self.assertIsNotNone(phys)
self.assertEqual(phys.shape, (3,))
self.assertEqual(phys.dtype, torch.int64)
self.assertEqual(phys.device.type, "cuda")
def test_fast_path_when_no_holes(self):
"""When `_free_phys_pages` is empty, the fast path fires.
Verifies: watermark advanced, v2p and p2v scattered correctly,
test/registered/unit/mem_cache/test_radix_cache_unit.py
-102
View File
@@ -17,7 +17,6 @@ Usage:
python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic
"""
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
# CPU-based unit test, runs quickly on any GPU runner
@@ -25,7 +24,6 @@ register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=5, suite="stage-b-test-1-gpu-small-amd")
import random
import time
import unittest
import unittest.mock
from array import array
@@ -49,13 +47,6 @@ DEFAULT_PAGE_SIZE = 4
class TestRadixKey(unittest.TestCase):
"""Test cases for RadixKey class."""
def test_init_basic(self):
"""Test basic initialization of RadixKey."""
token_ids = [1, 2, 3, 4]
key = RadixKey(array("q", token_ids))
self.assertEqual(list(key.token_ids), token_ids)
self.assertIsNone(key.extra_key)
def test_init_with_extra_key(self):
"""Test initialization with extra_key."""
token_ids = [1, 2, 3]
@@ -64,20 +55,6 @@ class TestRadixKey(unittest.TestCase):
self.assertEqual(list(key.token_ids), token_ids)
self.assertEqual(key.extra_key, extra_key)
def test_len(self):
"""Test __len__ method."""
key = RadixKey(array("q", [1, 2, 3]))
self.assertEqual(len(key), 3)
empty_key = RadixKey(array("q", []))
self.assertEqual(len(empty_key), 0)
def test_iter(self):
"""Test __iter__ method."""
token_ids = [1, 2, 3, 4]
key = RadixKey(array("q", token_ids))
self.assertEqual(list(key), token_ids)
def test_len_and_iter(self):
"""Test __len__ and __iter__ methods."""
test_cases = [
@@ -127,21 +104,6 @@ class TestRadixKey(unittest.TestCase):
with self.assertRaises(IndexError):
_ = key[10] # Out of bounds
def test_repr(self):
"""Test __repr__ method."""
key = RadixKey(array("q", [1, 2, 3]), "test")
repr_str = repr(key)
self.assertIn("RadixKey", repr_str)
self.assertIn("extra_key='test'", repr_str)
self.assertIn("[1, 2, 3]", repr_str)
def test_repr_long_token_ids(self):
"""Test __repr__ with long token_ids."""
long_tokens = list(range(15))
key = RadixKey(array("q", long_tokens))
repr_str = repr(key)
self.assertIn("...", repr_str) # Should be truncated
def _assert_match(self, a, b, page_size, expected, is_bigram=False):
key_a = RadixKey(array("q", a), is_bigram=is_bigram)
key_b = RadixKey(array("q", b), is_bigram=is_bigram)
@@ -225,13 +187,6 @@ class TestTreeNode(unittest.TestCase):
node2 = TreeNode()
self.assertEqual(node2.id, 1) # Counter was incremented
def test_counter_increment(self):
"""Test that counter increments properly."""
node1 = TreeNode()
node2 = TreeNode()
self.assertEqual(node1.id, 0)
self.assertEqual(node2.id, 1)
def test_evicted_backuped_properties(self):
"""Test evicted and backuped properties."""
test_cases = [
@@ -313,15 +268,6 @@ class TestTreeNode(unittest.TestCase):
n4.hash_value = ["h4"]
self.assertEqual(n4.get_prefix_hash_values(n3), ["h1", "h2", "h3"])
def test_lt_comparison(self):
"""Test less than comparison based on last_access_time."""
node1 = TreeNode()
time.sleep(0.001) # Small delay to ensure different timestamps
node2 = TreeNode()
self.assertTrue(node1 < node2)
self.assertFalse(node2 < node1)
class TestRadixCache(unittest.TestCase):
"""Test cases for RadixCache class."""
@@ -677,46 +623,6 @@ class TestRadixCache(unittest.TestCase):
match_len = len(result.device_indices)
self.assertEqual(match_len % page_size, 0)
def test_pretty_print_basic(self):
"""Test pretty_print produces output."""
cache = RadixCache.create_simulated()
cache.insert(
InsertParams(
key=RadixKey(array("q", [1, 2, 3])),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
# Just test that it doesn't crash
try:
cache.pretty_print()
except Exception as e:
self.fail(f"pretty_print raised an exception: {e}")
def test_all_values_flatten(self):
"""Test all_values_flatten method."""
cache = RadixCache.create_simulated()
cache.insert(
InsertParams(
key=RadixKey(array("q", [1, 2])),
value=torch.tensor([10, 20], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey(array("q", [3, 4])),
value=torch.tensor([30, 40], dtype=torch.int64),
)
)
all_values = cache.all_values_flatten()
self.assertEqual(len(all_values), 4)
# Values should contain all inserted values (order may vary)
values_set = set(all_values.tolist())
self.assertEqual(values_set, {10, 20, 30, 40})
def test_advanced_prefix_match_with_node_splits(self):
"""Advanced prefix matching: splits inside nodes and across pages."""
for page_size in [1, 2]:
@@ -895,14 +801,6 @@ class TestRadixCache(unittest.TestCase):
# The cache size should be within reasonable bounds of the actual allocated memory.
self.assertLess(torch_allocated, cache_size_bytes * 2)
def test_available_and_evictable_str(self):
mock_allocator = unittest.mock.Mock()
mock_allocator.available_size.return_value = 10
cache: RadixCache = RadixCache.create_simulated(mock_allocator=mock_allocator)
print(cache.available_and_evictable_str())
print(available_and_evictable_str(cache))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
-21
View File
@@ -740,27 +740,6 @@ _CI_BENCH_CONFIGS = [
num_seqs=5000,
kv_size=500_000,
),
dict(
label="FULL_SWA_ps1",
components=(ComponentType.FULL, ComponentType.SWA),
page_size=1,
num_seqs=1000,
kv_size=100_000,
),
dict(
label="FULL_ps16",
components=(ComponentType.FULL,),
page_size=16,
num_seqs=1000,
kv_size=100_000,
),
dict(
label="FULL_SWA_ps16",
components=(ComponentType.FULL, ComponentType.SWA),
page_size=16,
num_seqs=1000,
kv_size=100_000,
),
dict(
label="FULL_ps128",
components=(ComponentType.FULL,),
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
-361
View File
@@ -26,7 +26,6 @@ from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
EvictResult,
IncLockRefResult,
InitLoadBackParams,
InsertParams,
@@ -738,22 +737,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(len(m.device_indices), len(base))
cache.sanity_check()
def test_evict_basic(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_a = self._make_seq(1, 2)
seq_b = self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_a)
self._insert(cache, allocator, req_to_token_pool, seq_b)
total = len(seq_a) + len(seq_b)
self.assertEqual(cache.full_evictable_size(), total)
result = cache.evict(EvictParams(num_tokens=len(seq_a)))
self.assertIsInstance(result, EvictResult)
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a))
self.assertTrue(cache.full_evictable_size() <= len(seq_b))
cache.sanity_check()
def test_evict_respects_lock_ref(self):
"""Lock protects from eviction; unlock allows re-eviction."""
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
@@ -793,24 +776,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result.mamba_num_evicted, 0)
cache.sanity_check()
def test_evict_until_empty(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seqs = [self._make_seq(i * 100, 2) for i in range(5)]
for s in seqs:
self._insert(cache, allocator, req_to_token_pool, s)
total = sum(len(s) for s in seqs)
self.assertEqual(cache.full_evictable_size(), total)
result = cache.evict(EvictParams(num_tokens=total * 2))
self.assertGreaterEqual(result.num_tokens_evicted, total)
self.assertEqual(cache.full_evictable_size(), 0)
if self.cfg.has_mamba:
self.assertEqual(cache.mamba_evictable_size(), 0)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[0]))))
self.assertEqual(len(m.device_indices), 0)
cache.sanity_check()
def test_prev_prefix_len(self):
"""Three-step test: free overlap, free partial, no free."""
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
@@ -860,28 +825,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(allocator.available_size(), avail_before - len(seq_3p))
cache.sanity_check()
def test_node_split_at_boundary(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
base = self._make_seq(1, 3)
self._insert(cache, allocator, req_to_token_pool, base)
fork_a = base + self._make_seq(100, 1)
fork_b = base + self._make_seq(200, 1)
self._insert(cache, allocator, req_to_token_pool, fork_a)
result = self._insert(cache, allocator, req_to_token_pool, fork_b)
self.assertEqual(result.prefix_len, len(base))
for seq in (fork_a, fork_b):
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
m = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
)
self.assertEqual(len(m.device_indices), len(base))
cache.sanity_check()
def test_cache_finished_req_insert(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
ps = self.cfg.page_size
@@ -1086,34 +1029,6 @@ class UnifiedRadixCacheSuite:
cache.pretty_print()
cache.sanity_check()
def test_multi_branch_tree(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
base = self._make_seq(1, 2)
self._insert(cache, allocator, req_to_token_pool, base)
for suffix_start in [100, 200, 300]:
seq = base + self._make_seq(suffix_start, 2)
self._insert(cache, allocator, req_to_token_pool, seq)
for suffix_start in [100, 200, 300]:
seq = base + self._make_seq(suffix_start, 2)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
m = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", base + self._make_seq(999, 1))))
)
self.assertEqual(len(m.device_indices), len(base))
cache.sanity_check()
def test_paged_child_key_is_tuple(self):
if self.cfg.page_size == 1:
self.skipTest("page_size > 1 only")
cache, _, _ = build_fixture(self.cfg)
key = RadixKey(array("q", self._make_seq(1, 1)))
child_key = key.child_key(cache.page_size)
self.assertIsInstance(child_key, tuple)
def test_paged_match_truncates_unaligned_key(self):
"""match_prefix internally aligns keys to page boundary."""
if self.cfg.page_size == 1:
@@ -1227,18 +1142,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(len(m.device_indices), 0)
cache.sanity_check()
def test_mamba_evict_result_accounting(self):
if not self.cfg.has_mamba:
self.skipTest("requires Mamba component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 3)
self._insert(cache, allocator, req_to_token_pool, seq)
result = cache.evict(EvictParams(num_tokens=len(seq)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
self.assertGreaterEqual(result.mamba_num_evicted, 1)
cache.sanity_check()
def test_mamba_evict_cascades_on_full_leaf(self):
if not self.cfg.has_mamba:
self.skipTest("requires Mamba component")
@@ -1276,17 +1179,6 @@ class UnifiedRadixCacheSuite:
)
cache.sanity_check()
def test_swa_insert_and_match(self):
if not self.cfg.has_swa:
self.skipTest("requires SWA component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 3)
self._insert(cache, allocator, req_to_token_pool, seq)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), len(seq))
cache.sanity_check()
def test_swa_unfinished_recovery_preserves_locked_full_value(self):
if not self.cfg.has_swa or self.cfg.has_mamba:
self.skipTest("requires SWA without Mamba")
@@ -1391,34 +1283,6 @@ class UnifiedRadixCacheSuite:
self.assertIsNone(node.component_data[ComponentType.SWA].value)
cache.sanity_check()
def test_swa_evict_cascades(self):
"""Evict SWA tokens via swa_num_tokens — cascades to lower-priority components."""
if not self.cfg.has_swa:
self.skipTest("requires SWA component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_short = self._make_seq(1, 2)
seq_long = seq_short + self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_short)
self._insert(cache, allocator, req_to_token_pool, seq_long)
result = cache.evict(EvictParams(num_tokens=0, swa_num_tokens=len(seq_short)))
self.assertGreater(result.swa_num_tokens_evicted, 0)
cache.sanity_check()
def test_swa_evict_cascades_mamba(self):
"""SWA eviction on an internal node cascades to Mamba."""
if not self.cfg.has_swa or not self.cfg.has_mamba:
self.skipTest("requires SWA and Mamba components")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_short = self._make_seq(1, 3)
seq_long = seq_short + self._make_seq(500, 4)
self._insert(cache, allocator, req_to_token_pool, seq_short)
self._insert(cache, allocator, req_to_token_pool, seq_long)
result = cache.evict(EvictParams(num_tokens=0, swa_num_tokens=len(seq_short)))
self.assertGreaterEqual(result.swa_num_tokens_evicted, 0)
cache.sanity_check()
def test_leaf_transition_swa_evict_spares_locked_full(self):
if not self.cfg.has_swa or not self.cfg.has_mamba:
self.skipTest("requires SWA and Mamba components")
@@ -1678,46 +1542,6 @@ class UnifiedRadixCacheSuite:
self.assertTrue(cache._is_device_leaf(node_a))
cache.sanity_check()
def test_swa_evict_full_leaf_cascades_all(self):
if not self.cfg.has_swa:
self.skipTest("requires SWA component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_a = self._make_seq(1, 2)
seq_b = self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_a)
self._insert(cache, allocator, req_to_token_pool, seq_b)
result = cache.evict(EvictParams(num_tokens=len(seq_a)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a))
self.assertGreater(result.swa_num_tokens_evicted, 0)
if self.cfg.has_mamba:
self.assertGreaterEqual(result.mamba_num_evicted, 1)
cache.sanity_check()
def test_swa_lock_protects_from_eviction(self):
if not self.cfg.has_swa:
self.skipTest("requires SWA component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_a = self._make_seq(1, 2)
seq_b = self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_a)
self._insert(cache, allocator, req_to_token_pool, seq_b)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lock_result = cache.inc_lock_ref(m.last_device_node)
result = cache.evict(EvictParams(num_tokens=len(seq_a) + len(seq_b)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_b))
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m.device_indices), len(seq_a))
cache.dec_lock_ref(
m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock),
)
cache.sanity_check()
def test_swa_leaf_capped_to_window_on_insert(self):
"""A long SWA leaf is split so locking it protects one window of SWA
while full attention still protects the whole sequence."""
@@ -1909,44 +1733,6 @@ class UnifiedRadixCacheSuite:
)
cache.sanity_check()
def test_swa_lru_cushion_bound_is_sliding_window_plus_page_size(self):
if not self._swa_pinning_cfg_supported():
self.skipTest("requires SWA-only config with node size >= cushion")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_a = self._make_seq(1, 8)
seq_ab = seq_a + self._make_seq(100, 8)
seq_abc = seq_ab + self._make_seq(200, 8)
self._insert(cache, allocator, req_to_token_pool, seq_a)
self._insert(cache, allocator, req_to_token_pool, seq_ab)
self._insert(cache, allocator, req_to_token_pool, seq_abc)
seq_side = self._make_seq(900, 5)
self._insert(cache, allocator, req_to_token_pool, seq_side)
pre = self._swa_lru_order(cache)
self.assertEqual(len(pre), 8)
side_node, c_node, b_node, a_node = pre[0], pre[2], pre[4], pre[6]
c_prefix = pre[3] # C's prefix pairs with its tail (c_node) at pre[2:4]
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_abc))))
self.assertEqual(len(m.device_indices), len(seq_abc))
post = self._swa_lru_order(cache)
cushion = self.cfg.sliding_window_size + self.cfg.page_size
# Under leaf-cap no single node exceeds the cushion; it spans C's capped
# tail plus its prefix, so both of C's nodes are refreshed to the MRU
# side while B and A keep their relative order below.
self.assertLess(len(c_node.key), cushion)
self.assertIn(c_node, post[:2])
self.assertIn(c_prefix, post[:2])
side_pos = post.index(side_node)
b_pos = post.index(b_node)
a_pos = post.index(a_node)
self.assertLess(side_pos, b_pos, "B was below side in pre, must stay below")
self.assertLess(b_pos, a_pos, "A was below B in pre, must stay below")
cache.sanity_check()
def test_swa_eager_eviction_on_unfinished_req(self):
if not self.cfg.has_swa or self.cfg.has_mamba:
self.skipTest(
@@ -2362,51 +2148,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result.num_tokens_evicted, before - after)
cache.sanity_check()
def test_evict_locked_subtree_skipped(self):
"""All nodes in a locked path are skipped during eviction."""
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq_a = self._make_seq(1, 3)
seq_b = self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_a)
self._insert(cache, allocator, req_to_token_pool, seq_b)
# Lock seq_a
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
lr = cache.inc_lock_ref(m.last_device_node)
# Try to evict everything
total = cache.full_evictable_size() + cache.full_protected_size()
result = cache.evict(EvictParams(num_tokens=total))
# seq_a should still be matchable (protected)
m2 = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_a))))
self.assertEqual(len(m2.device_indices), len(seq_a))
cache.dec_lock_ref(
m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)),
)
cache.sanity_check()
def test_mamba_internal_tombstone_evict(self):
"""Mamba eviction on internal node tombstones mamba only, keeps Full."""
if not self.cfg.has_mamba:
self.skipTest("requires Mamba component")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
# Create internal node with mamba and leaf extending it
seq_short = self._make_seq(1, 2)
seq_long = seq_short + self._make_seq(500, 2)
self._insert(cache, allocator, req_to_token_pool, seq_short)
self._insert(cache, allocator, req_to_token_pool, seq_long)
# Evict only mamba
result = cache.evict(EvictParams(num_tokens=0, mamba_num=10))
self.assertEqual(cache.mamba_evictable_size(), 0)
# Full should still be accessible for at least the long seq base
# (mamba gone breaks match, but full data might still be in tree)
cache.sanity_check()
def test_evict_reinsert_after_full_eviction(self):
"""After evicting everything, new inserts work correctly."""
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
@@ -2990,33 +2731,6 @@ class UnifiedRadixCacheSuite:
cache.sanity_check()
def test_hicache_node_states(self):
"""Verify device-only to device+host transition after real backup."""
if self._skip_unsupported_hicache_test():
return
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
seq = self._make_seq(1, 2)
self._insert(cache, allocator, req_to_token_pool, seq)
# Find the leaf node
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self.assertIsNot(node, cache.root_node)
ct = ComponentType.FULL
# S1: device only
self.assertIsNotNone(node.component_data[ct].value)
self.assertIsNone(node.component_data[ct].host_value)
self.assertFalse(node.backuped)
self.assertFalse(node.evicted)
self._backup_node(cache, node)
self.assertIsNotNone(node.component_data[ct].value)
self.assertIsNotNone(node.component_data[ct].host_value)
self.assertTrue(node.backuped)
self.assertFalse(node.evicted)
cache.sanity_check()
def test_hicache_evict_to_host(self):
"""Evicting a backed-up device leaf demotes it to host-only state."""
if self._skip_unsupported_hicache_test():
@@ -3130,27 +2844,6 @@ class UnifiedRadixCacheSuite:
self.assertTrue(cur.evicted and cur.backuped)
cache.sanity_check()
def test_hicache_d_leaf_h_leaf_mutual_exclusion(self):
"""D-leaf and H-leaf sets are always disjoint."""
if self._skip_unsupported_hicache_test():
return
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
seqs = [self._make_seq(i * 100, 2) for i in range(4)]
for s in seqs:
self._insert(cache, allocator, req_to_token_pool, s)
for i in range(2):
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seqs[i]))))
self._backup_node(cache, m.last_device_node)
# Evict one backed-up node
cache.evict(EvictParams(num_tokens=len(seqs[0])))
# Check mutual exclusion
overlap = cache.evictable_device_leaves & cache.evictable_host_leaves
self.assertEqual(len(overlap), 0)
cache.sanity_check()
def test_hicache_host_leaf_eviction(self):
"""Evicting a host leaf removes the node from the tree entirely."""
if self._skip_unsupported_hicache_test():
@@ -3713,44 +3406,6 @@ class UnifiedRadixCacheSuite:
self.assertGreaterEqual(int(xfer.host_indices.numel()), sw)
self.assertEqual(xfer.nodes_to_load, chain[-expected_pages:])
def test_hicache_swa_host_independent_of_full(self):
"""FULL host and SWA host are physically independent.
Freeing one component's host_value must not touch the other.
"""
if not self.cfg.has_swa:
self.skipTest("requires SWA")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2)
self._insert(cache, allocator, req_to_token_pool, seq)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self._simulate_backup(cache, node)
cache.evict(EvictParams(num_tokens=len(seq)))
cd_full = node.component_data[ComponentType.FULL]
cd_swa = node.component_data[ComponentType.SWA]
self.assertIsNotNone(cd_full.host_value)
self.assertIsNotNone(cd_swa.host_value)
self.assertIn(node, cache.evictable_host_leaves)
self.assertTrue(cache.host_lru_lists[ComponentType.SWA].in_list(node))
# Drop FULL host bookkeeping. SWA side must stay intact.
cache.evictable_host_leaves.discard(node)
cd_full.host_value = None
self.assertIsNotNone(cd_swa.host_value)
self.assertTrue(cache.host_lru_lists[ComponentType.SWA].in_list(node))
self.assertNotIn(node, cache.evictable_host_leaves)
# Drop SWA host bookkeeping. FULL side (already cleared) stays cleared.
cache.host_lru_lists[ComponentType.SWA].remove_node(node)
cd_swa.host_value = None
self.assertIsNone(cd_full.host_value)
self.assertIsNone(cd_swa.host_value)
self.assertFalse(cache.host_lru_lists[ComponentType.SWA].in_list(node))
self.assertNotIn(node, cache.evictable_host_leaves)
def _swa_finalize_setup(self):
"""Build a SWA chain long enough to fill at least the window
plus one extra page, and host-back every node so we can flip
@@ -4325,22 +3980,6 @@ class UnifiedLRUListBoundedRefreshTest(CustomTestCase):
cur = cur.lru_next[pt]
return out
def test_bounded_refresh_stops_after_accumulated_meets_window(self):
root, [a, b, c, d] = self._build_chain([2, 2, 2, 2])
lru = UnifiedLRUList(ComponentType.SWA, self.components)
for n in (a, b, c, d):
lru.insert_mru(n)
self.assertEqual(self._lru_order(lru), [d, c, b, a])
# window=5, page_size=1 implicit; nodes are size 2 each
# Walking up from D: visit D(acc=2<5) -> visit C(acc=4<5) -> visit
# B(acc=6>=5, refresh and stop). A is NOT touched.
lru.reset_node_and_window_ancestors_mru(
d, root, window_size=5, should_include=lambda _n: True
)
# Expected MRU->LRU: D, C, B (refreshed in walk-up order), A (untouched)
self.assertEqual(self._lru_order(lru), [d, c, b, a])
def test_bounded_refresh_skips_non_included(self):
root, [a, b, c, d] = self._build_chain([2, 2, 2, 2])
lru = UnifiedLRUList(ComponentType.SWA, self.components)
test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py
+29 -163
View File
@@ -76,16 +76,6 @@ class TestGraphSlot(unittest.TestCase):
axis="garbage",
)
def test_slice_for_before_buffer_alloc_raises(self):
slot = GraphSlot(
name="x",
shape_fn=lambda bs, mt: (bs,),
dtype=torch.int32,
axis="bs",
)
with self.assertRaises(RuntimeError):
slot.slice_for(padded_bs=1, padded_num_tokens=1)
class TestRegistryRegister(unittest.TestCase):
def test_register_allocates_zero_buffer(self):
@@ -102,22 +92,6 @@ class TestRegistryRegister(unittest.TestCase):
self.assertEqual(slot.buffer.dtype, torch.int64)
self.assertTrue(torch.equal(slot.buffer, torch.zeros(16, dtype=torch.int64)))
def test_register_fill_sentinel_init(self):
r = _make_registry()
slot = r.register_slot(
GraphSlot(
name="seq_lens",
shape_fn=lambda bs, mt: (bs,),
dtype=torch.int32,
axis="bs",
padding_policy=PaddingPolicy.FILL_SENTINEL,
pad_value=7,
)
)
self.assertTrue(
torch.equal(slot.buffer, torch.full((8,), 7, dtype=torch.int32))
)
def test_register_duplicate_raises(self):
r = _make_registry()
r.register_slot(
@@ -153,22 +127,6 @@ class TestRegistryRegister(unittest.TestCase):
self.assertFalse(r.has_slot("off"))
self.assertNotIn("off", r.slot_names())
def test_cpu_device_override(self):
r = _make_registry()
slot = r.register_slot(
GraphSlot(
name="seq_lens_cpu",
shape_fn=lambda bs, mt: (bs,),
dtype=torch.int32,
axis="bs",
device=torch.device("cpu"),
padding_policy=PaddingPolicy.FILL_SENTINEL,
pad_value=11,
)
)
self.assertEqual(slot.buffer.device.type, "cpu")
self.assertEqual(int(slot.buffer[0].item()), 11)
class TestFillFromAndExtract(unittest.TestCase):
"""End-to-end exercise: register a representative slot set, fill from
@@ -233,35 +191,6 @@ class TestFillFromAndExtract(unittest.TestCase):
)
return r
def test_basic_fill_no_padding(self):
r = self._build_registry()
fb = _MiniForwardBatch(
batch_size=4,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([3, 1, 4, 2], dtype=torch.int64),
seq_lens=torch.tensor([10, 11, 12, 13], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64) + 100,
positions=torch.arange(8, dtype=torch.int64),
seq_lens_cpu=torch.tensor([10, 11, 12, 13], dtype=torch.int32),
)
r.fill_from(
fb,
raw_bs=4,
padded_bs=4,
raw_num_tokens=8,
padded_num_tokens=8,
)
self.assertTrue(torch.equal(r.get_slot("input_ids").buffer, fb.input_ids))
self.assertTrue(
torch.equal(r.get_slot("req_pool_indices").buffer, fb.req_pool_indices)
)
self.assertTrue(torch.equal(r.get_slot("seq_lens").buffer, fb.seq_lens))
self.assertTrue(
torch.equal(r.get_slot("out_cache_loc").buffer, fb.out_cache_loc)
)
self.assertTrue(torch.equal(r.get_slot("positions").buffer, fb.positions))
self.assertTrue(torch.equal(r.get_slot("seq_lens_cpu").buffer, fb.seq_lens_cpu))
def test_fill_with_padding_resets_zero_and_sentinel(self):
r = self._build_registry()
# Pre-poison the padded tail to a non-zero value so we can prove
@@ -404,39 +333,6 @@ class TestFillFromAndExtract(unittest.TestCase):
class TestMissingAndOptionalSlots(unittest.TestCase):
def test_missing_fb_attr_is_skipped(self):
r = _make_registry()
r.register_slot(
GraphSlot(
name="encoder_lens",
shape_fn=lambda bs, mt: (bs,),
dtype=torch.int32,
axis="bs",
padding_policy=PaddingPolicy.FILL_SENTINEL,
pad_value=0,
)
)
fb = _MiniForwardBatch(
batch_size=2,
input_ids=torch.arange(4, dtype=torch.int64),
encoder_lens=None, # FB doesn't carry this for this request.
)
# Should NOT raise; encoder_lens buffer stays at the FILL_SENTINEL
# init value.
r.fill_from(
fb,
raw_bs=2,
padded_bs=4,
raw_num_tokens=4,
padded_num_tokens=8,
)
self.assertTrue(
torch.equal(
r.get_slot("encoder_lens").buffer,
torch.zeros(8, dtype=torch.int32),
)
)
def test_extract_carries_none_for_absent_plain_slot(self):
# A plain copy slot absent this iter (mrope on a non-multimodal batch)
# must be carried as None, not exposed as the stale/zero buffer.
@@ -468,6 +364,35 @@ class TestMissingAndOptionalSlots(unittest.TestCase):
)
self.assertIsNone(fb_view.mrope_positions)
def test_plain_slot_with_missing_fb_attr_keeps_sentinel(self):
# A plain copy slot whose FB field is None must be skipped, leaving its
# buffer at the FILL_SENTINEL init value rather than raising.
r = _make_registry()
slot = r.register_slot(
GraphSlot(
"encoder_lens",
lambda bs, mt: (bs,),
torch.int32,
axis="bs",
padding_policy=PaddingPolicy.FILL_SENTINEL,
pad_value=0,
)
)
fb = _MiniForwardBatch(
batch_size=2,
input_ids=torch.arange(4, dtype=torch.int64),
encoder_lens=None,
)
r.fill_from(
fb,
raw_bs=2,
padded_bs=4,
raw_num_tokens=4,
padded_num_tokens=8,
)
# Buffer untouched by the copy; stays at the sentinel pad value.
self.assertTrue(torch.equal(slot.buffer, torch.zeros_like(slot.buffer)))
def test_extract_exposes_computed_slot_even_when_fb_field_none(self):
# A computed slot (copy_from_fb=False) is always exposed, even when its
# FB field is None — the None-skip carry applies only to plain copies.
@@ -635,40 +560,6 @@ class TestSourceFnSlots(unittest.TestCase):
r.fill_from(fb, raw_bs=3, padded_bs=8, raw_num_tokens=3, padded_num_tokens=16)
self.assertTrue(torch.all(buf == 7)) # untouched
def test_side_input_source_via_fill_context(self):
r = _make_registry(max_bs=8, max_num_tokens=16)
r.register_slot(
GraphSlot(
name="pp_proxy_tensors.hidden_states",
shape_fn=lambda _bs, mt: (mt,),
dtype=torch.int32,
axis="none",
padding_policy=PaddingPolicy.KEEP_PAD,
source_fn=lambda fb, ctx: (
None
if ctx.pp_proxy_tensors is None
else ctx.pp_proxy_tensors.tensors["hidden_states"]
),
)
)
buf = r.get_slot("pp_proxy_tensors.hidden_states").buffer
buf.zero_()
fb = _MiniForwardBatch(batch_size=4)
pp = SimpleNamespace(
tensors={"hidden_states": torch.tensor([5, 6, 7, 8], dtype=torch.int32)}
)
r.fill_from(
fb,
raw_bs=4,
padded_bs=8,
raw_num_tokens=4,
padded_num_tokens=16,
pp_proxy_tensors=pp,
)
self.assertTrue(
torch.equal(buf[:4], torch.tensor([5, 6, 7, 8], dtype=torch.int32))
)
def test_extract_buffer_skips_dotted_slots(self):
r = _make_registry(max_bs=8, max_num_tokens=16)
r.register_slot(
@@ -725,31 +616,6 @@ class TestPoolBackedAlloc(unittest.TestCase):
r2.get_slot("ids").buffer.data_ptr(),
)
def test_same_size_shares_one_allocation(self):
a = self._reg(max_num_tokens=16, share_pool=True)
b = self._reg(max_num_tokens=16, share_pool=True)
a.register_slot(self._ids_slot("ids"))
b.register_slot(self._ids_slot("ids"))
# Identical (name, size, dtype, device) -> one shared allocation.
self.assertEqual(
a.get_slot("ids").buffer.data_ptr(),
b.get_slot("ids").buffer.data_ptr(),
)
def test_different_sizes_do_not_share(self):
big = self._reg(max_num_tokens=32, share_pool=True)
small = self._reg(max_num_tokens=16, share_pool=True)
big.register_slot(self._ids_slot("ids"))
small.register_slot(self._ids_slot("ids"))
# Different sizes -> different pool keys -> independent storage (no
# aliasing a smaller request onto a larger buffer).
self.assertEqual(tuple(small.get_slot("ids").buffer.shape), (16,))
self.assertEqual(tuple(big.get_slot("ids").buffer.shape), (32,))
self.assertNotEqual(
small.get_slot("ids").buffer.data_ptr(),
big.get_slot("ids").buffer.data_ptr(),
)
def test_sharing_is_independent_of_registration_order(self):
from sglang.srt.model_executor import input_buffers
test/registered/unit/parser/test_conversation.py
-230
View File
@@ -119,19 +119,6 @@ class TestConversationGetPrompt(CustomTestCase):
self.assertIn("[USER]Hello\n", prompt)
self.assertTrue(prompt.endswith("[ASST]"))
def test_none_message_in_prompt(self):
"""Test that None message produces role-only output (no content)."""
conv = Conversation(
name="test",
system_message="",
roles=("User", "Assistant"),
messages=[["User", "Q"], ["Assistant", None]],
sep_style=SeparatorStyle.ADD_COLON_SINGLE,
sep="\n",
)
prompt = conv.get_prompt()
self.assertTrue(prompt.endswith("Assistant:"))
def test_empty_system_message(self):
"""Test that empty system message produces empty prefix for LLAMA3."""
conv = Conversation(
@@ -189,22 +176,6 @@ class TestConversationGetPrompt(CustomTestCase):
self.assertIn("[A]A<s2>", prompt)
self.assertTrue(prompt.endswith("[U]"))
def test_llama2_with_system(self):
"""Test LLAMA2 with system message."""
conv = Conversation(
name="test",
system_message="<<SYS>>\nBe helpful\n<</SYS>>\n\n",
system_template="[INST] {system_message}",
roles=("[INST]", "[/INST]"),
messages=[["[INST]", "Hi"], ["[/INST]", None]],
sep_style=SeparatorStyle.LLAMA2,
sep=" ",
sep2=" </s><s>",
)
prompt = conv.get_prompt()
self.assertIn("Be helpful", prompt)
self.assertIn("Hi ", prompt)
def test_llama2_without_system(self):
"""Test LLAMA2 without system message falls back to '[INST] ' prefix."""
conv = Conversation(
@@ -570,23 +541,6 @@ class TestConversationGetPrompt(CustomTestCase):
self.assertIn("USER: <image>Describe this\n", prompt)
self.assertIn("ASSISTANT: It shows a cat<eos>", prompt)
def test_mpt_with_tuple_message(self):
"""Test MPT style extracts first element from tuple messages."""
conv = Conversation(
name="test",
system_message="<|system|>",
roles=("<|user|>", "<|assistant|>"),
messages=[
["<|user|>", ("Hello", "extra1", "extra2")],
["<|assistant|>", None],
],
sep_style=SeparatorStyle.MPT,
sep="\n",
)
prompt = conv.get_prompt()
self.assertIn("<|user|>Hello\n", prompt)
self.assertNotIn("extra1", prompt)
def test_invalid_sep_style_raises(self):
"""Test that an invalid SeparatorStyle raises ValueError."""
conv = Conversation(
@@ -611,100 +565,6 @@ class TestConversationMethods(CustomTestCase):
sep="\n",
)
def test_append_message(self):
"""Test appending messages to conversation."""
conv = self._make_conv()
conv.append_message("User", "Hello")
conv.append_message("Assistant", "Hi")
self.assertEqual(len(conv.messages), 2)
self.assertEqual(conv.messages[0], ["User", "Hello"])
def test_set_system_message(self):
"""Test setting the system message."""
conv = self._make_conv()
conv.set_system_message("Be helpful")
self.assertEqual(conv.system_message, "Be helpful")
def test_update_last_message(self):
"""Test updating the last message in-place."""
conv = self._make_conv()
conv.append_message("User", "Q")
conv.append_message("Assistant", None)
conv.update_last_message("Answer")
self.assertEqual(conv.messages[-1][1], "Answer")
def test_to_openai_api_messages_with_system(self):
"""Test conversion to OpenAI format with system message."""
conv = self._make_conv()
conv.system_message = "Be helpful"
conv.append_message("User", "Hello")
conv.append_message("Assistant", "Hi")
result = conv.to_openai_api_messages()
self.assertEqual(result[0], {"role": "system", "content": "Be helpful"})
self.assertEqual(result[1], {"role": "user", "content": "Hello"})
self.assertEqual(result[2], {"role": "assistant", "content": "Hi"})
def test_to_openai_api_messages_without_system(self):
"""Test conversion to OpenAI format without system message."""
conv = self._make_conv()
conv.append_message("User", "Hello")
result = conv.to_openai_api_messages()
self.assertEqual(len(result), 1)
self.assertEqual(result[0]["role"], "user")
def test_to_openai_api_messages_skips_none_assistant(self):
"""Test that None assistant message is omitted from OpenAI format."""
conv = self._make_conv()
conv.append_message("User", "Hello")
conv.append_message("Assistant", None)
result = conv.to_openai_api_messages()
self.assertEqual(len(result), 1) # only user message
def test_to_gradio_chatbot(self):
"""Test conversion to Gradio chatbot format (user/assistant pairs)."""
conv = self._make_conv()
conv.append_message("User", "Q1")
conv.append_message("Assistant", "A1")
conv.append_message("User", "Q2")
conv.append_message("Assistant", "A2")
result = conv.to_gradio_chatbot()
self.assertEqual(len(result), 2)
self.assertEqual(result[0], ["Q1", "A1"])
self.assertEqual(result[1], ["Q2", "A2"])
def test_to_gradio_chatbot_pending_response(self):
"""Test Gradio format with pending assistant response (None)."""
conv = self._make_conv()
conv.append_message("User", "Q1")
conv.append_message("Assistant", None)
result = conv.to_gradio_chatbot()
self.assertEqual(result, [["Q1", None]])
def test_append_image(self):
"""Test appending image data to conversation."""
conv = self._make_conv()
conv.image_data = []
conv.append_image("http://example.com/img.jpg", "auto")
self.assertEqual(len(conv.image_data), 1)
self.assertEqual(conv.image_data[0].url, "http://example.com/img.jpg")
self.assertEqual(conv.image_data[0].detail, "auto")
def test_append_video(self):
"""Test appending video data to conversation."""
conv = self._make_conv()
conv.video_data = []
conv.append_video("http://example.com/vid.mp4")
self.assertEqual(len(conv.video_data), 1)
self.assertEqual(conv.video_data[0], "http://example.com/vid.mp4")
def test_append_audio(self):
"""Test appending audio data to conversation."""
conv = self._make_conv()
conv.audio_data = []
conv.append_audio("http://example.com/audio.wav")
self.assertEqual(len(conv.audio_data), 1)
self.assertEqual(conv.audio_data[0], "http://example.com/audio.wav")
def test_copy_is_independent(self):
"""Test that copy() creates an independent conversation."""
conv = self._make_conv()
@@ -714,15 +574,6 @@ class TestConversationMethods(CustomTestCase):
self.assertEqual(len(conv.messages), 1)
self.assertEqual(len(copied.messages), 2)
def test_dict_serialization(self):
"""Test dict() returns expected keys."""
conv = self._make_conv()
conv.append_message("User", "Hello")
d = conv.dict()
self.assertEqual(d["template_name"], "test")
self.assertIn("messages", d)
self.assertIn("roles", d)
class TestTemplateRegistry(CustomTestCase):
def test_builtin_templates_exist(self):
@@ -730,24 +581,6 @@ class TestTemplateRegistry(CustomTestCase):
self.assertTrue(chat_template_exists("chatml"))
self.assertTrue(chat_template_exists("llama-2"))
def test_unregistered_template_not_found(self):
"""Test that non-existent template returns False."""
self.assertFalse(chat_template_exists("_nonexistent_template_xyz"))
def test_register_and_lookup(self):
"""Test registering and looking up a custom template."""
t = Conversation(
name="_test_conv_template",
roles=("A", "B"),
messages=[],
sep_style=SeparatorStyle.ADD_COLON_SINGLE,
sep="\n",
)
register_conv_template(t)
self.assertTrue(chat_template_exists("_test_conv_template"))
# Cleanup
del chat_templates["_test_conv_template"]
def test_register_duplicate_raises(self):
"""Test that registering a duplicate name without override raises."""
with self.assertRaises(AssertionError):
@@ -841,32 +674,6 @@ class TestGenerateEmbeddingConvs(CustomTestCase):
self.assertIn("Hello world", convs[0].messages[0][1])
self.assertIsNone(convs[0].messages[1][1]) # assistant placeholder
def test_with_image(self):
"""Test generating embedding conversations with image."""
convs = generate_embedding_convs(
texts=["Describe"],
images=["http://example.com/img.jpg"],
videos=[None],
template_name="chatml",
)
self.assertEqual(len(convs), 1)
msg = convs[0].messages[0][1]
self.assertIn("<image>", msg)
self.assertIn("Describe", msg)
def test_with_video(self):
"""Test generating embedding conversations with video."""
convs = generate_embedding_convs(
texts=["Describe"],
images=[None],
videos=["http://example.com/vid.mp4"],
template_name="chatml",
)
self.assertEqual(len(convs), 1)
msg = convs[0].messages[0][1]
self.assertIn("<video>", msg)
self.assertIn("Describe", msg)
def test_with_image_and_video(self):
"""Test embedding conv with both image and video."""
convs = generate_embedding_convs(
@@ -892,24 +699,8 @@ class TestGenerateEmbeddingConvs(CustomTestCase):
# None text should not produce "None" string
self.assertNotIn("None", msg)
def test_multiple_items(self):
"""Test generating multiple embedding conversations."""
convs = generate_embedding_convs(
texts=["text1", "text2"],
images=[None, None],
videos=[None, None],
template_name="chatml",
)
self.assertEqual(len(convs), 2)
class TestGetFullMultimodalTextPrompt(CustomTestCase):
def test_adds_missing_image_tokens(self):
"""Test adding missing image tokens to prompt."""
result = _get_full_multimodal_text_prompt("<image>", 3, "Describe this.")
self.assertEqual(result.count("<image>"), 3)
self.assertIn("Describe this.", result)
def test_preserves_existing_tokens(self):
"""Test that existing tokens in prompt are preserved."""
result = _get_full_multimodal_text_prompt(
@@ -927,17 +718,6 @@ class TestGetFullMultimodalTextPrompt(CustomTestCase):
with self.assertRaises(ValueError):
_get_full_multimodal_text_prompt("<image>", 1, "<image> <image>")
def test_zero_count_with_no_tokens(self):
"""Test zero modality count with no tokens in prompt."""
result = _get_full_multimodal_text_prompt("<image>", 0, "Just text")
self.assertEqual(result, "Just text")
def test_video_tokens(self):
"""Test adding missing video tokens."""
result = _get_full_multimodal_text_prompt("<video>", 2, "Describe:")
self.assertEqual(result.count("<video>"), 2)
self.assertIn("Describe:", result)
def test_tokens_joined_with_newline(self):
"""Test that missing tokens are joined with newlines before prompt."""
result = _get_full_multimodal_text_prompt("<image>", 3, "text")
@@ -1071,16 +851,6 @@ class TestGenerateChatConv(CustomTestCase):
with self.assertRaises(ValueError):
generate_chat_conv(request, "chatml")
def test_string_messages_raises(self):
"""Test that passing messages as a raw string raises ValueError."""
request = self._make_request(
[ChatCompletionMessageUserParam(role="user", content="Hi")]
)
# Manually override messages to be a string to trigger validation
request.__dict__["messages"] = "not a list"
with self.assertRaises(ValueError):
generate_chat_conv(request, "chatml")
def test_user_message_with_image(self):
"""Test user message with image content part."""
request = self._make_request(
test/registered/unit/parser/test_harmony_parser.py
-332
View File
@@ -4,10 +4,8 @@ import unittest
from sglang.srt.parser.harmony_parser import (
CanonicalStrategy,
Event,
HarmonyParser,
TextStrategy,
Token,
iter_tokens,
prefix_hold,
)
@@ -18,36 +16,7 @@ register_cpu_ci(est_time=7, suite="base-a-test-cpu")
register_cpu_ci(est_time=7, suite="base-c-test-cpu")
class TestEvent(CustomTestCase):
def test_init(self):
"""Test Event dataclass initialization."""
event = Event("reasoning", "content")
self.assertEqual(event.event_type, "reasoning")
self.assertEqual(event.content, "content")
class TestToken(CustomTestCase):
def test_init(self):
"""Test Token dataclass initialization."""
token = Token("START", 0, 7)
self.assertEqual(token.type, "START")
self.assertEqual(token.start, 0)
self.assertEqual(token.end, 7)
class TestPrefixHold(CustomTestCase):
def test_empty_text(self):
"""Test prefix_hold with empty text."""
emit, hold = prefix_hold("", ["<|start|>"])
self.assertEqual(emit, "")
self.assertEqual(hold, "")
def test_no_matching_prefixes(self):
"""Test prefix_hold with no matching prefixes."""
emit, hold = prefix_hold("hello world", ["<|start|>", "<|end|>"])
self.assertEqual(emit, "hello world")
self.assertEqual(hold, "")
def test_partial_token_suffix(self):
"""Test prefix_hold with partial token at end."""
emit, hold = prefix_hold("hello <|ret", ["<|return|>"])
@@ -68,11 +37,6 @@ class TestPrefixHold(CustomTestCase):
class TestIterTokens(CustomTestCase):
def test_empty_text(self):
"""Test iter_tokens with empty text."""
tokens = list(iter_tokens(""))
self.assertEqual(tokens, [])
def test_plain_text(self):
"""Test iter_tokens with plain text."""
tokens = list(iter_tokens("hello world"))
@@ -81,14 +45,6 @@ class TestIterTokens(CustomTestCase):
self.assertEqual(tokens[0].start, 0)
self.assertEqual(tokens[0].end, 11)
def test_single_token(self):
"""Test iter_tokens with single structural token."""
tokens = list(iter_tokens("<|start|>"))
self.assertEqual(len(tokens), 1)
self.assertEqual(tokens[0].type, "START")
self.assertEqual(tokens[0].start, 0)
self.assertEqual(tokens[0].end, 9)
def test_mixed_content(self):
"""Test iter_tokens with mixed text and tokens."""
tokens = list(iter_tokens("text<|start|>more text"))
@@ -154,11 +110,6 @@ class TestCanonicalStrategy(CustomTestCase):
def setUp(self):
self.strategy = CanonicalStrategy()
def test_init(self):
"""Test CanonicalStrategy initialization."""
self.assertIn("<|start|>", self.strategy.guard_tokens)
self.assertIn("<|constrain|>", self.strategy.guard_tokens)
def test_extract_channel_type(self):
"""Test _extract_channel_type method."""
self.assertEqual(self.strategy._extract_channel_type("analysis"), "analysis")
@@ -170,72 +121,6 @@ class TestCanonicalStrategy(CustomTestCase):
self.assertEqual(self.strategy._extract_channel_type("ANALYSIS"), "analysis")
self.assertIsNone(self.strategy._extract_channel_type("unknown"))
def test_parse_single_analysis_block(self):
"""Test parsing single analysis block."""
text = "<|channel|>analysis<|message|>Let me think about this<|end|>"
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "reasoning")
self.assertEqual(events[0].content, "Let me think about this")
self.assertEqual(remaining, "")
def test_parse_single_commentary_block(self):
"""Test parsing single commentary block."""
text = "<|channel|>commentary<|message|>User-visible message<|end|>"
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "normal")
self.assertEqual(events[0].content, "User-visible message")
self.assertEqual(remaining, "")
def test_parse_single_final_block(self):
"""Test parsing single final block."""
text = "<|start|>assistant<|channel|>final<|message|>The answer is 42<|return|>"
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "normal")
self.assertEqual(events[0].content, "The answer is 42")
self.assertEqual(remaining, "")
def test_parse_tool_call_commentary(self):
"""Test parsing tool call on commentary channel."""
text = '<|channel|>commentary to=functions.get_weather<|message|>{"location": "SF"}<|call|>'
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "tool_call")
self.assertEqual(events[0].content, '{"location": "SF"}')
self.assertEqual(remaining, "")
def test_parse_tool_call_analysis(self):
"""Test parsing built-in tool call on analysis channel."""
text = '<|channel|>analysis to=browser.search<|message|>{"query": "SGLang"}<|call|>'
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "tool_call")
self.assertEqual(events[0].content, '{"query": "SGLang"}')
self.assertEqual(remaining, "")
def test_parse_complex_sequence(self):
"""Test parsing complex sequence with multiple blocks."""
text = (
"<|channel|>analysis<|message|>Need to use function get_weather.<|end|>"
"<|start|>assistant<|channel|>commentary to=functions.get_weather<|message|>"
'{"location":"San Francisco"}<|call|>'
)
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 2)
self.assertEqual(events[0].event_type, "reasoning")
self.assertEqual(events[0].content, "Need to use function get_weather.")
self.assertEqual(events[1].event_type, "tool_call")
self.assertEqual(events[1].content, '{"location":"San Francisco"}')
self.assertEqual(remaining, "")
def test_parse_with_interspersed_text(self):
"""Test parsing with plain text between blocks."""
text = (
@@ -298,42 +183,11 @@ class TestCanonicalStrategy(CustomTestCase):
self.assertEqual(events[0].content, "")
self.assertEqual(remaining, "")
def test_parse_commentary_filler_between_blocks(self):
"""Test that 'commentary' filler between <|call|> and <|channel|> is filtered out."""
# This pattern occurs when the model generates malformed output
text = (
'<|channel|>commentary to=functions.get_weather<|message|>{"location":"SF"}<|call|>'
"commentary" # This should be filtered out
'<|channel|>commentary to=functions.get_temp<|message|>{"location":"NYC"}<|call|>'
)
events, remaining = self.strategy.parse(text)
# Should have 2 tool calls, no "commentary" normal text
self.assertEqual(len(events), 2)
self.assertEqual(events[0].event_type, "tool_call")
self.assertEqual(events[0].content, '{"location":"SF"}')
self.assertEqual(events[1].event_type, "tool_call")
self.assertEqual(events[1].content, '{"location":"NYC"}')
self.assertEqual(remaining, "")
# Verify no "commentary" text was emitted as normal content
normal_events = [e for e in events if e.event_type == "normal"]
commentary_events = [
e for e in normal_events if "commentary" in e.content.lower()
]
self.assertEqual(
len(commentary_events), 0, "Commentary filler should be filtered out"
)
class TestTextStrategy(CustomTestCase):
def setUp(self):
self.strategy = TextStrategy()
def test_init(self):
"""Test TextStrategy initialization."""
self.assertIn("analysis_then_final", self.strategy.patterns)
def test_parse_analysis_then_final(self):
"""Test parsing analysis then final format."""
text = "analysis I need to think about this. assistantfinal The answer is 42."
@@ -387,16 +241,6 @@ class TestTextStrategy(CustomTestCase):
self.assertEqual(len(events), 0)
self.assertEqual(remaining, text) # Hold entire buffer
def test_parse_partial_analysis_streaming(self):
"""Test streaming partial analysis content."""
text = "analysis partial content"
events, remaining = self.strategy.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "reasoning")
self.assertEqual(events[0].content, " partial content") # Space preserved
self.assertEqual(remaining, "analysis") # Hold header
def test_parse_case_insensitive(self):
"""Test case insensitive parsing."""
text = "ANALYSIS reasoning ASSISTANTFINAL answer"
@@ -437,11 +281,6 @@ class TestHarmonyParser(CustomTestCase):
def setUp(self):
self.parser = HarmonyParser()
def test_init(self):
"""Test HarmonyParser initialization."""
self.assertIsNone(self.parser.strategy)
self.assertEqual(self.parser._buffer, "")
def test_strategy_selection_canonical(self):
"""Test automatic strategy selection for canonical format."""
events = self.parser.parse("<|channel|>analysis<|message|>test<|end|>")
@@ -470,56 +309,6 @@ class TestHarmonyParser(CustomTestCase):
self.assertIsInstance(self.parser.strategy, TextStrategy)
self.assertEqual(len(events2), 1)
def test_streaming_canonical_format(self):
"""Test streaming with canonical format."""
chunks = [
"<|channel|>analysis<|message|>",
"reasoning content",
"<|end|>",
"<|start|>assistant<|channel|>final<|message|>",
"final answer",
"<|return|>",
]
all_events = []
for chunk in chunks:
events = self.parser.parse(chunk)
all_events.extend(events)
# Verify we get both reasoning and normal events
reasoning_events = [e for e in all_events if e.event_type == "reasoning"]
self.assertGreater(len(reasoning_events), 0)
normal_events = [e for e in all_events if e.event_type == "normal"]
self.assertGreater(len(normal_events), 0)
# Verify content is eventually parsed correctly
combined_reasoning = "".join(e.content for e in reasoning_events)
combined_normal = "".join(
e.content
for e in normal_events
if e.content and "<|return|>" not in e.content
)
self.assertIn("reasoning content", combined_reasoning)
self.assertIn("final answer", combined_normal)
def test_streaming_text_format(self):
"""Test streaming with text format."""
chunks = ["analysis reasoning", " content assistantfinal", " the answer"]
all_events = []
for chunk in chunks:
events = self.parser.parse(chunk)
all_events.extend(events)
# Should have reasoning and normal events
reasoning_events = [e for e in all_events if e.event_type == "reasoning"]
normal_events = [e for e in all_events if e.event_type == "normal"]
self.assertGreater(len(reasoning_events), 0)
self.assertGreater(len(normal_events), 0)
def test_streaming_commentary_filler(self):
"""Test that 'commentary' filler is filtered in streaming case."""
# Test when commentary arrives as a separate chunk after <|call|>
@@ -679,35 +468,6 @@ class TestIntegrationScenarios(CustomTestCase):
self.assertEqual(events[0].event_type, "tool_call")
self.assertEqual(events[0].content, '{"query": "SGLang"}')
def test_tool_response_handling(self):
"""Test tool response message handling."""
parser = HarmonyParser()
text = '<|start|>functions.get_weather to=assistant<|channel|>commentary<|message|>{"sunny": true, "temperature": 20}<|end|>'
events = parser.parse(text)
self.assertEqual(len(events), 1)
self.assertEqual(events[0].event_type, "normal")
self.assertEqual(events[0].content, '{"sunny": true, "temperature": 20}')
def test_text_fallback_formats(self):
"""Test various text fallback formats."""
parser = HarmonyParser()
# Test analysis then final
events1 = parser.parse("analysis thinking assistantfinal answer")
self.assertEqual(len([e for e in events1 if e.event_type == "reasoning"]), 1)
self.assertEqual(len([e for e in events1 if e.event_type == "normal"]), 1)
# Reset parser for next test
parser = HarmonyParser()
# Test final only
events2 = parser.parse("assistantfinal direct answer")
self.assertEqual(len(events2), 1)
self.assertEqual(events2[0].event_type, "normal")
def test_streaming_property_canonical(self):
"""Test streaming property: chunked parsing produces same semantic content as one-shot parsing."""
full_text = (
@@ -819,12 +579,6 @@ class TestEdgeCases(CustomTestCase):
self.assertEqual(len(reasoning_events), 1)
self.assertGreater(len(normal_events), 0)
def test_empty_input(self):
"""Test handling of empty input."""
parser = HarmonyParser()
events = parser.parse("")
self.assertEqual(len(events), 0)
def test_whitespace_preservation(self):
"""Test that whitespace is preserved correctly."""
parser = HarmonyParser()
@@ -878,14 +632,6 @@ class TestEdgeCases(CustomTestCase):
class TestAdditionalEdgeCases(CustomTestCase):
"""Additional tests to cover remaining edge cases."""
def test_prefix_hold_with_empty_token_in_list(self):
"""Test that empty string token in the list is skipped."""
from sglang.srt.parser.harmony_parser import prefix_hold
emit, hold = prefix_hold("hello", ["", "world"])
self.assertEqual(emit, "hello")
self.assertEqual(hold, "")
def test_iter_tokens_unknown_token_no_closing(self):
"""Test iter_tokens with <| that has no closing |>."""
from sglang.srt.parser.harmony_parser import iter_tokens
@@ -894,72 +640,6 @@ class TestAdditionalEdgeCases(CustomTestCase):
# Should emit TEXT tokens for the content after <|
self.assertTrue(any(t.type == "TEXT" for t in tokens))
def test_canonical_commentary_filler_after_call(self):
"""Test that MESSAGE token after CALL is filtered as commentary filler."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
text = "<|start|><|channel|>analysis<|message|>thinking<|end|><|call|><|message|>noise<|return|><|channel|>final<|message|>answer<|end|>"
events, remainder = strategy.parse(text)
# The MESSAGE after CALL should be filtered, final answer should appear
answers = [e.content for e in events if e.event_type == "normal"]
self.assertTrue(any("answer" in a for a in answers))
def test_canonical_standalone_structural_token_filtered(self):
"""Test that standalone structural tokens like <|end|> in TEXT position are filtered."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
# A malformed sequence where an END token appears in an unexpected position
text = "<|start|><|channel|>analysis<|message|>content<|end|>"
events, remainder = strategy.parse(text)
# Should parse without error
self.assertTrue(len(events) >= 0)
def test_canonical_incomplete_block_returns_partial(self):
"""Test parsing an incomplete channel block (no END token)."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
text = "<|start|><|channel|>analysis<|message|>partial content"
events, remainder = strategy.parse(text)
# Incomplete block: should hold content as remainder or emit partial
reasoning_events = [e for e in events if e.event_type == "reasoning"]
# The partial content may be in events or remainder
total = "".join(e.content for e in reasoning_events) + remainder
self.assertIn("partial", total)
def test_text_strategy_commentary_channel(self):
"""Test TextStrategy parsing commentary channel."""
from sglang.srt.parser.harmony_parser import TextStrategy
strategy = TextStrategy()
text = "commentary: some discussion\nassistantfinal: the answer"
events, remainder = strategy.parse(text)
normal = [e for e in events if e.event_type == "normal"]
self.assertTrue(any("the answer" in e.content for e in normal))
def test_canonical_call_with_text_commentary_after(self):
"""Test filtering of 'commentary' text after CALL token."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
text = "<|start|><|channel|>analysis<|message|>think<|end|><|call|>commentary<|return|><|channel|>final<|message|>result<|end|>"
events, remainder = strategy.parse(text)
normal = [e for e in events if e.event_type == "normal"]
self.assertTrue(any("result" in e.content for e in normal))
def test_canonical_return_without_final(self):
"""Test that _parse_block returns None for block without proper end."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
# Channel block that has no message content before end
text = "<|start|><|channel|>final<|end|>"
events, remainder = strategy.parse(text)
# Should handle gracefully
self.assertIsInstance(events, list)
def test_iter_tokens_unknown_at_end_no_next_marker(self):
"""Test unknown token with |> close but no next <| marker after it."""
from sglang.srt.parser.harmony_parser import iter_tokens
@@ -984,18 +664,6 @@ class TestAdditionalEdgeCases(CustomTestCase):
normal = [e.content for e in events if e.event_type == "normal"]
self.assertTrue(any("answer" in c for c in normal))
def test_canonical_incomplete_parse_block_no_end(self):
"""Test that a channel block without END/CALL/RETURN returns None (incomplete)."""
from sglang.srt.parser.harmony_parser import CanonicalStrategy
strategy = CanonicalStrategy()
# Channel with message but no end token
text = "<|start|><|channel|>final<|message|>partial"
events, remainder = strategy.parse(text)
# Should be treated as incomplete
total = "".join(e.content for e in events) + remainder
self.assertIn("partial", total)
def test_text_strategy_commentary_only(self):
"""Test TextStrategy with commentary-only pattern (no 'assistantfinal')."""
from sglang.srt.parser.harmony_parser import TextStrategy
test/registered/unit/parser/test_reasoning_parser.py
-556
View File
@@ -13,7 +13,6 @@ from sglang.srt.parser.reasoning_parser import (
Nemotron3Detector,
Qwen3Detector,
ReasoningParser,
StreamingParseResult,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -21,20 +20,6 @@ from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=7, suite="base-a-test-cpu")
class TestStreamingParseResult(CustomTestCase):
def test_init_default(self):
"""Test default initialization of StreamingParseResult."""
result = StreamingParseResult()
self.assertEqual(result.normal_text, "")
self.assertEqual(result.reasoning_text, "")
def test_init_with_values(self):
"""Test initialization with specific values."""
result = StreamingParseResult("normal", "reasoning")
self.assertEqual(result.normal_text, "normal")
self.assertEqual(result.reasoning_text, "reasoning")
class TestBaseReasoningFormatDetector(CustomTestCase):
def setUp(self):
self.detector = BaseReasoningFormatDetector(
@@ -44,15 +29,6 @@ class TestBaseReasoningFormatDetector(CustomTestCase):
stream_reasoning=True,
)
def test_init(self):
"""Test initialization of BaseReasoningFormatDetector."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
self.assertEqual(self.detector._buffer, "")
self.assertFalse(self.detector.stripped_think_start)
def test_detect_and_parse_normal_text(self):
"""Test parsing normal text without reasoning."""
text = "This is normal text"
@@ -161,28 +137,6 @@ class TestDeepSeekR1Detector(CustomTestCase):
def setUp(self):
self.detector = DeepSeekR1Detector()
def test_init(self):
"""Test DeepSeekR1Detector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertTrue(self.detector._in_reasoning) # force_reasoning=True
self.assertTrue(self.detector.stream_reasoning)
def test_init_no_stream_reasoning(self):
"""Test DeepSeekR1Detector with stream_reasoning=False."""
detector = DeepSeekR1Detector(stream_reasoning=False)
self.assertFalse(detector.stream_reasoning)
def test_detect_and_parse_r1_format(self):
"""Test parsing DeepSeek-R1 format."""
text = "I need to think about this. The answer is 42."
result = self.detector.detect_and_parse(text)
# Should be treated as reasoning because force_reasoning=True
self.assertEqual(
result.reasoning_text, "I need to think about this. The answer is 42."
)
self.assertEqual(result.normal_text, "")
def test_detect_and_parse_with_end_token(self):
"""Test parsing with end token."""
text = "I think this is the answer</think>The final answer is 42."
@@ -203,20 +157,6 @@ class TestQwen3Detector(CustomTestCase):
def setUp(self):
self.detector = Qwen3Detector()
def test_init(self):
"""Test Qwen3Detector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertFalse(self.detector._in_reasoning) # force_reasoning=False
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_qwen3_format(self):
"""Test parsing Qwen3 format."""
text = "<think>Let me think about this problem</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Let me think about this problem")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_without_thinking(self):
"""Test parsing without thinking (enable_thinking=False case)."""
text = "Direct answer without thinking."
@@ -225,70 +165,10 @@ class TestQwen3Detector(CustomTestCase):
self.assertEqual(result.reasoning_text, "")
class TestQwen3ForcedReasoningDetector(CustomTestCase):
def setUp(self):
self.detector = Qwen3Detector(force_reasoning=True)
def test_init(self):
"""Test Qwen3ForcedReasoningDetector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertTrue(self.detector._in_reasoning) # force_reasoning=True
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_qwen3_forced_reasoning_format(self):
"""Test parsing Qwen3-ForcedReasoning format (no <think> start tag)."""
text = "I need to think about this step by step.</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(
result.reasoning_text, "I need to think about this step by step."
)
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_with_start_token(self):
"""Test parsing Qwen3-ForcedReasoning with optional <think> start tag."""
text = "<think>I need to think about this.</think>The answer is 42."
result = self.detector.detect_and_parse(text)
# Should work because base class logic handles both force_reasoning=True OR start token
self.assertEqual(result.reasoning_text, "I need to think about this.")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_streaming_qwen3_forced_reasoning_format(self):
"""Test streaming parse of Qwen3-ForcedReasoning format."""
# First chunk without <think> start
result = self.detector.parse_streaming_increment("I need to")
self.assertEqual(result.reasoning_text, "I need to")
self.assertEqual(result.normal_text, "")
# More reasoning content
result = self.detector.parse_streaming_increment(" think about this.")
self.assertEqual(result.reasoning_text, " think about this.")
self.assertEqual(result.normal_text, "")
# End token with normal text
result = self.detector.parse_streaming_increment("</think>The answer is 42.")
self.assertEqual(result.reasoning_text, "") # Buffer cleared
self.assertEqual(result.normal_text, "The answer is 42.")
class TestKimiDetector(CustomTestCase):
def setUp(self):
self.detector = KimiDetector()
def test_init(self):
"""Test KimiDetector initialization."""
self.assertEqual(self.detector.think_start_token, "◁think▷")
self.assertEqual(self.detector.think_end_token, "◁/think▷")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_kimi_format(self):
"""Test parsing Kimi format."""
text = "◁think▷Let me consider this carefully◁/think▷The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Let me consider this carefully")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_kimi_no_thinking(self):
"""Test parsing Kimi format without thinking."""
text = "Direct answer without thinking tokens."
@@ -296,29 +176,6 @@ class TestKimiDetector(CustomTestCase):
self.assertEqual(result.normal_text, text)
self.assertEqual(result.reasoning_text, "")
def test_streaming_kimi_format(self):
"""Test streaming parse of Kimi format."""
# Test partial token
result = self.detector.parse_streaming_increment("◁thi")
self.assertEqual(result.normal_text, "")
self.assertEqual(result.reasoning_text, "")
# Complete start token
result = self.detector.parse_streaming_increment("nk▷Start")
self.assertEqual(result.normal_text, "")
self.assertEqual(result.reasoning_text, "Start")
self.assertTrue(self.detector._in_reasoning)
# Add reasoning content
result = self.detector.parse_streaming_increment("thinking...")
self.assertEqual(result.reasoning_text, "thinking...")
self.assertEqual(result.normal_text, "")
# End token - reasoning content is cleared when end token is processed
result = self.detector.parse_streaming_increment("◁/think▷answer")
self.assertEqual(result.reasoning_text, "") # Buffer cleared
self.assertEqual(result.normal_text, "answer")
class TestKimiK2Detector(CustomTestCase):
"""Test cases for KimiK2 detector with tool interruption support."""
@@ -334,36 +191,6 @@ class TestKimiK2Detector(CustomTestCase):
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_tool_interrupt(self):
"""Test parsing with Kimi-K2 tool-section interruption."""
text = "<think>thinking<|tool_calls_section_begin|><|tool_call_begin|>"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "thinking")
self.assertEqual(
result.normal_text, "<|tool_calls_section_begin|><|tool_call_begin|>"
)
def test_streaming_tool_interrupt(self):
"""Test streaming parse interrupted by tool section."""
self.detector.parse_streaming_increment("<think>")
result1 = self.detector.parse_streaming_increment("reasoning")
self.assertEqual(result1.reasoning_text, "reasoning")
self.assertEqual(result1.normal_text, "")
result2 = self.detector.parse_streaming_increment(
"<|tool_calls_section_begin|>"
)
self.assertEqual(result2.reasoning_text, "")
self.assertEqual(result2.normal_text, "<|tool_calls_section_begin|>")
def test_streaming_after_interrupt_is_normal(self):
"""After interruption, subsequent chunks should be normal text."""
self.detector.parse_streaming_increment("<think>")
self.detector.parse_streaming_increment("reasoning<|tool_calls_section_begin|>")
result = self.detector.parse_streaming_increment("<|tool_call_begin|>")
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, "<|tool_call_begin|>")
class TestGlm45Detector(CustomTestCase):
"""Test cases for GLM45 detector with tool interruption support."""
@@ -371,33 +198,6 @@ class TestGlm45Detector(CustomTestCase):
def setUp(self):
self.detector = Glm45Detector()
def test_init(self):
"""Test Glm45Detector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertEqual(self.detector.tool_start_token, "<tool_call>")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_normal_reasoning(self):
"""Test parsing normal reasoning block without tool interruption."""
text = "<think>Let me think about this step by step</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Let me think about this step by step")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_tool_interrupt(self):
"""
Test parsing with tool interruption.
GLM45 can interrupt reasoning with tool token (<tool_call>) without closing </think>.
Should split at the first occurrence of tool_start_token using find().
"""
text = "<think>I need to think<tool_call>tool call data"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "I need to think")
self.assertEqual(result.normal_text, "<tool_call>tool call data")
def test_detect_and_parse_multiple_tool_calls_find(self):
"""
Test that find() finds the FIRST occurrence of tool_start_token.
@@ -413,17 +213,6 @@ class TestGlm45Detector(CustomTestCase):
"<tool_call>first tool<tool_call>second tool<tool_call>final tool",
)
def test_detect_and_parse_truncated_reasoning(self):
"""
Test truncated reasoning without tool or end tag.
Should return all content as reasoning_text.
"""
text = "<think>This is incomplete"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "This is incomplete")
self.assertEqual(result.normal_text, "")
def test_detect_and_parse_normal_text_only(self):
"""Test parsing text without reasoning block."""
text = "Just the answer without any reasoning."
@@ -431,50 +220,6 @@ class TestGlm45Detector(CustomTestCase):
self.assertEqual(result.normal_text, text)
self.assertEqual(result.reasoning_text, "")
def test_streaming_normal_flow(self):
"""Test streaming with normal reasoning flow."""
# Start reasoning
result1 = self.detector.parse_streaming_increment("<think>")
self.assertEqual(result1.normal_text, "")
self.assertEqual(result1.reasoning_text, "")
self.assertTrue(self.detector._in_reasoning)
# Reasoning content
result2 = self.detector.parse_streaming_increment("thinking...")
self.assertEqual(result2.normal_text, "")
self.assertEqual(result2.reasoning_text, "thinking...")
# End reasoning
result3 = self.detector.parse_streaming_increment("</think>answer")
self.assertEqual(result3.normal_text, "answer")
self.assertEqual(result3.reasoning_text, "")
self.assertFalse(self.detector._in_reasoning)
def test_streaming_tool_interrupt_split_tokens(self):
"""
Test streaming with tool interruption where tool token is split across chunks.
This tests the buffer prefix logic that prevents partial emission of tool token.
"""
# Start reasoning
self.detector.parse_streaming_increment("<think>")
# Add reasoning
result1 = self.detector.parse_streaming_increment("thinking")
self.assertEqual(result1.reasoning_text, "thinking")
# Send partial tool token (should be buffered, not emitted)
result2 = self.detector.parse_streaming_increment("<tool_call>")
# Tool token is in buffer, causing switch to normal mode
self.assertEqual(result2.reasoning_text, "")
self.assertEqual(result2.normal_text, "<tool_call>")
self.assertFalse(self.detector._in_reasoning)
# Send tool args
result3 = self.detector.parse_streaming_increment("tool args")
self.assertEqual(result3.reasoning_text, "")
self.assertEqual(result3.normal_text, "tool args")
def test_streaming_no_stream_reasoning(self):
"""Test streaming without stream_reasoning enabled."""
detector = Glm45Detector(stream_reasoning=False)
@@ -526,21 +271,6 @@ class TestHunyuanDetector(CustomTestCase):
def setUp(self):
self.detector = HunyuanDetector()
def test_init(self):
"""Test HunyuanDetector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertEqual(self.detector.tool_start_token, "<tool_calls>")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_normal_reasoning(self):
"""Test parsing normal reasoning block without tool interruption."""
text = "<think>Let me think about this</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Let me think about this")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_without_thinking(self):
"""Test parsing without thinking tokens (no_think mode)."""
text = "Direct answer without thinking."
@@ -548,42 +278,6 @@ class TestHunyuanDetector(CustomTestCase):
self.assertEqual(result.normal_text, text)
self.assertEqual(result.reasoning_text, "")
def test_detect_and_parse_tool_interrupt(self):
"""Test parsing with tool call interruption during reasoning."""
text = "<think>I need to check<tool_calls><tool_call>get_weather<tool_sep></tool_call></tool_calls>"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "I need to check")
self.assertIn("<tool_calls>", result.normal_text)
def test_streaming_normal_reasoning(self):
"""Test streaming parse of normal reasoning block."""
self.detector.parse_streaming_increment("<think>")
result1 = self.detector.parse_streaming_increment("reasoning content")
self.assertEqual(result1.reasoning_text, "reasoning content")
result2 = self.detector.parse_streaming_increment("</think>answer")
self.assertEqual(result2.normal_text, "answer")
self.assertFalse(self.detector._in_reasoning)
def test_streaming_tool_interrupt(self):
"""Test streaming parse interrupted by tool call section."""
self.detector.parse_streaming_increment("<think>")
result1 = self.detector.parse_streaming_increment("thinking")
self.assertEqual(result1.reasoning_text, "thinking")
result2 = self.detector.parse_streaming_increment("<tool_calls>")
self.assertEqual(result2.reasoning_text, "")
self.assertEqual(result2.normal_text, "<tool_calls>")
self.assertFalse(self.detector._in_reasoning)
def test_streaming_after_interrupt_is_normal(self):
"""After tool interruption, subsequent chunks should be normal text."""
self.detector.parse_streaming_increment("<think>")
self.detector.parse_streaming_increment("reasoning<tool_calls>")
result = self.detector.parse_streaming_increment("<tool_call>data")
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, "<tool_call>data")
def test_reasoning_parser_integration(self):
"""Test Hunyuan through ReasoningParser API."""
parser = ReasoningParser("hunyuan")
@@ -617,21 +311,6 @@ class TestNemotron3Detector(CustomTestCase):
def setUp(self):
self.detector = Nemotron3Detector()
def test_init(self):
"""Test Nemotron3Detector initialization."""
self.assertEqual(self.detector.think_start_token, "<think>")
self.assertEqual(self.detector.think_end_token, "</think>")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
self.assertFalse(self.detector._force_nonempty_content)
def test_detect_and_parse_complete_reasoning(self):
"""Test parsing complete reasoning block."""
text = "<think>Let me think about this</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Let me think about this")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_detect_and_parse_no_thinking(self):
"""Test parsing without thinking tokens."""
text = "Direct answer without thinking."
@@ -671,28 +350,11 @@ class TestNemotron3Detector(CustomTestCase):
self.assertEqual(result.normal_text, "Truncated reasoning without end token")
self.assertEqual(result.reasoning_text, "")
def test_force_nonempty_content_no_thinking_tokens(self):
"""Test force_nonempty_content with plain text (no thinking tokens)."""
detector = Nemotron3Detector(force_nonempty_content=True)
text = "Plain text without any thinking."
result = detector.detect_and_parse(text)
# Normal text already exists, no swap needed
self.assertEqual(result.normal_text, text)
self.assertEqual(result.reasoning_text, "")
class TestGemma4Detector(CustomTestCase):
def setUp(self):
self.detector = Gemma4Detector()
def test_init(self):
"""Test Gemma4Detector initialization."""
self.assertEqual(self.detector.think_start_token, "<|channel>")
self.assertEqual(self.detector.think_end_token, "<channel|>")
self.assertEqual(self.detector.think_start_self_label, "thought\n")
self.assertFalse(self.detector._in_reasoning)
self.assertTrue(self.detector.stream_reasoning)
def test_detect_and_parse_complete_reasoning(self):
"""Test parsing complete Gemma4 reasoning block (think_start_self_label is stripped)."""
text = "<|channel>thought\nLet me think about this<channel|>The answer is 42."
@@ -707,49 +369,6 @@ class TestGemma4Detector(CustomTestCase):
self.assertEqual(result.normal_text, text)
self.assertEqual(result.reasoning_text, "")
def test_detect_and_parse_reasoning_only(self):
"""Test parsing when output is all reasoning (no end token yet)."""
text = "<|channel>thought\nStill thinking..."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Still thinking...")
self.assertEqual(result.normal_text, "")
def test_streaming_complete_flow(self):
"""Test streaming parse of Gemma4 reasoning flow."""
chunks = [
"<|channel>",
"thought\nreasoning content",
"<channel|>",
"final answer",
]
all_reasoning = ""
all_normal = ""
for chunk in chunks:
result = self.detector.parse_streaming_increment(chunk)
all_reasoning += result.reasoning_text
all_normal += result.normal_text
self.assertIn("reasoning content", all_reasoning)
self.assertIn("final answer", all_normal)
def test_streaming_full_start_sequence(self):
"""Test streaming with the full start sequence (token + self_label)."""
# Gemma4 start sequence is "<|channel>thought\n", not just "<|channel>"
result = self.detector.parse_streaming_increment("<|channel>thought\n")
self.assertEqual(result.normal_text, "")
self.assertEqual(result.reasoning_text, "")
self.assertTrue(self.detector._in_reasoning)
result = self.detector.parse_streaming_increment("reasoning content")
self.assertEqual(result.reasoning_text, "reasoning content")
self.assertEqual(result.normal_text, "")
def test_streaming_partial_start_buffered(self):
"""Test that partial start sequence is buffered."""
# "<|channel>" alone is a prefix of "<|channel>thought\n", so it's buffered
result = self.detector.parse_streaming_increment("<|channel>")
self.assertEqual(result.normal_text, "")
self.assertEqual(result.reasoning_text, "")
def test_streaming_end_token_mid_chunk(self):
"""Test end token arriving in the same chunk as reasoning content."""
self.detector.parse_streaming_increment("<|channel>thought\n")
@@ -760,18 +379,6 @@ class TestGemma4Detector(CustomTestCase):
self.assertEqual(result.normal_text, "the answer")
self.assertFalse(self.detector._in_reasoning)
def test_streaming_split_end_token(self):
"""Test end token split across two chunks."""
self.detector.parse_streaming_increment("<|channel>thought\n")
self.detector.parse_streaming_increment("reasoning content")
result1 = self.detector.parse_streaming_increment("<chan")
self.assertEqual(result1.normal_text, "")
result2 = self.detector.parse_streaming_increment("nel|>final answer")
self.assertFalse(self.detector._in_reasoning)
self.assertIn("final answer", result2.normal_text)
def test_streaming_self_label_split_across_chunks(self):
"""Test self_label ('thought\\n') arriving separately from start token."""
result1 = self.detector.parse_streaming_increment("<|channel>")
@@ -784,37 +391,6 @@ class TestGemma4Detector(CustomTestCase):
result3 = self.detector.parse_streaming_increment("reasoning here")
self.assertEqual(result3.reasoning_text, "reasoning here")
def test_streaming_force_reasoning(self):
"""Test streaming with force_reasoning=True (no start token needed)."""
detector = Gemma4Detector(force_reasoning=True)
result1 = detector.parse_streaming_increment("reasoning content")
self.assertEqual(result1.reasoning_text, "reasoning content")
self.assertEqual(result1.normal_text, "")
result2 = detector.parse_streaming_increment("<channel|>the answer")
self.assertFalse(detector._in_reasoning)
self.assertIn("the answer", result2.normal_text)
def test_streaming_multiple_reasoning_chunks(self):
"""Test reasoning content arriving in many small chunks."""
self.detector.parse_streaming_increment("<|channel>thought\n")
all_reasoning = ""
for chunk in ["Think", "ing ", "step ", "by ", "step."]:
result = self.detector.parse_streaming_increment(chunk)
all_reasoning += result.reasoning_text
self.assertEqual(result.normal_text, "")
self.assertEqual(all_reasoning, "Thinking step by step.")
def test_force_reasoning(self):
"""Test Gemma4Detector with force_reasoning=True."""
detector = Gemma4Detector(force_reasoning=True)
text = "This should be reasoning<channel|>The answer."
result = detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "This should be reasoning")
self.assertEqual(result.normal_text, "The answer.")
class TestReasoningParser(CustomTestCase):
def test_init_valid_model(self):
@@ -990,42 +566,6 @@ class TestReasoningParser(CustomTestCase):
class TestIntegrationScenarios(CustomTestCase):
"""Integration tests for realistic usage scenarios."""
def test_deepseek_r1_complete_response(self):
"""Test complete DeepSeek-R1 response parsing."""
parser = ReasoningParser("deepseek-r1")
text = "I need to solve this step by step. First, I'll analyze the problem. The given equation is x + 2 = 5. To solve for x, I subtract 2 from both sides: x = 5 - 2 = 3.</think>The answer is x = 3."
reasoning, normal = parser.parse_non_stream(text)
self.assertIn("step by step", reasoning)
self.assertIn(
"= 3", reasoning
) # The reasoning contains "x = 5 - 2 = 3" which has "= 3"
self.assertEqual(normal, "The answer is x = 3.")
def test_qwen3_streaming_scenario(self):
"""Test Qwen3 streaming scenario."""
parser = ReasoningParser("qwen3")
chunks = [
"<think>",
"Let me analyze this problem.",
" I need to consider multiple factors.",
"</think>",
"Based on my analysis, the solution is to use a different approach.",
]
all_reasoning = ""
all_normal = ""
for chunk in chunks:
reasoning, normal = parser.parse_stream_chunk(chunk)
all_reasoning += reasoning
all_normal += normal
self.assertIn("analyze", all_reasoning)
self.assertIn("multiple factors", all_reasoning)
self.assertIn("different approach", all_normal)
def test_kimi_streaming_scenario(self):
"""Test Kimi streaming scenario."""
parser = ReasoningParser("kimi")
@@ -1049,35 +589,6 @@ class TestIntegrationScenarios(CustomTestCase):
self.assertIn("multiple factors", all_reasoning)
self.assertIn("42", all_normal)
def test_gemma4_complete_response(self):
"""Test complete Gemma4 response parsing (think_start_self_label stripped)."""
parser = ReasoningParser("gemma4")
text = "<|channel>thought\nI need to solve x + 2 = 5. Subtracting 2: x = 3.<channel|>The answer is x = 3."
reasoning, normal = parser.parse_non_stream(text)
self.assertIn("x = 3", reasoning)
self.assertNotIn("thought\n", reasoning)
self.assertEqual(normal, "The answer is x = 3.")
def test_gemma4_streaming_scenario(self):
"""Test Gemma4 streaming scenario."""
parser = ReasoningParser("gemma4")
chunks = [
"<|channel>",
"thought\nLet me analyze.",
" Multiple factors.",
"<channel|>",
"The solution is 42.",
]
all_reasoning = ""
all_normal = ""
for chunk in chunks:
reasoning, normal = parser.parse_stream_chunk(chunk)
all_reasoning += reasoning
all_normal += normal
self.assertIn("analyze", all_reasoning)
self.assertIn("Multiple factors", all_reasoning)
self.assertIn("42", all_normal)
def test_empty_reasoning_blocks(self):
"""Test handling of empty reasoning blocks."""
parser = ReasoningParser("qwen3")
@@ -1157,22 +668,6 @@ class TestBufferLossBugFix(CustomTestCase):
self.assertEqual(result2.normal_text, "</answer")
self.assertEqual(result2.reasoning_text, "")
def test_partial_start_tag_buffer_preservation(self):
"""
Test that partial start tag fragments are properly preserved.
"""
detector = BaseReasoningFormatDetector("<think>", "</think>")
# Send partial start tag
result1 = detector.parse_streaming_increment("<th")
self.assertEqual(result1.normal_text, "")
self.assertEqual(result1.reasoning_text, "")
# Complete with non-matching text
result2 = detector.parse_streaming_increment("is is text")
self.assertEqual(result2.normal_text, "<this is text")
self.assertEqual(result2.reasoning_text, "")
def test_partial_end_tag_in_reasoning_mode(self):
"""
Test partial end tag handling when already in reasoning mode.
@@ -1194,25 +689,6 @@ class TestBufferLossBugFix(CustomTestCase):
# The reasoning text should be empty since buffer was cleared when end tag was processed
self.assertEqual(result2.reasoning_text, "")
def test_multiple_partial_fragments(self):
"""
Test handling of multiple partial fragments that don't match any tokens.
"""
detector = BaseReasoningFormatDetector("<think>", "</think>")
# Send multiple partial fragments
result1 = detector.parse_streaming_increment("<")
self.assertEqual(result1.normal_text, "")
self.assertEqual(result1.reasoning_text, "")
result2 = detector.parse_streaming_increment("/")
self.assertEqual(result2.normal_text, "")
self.assertEqual(result2.reasoning_text, "")
result3 = detector.parse_streaming_increment("random>")
self.assertEqual(result3.normal_text, "</random>")
self.assertEqual(result3.reasoning_text, "")
def test_edge_case_exact_token_match(self):
"""
Test edge case where buffer content exactly matches a token.
@@ -1242,19 +718,6 @@ class TestGptOssDetector(CustomTestCase):
self.detector = GptOssDetector()
def test_detect_and_parse_with_analysis_and_final(self):
"""Test one-shot parsing with analysis (reasoning) and final (normal) blocks."""
text = "<|start|><|channel|>analysis<|message|>thinking hard<|end|><|channel|>final<|message|>the answer<|end|>"
result = self.detector.detect_and_parse(text)
self.assertIn("thinking hard", result.reasoning_text)
self.assertIn("the answer", result.normal_text)
def test_detect_and_parse_normal_only(self):
"""Test one-shot parsing with only final block."""
text = "<|start|><|channel|>final<|message|>just the answer<|end|>"
result = self.detector.detect_and_parse(text)
self.assertIn("just the answer", result.normal_text)
def test_streaming_analysis_then_final(self):
"""Test streaming parse across multiple chunks."""
chunks = [
@@ -1273,13 +736,6 @@ class TestGptOssDetector(CustomTestCase):
self.assertIn("reasoning part", all_reasoning)
self.assertIn("answer", all_normal)
def test_streaming_with_tool_call(self):
"""Test streaming parse with tool call events."""
text = "<|start|><|channel|>analysis<|message|>think<|end|><|call|>tool_data<|return|><|channel|>final<|message|>result<|end|>"
result = self.detector.detect_and_parse(text)
self.assertIn("think", result.reasoning_text)
self.assertIn("result", result.normal_text)
class TestMiniMaxAppendThinkDetector(CustomTestCase):
"""Test cases for MiniMaxAppendThinkDetector."""
@@ -1478,18 +934,6 @@ class TestContinueFinalMessage(CustomTestCase):
self.assertEqual(result.reasoning_text, "new reasoning")
self.assertEqual(result.normal_text, "new answer")
def test_streaming_returns_empty_when_in_reasoning_and_end_buffered(self):
"""Test that streaming returns empty when buffer could be partial end token."""
detector = BaseReasoningFormatDetector(
"<think>", "</think>", force_reasoning=True, stream_reasoning=True
)
# In reasoning mode, send partial end token
result = detector.parse_streaming_increment("</")
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, "")
# This goes through the path where _in_reasoning is True but buffer
# is a prefix of think_end_token → returns empty
class TestGptOssDetectorToolCall(CustomTestCase):
"""Test GptOssDetector tool_call raw_text handling."""
test/registered/unit/platforms/test_platform_interface.py
+2 -182
View File
@@ -10,8 +10,8 @@ from unittest.mock import MagicMock, patch
import torch
from sglang.srt.platforms import _load_platform_class, _resolve_platform
from sglang.srt.platforms.cpu import CpuDeviceMixin, CpuSRTPlatform
from sglang.srt.platforms.cuda import CudaDeviceMixin, CudaSRTPlatform
from sglang.srt.platforms.cpu import CpuSRTPlatform
from sglang.srt.platforms.cuda import CudaSRTPlatform
from sglang.srt.platforms.device_mixin import (
CpuArchEnum,
DeviceCapability,
@@ -47,41 +47,6 @@ def _make_device_mixin(enum, name, dtype):
return M()
class _StubPlatform(SRTPlatform):
"""Concrete SRTPlatform with minimal defaults for testing overrides."""
_enum = PlatformEnum.CUDA
device_name = "cuda"
device_type = "cuda"
def get_device_total_memory(self, device_id=0):
return 10**9
def get_current_memory_usage(self, device=None):
return 5 * 10**8
def get_default_attention_backend(self):
return "flashinfer"
def get_graph_runner_cls(self):
return object
def get_mha_kv_pool_cls(self):
return object
def get_mla_kv_pool_cls(self):
return object
def get_dsa_kv_pool_cls(self):
return object
def get_paged_allocator_cls(self):
return object
def get_piecewise_backend_cls(self):
return object
def _make_platform_ep(name, load_fn=None):
"""Create a mock entry point for platform plugins."""
ep = MagicMock()
@@ -240,78 +205,12 @@ class TestCudaDeviceMixin(CustomTestCase):
base = CudaSRTPlatform()
self.assertEqual(base.get_device(2), torch.device("cuda", 2))
def test_cuda_platform_identity(self):
base = CudaSRTPlatform()
self.assertTrue(base.is_cuda())
self.assertTrue(base.is_cuda_alike())
self.assertIsInstance(base, CudaDeviceMixin)
@patch("torch.cuda.get_device_properties")
def test_default_get_device_total_memory_uses_cuda(
self, mock_get_device_properties
):
mock_get_device_properties.return_value.total_memory = 123
base = CudaSRTPlatform()
self.assertEqual(base.get_device_total_memory(1), 123)
mock_get_device_properties.assert_called_once_with(1)
@patch("torch.cuda.max_memory_allocated", return_value=456)
def test_default_get_current_memory_usage_uses_cuda(
self, mock_max_memory_allocated
):
base = CudaSRTPlatform()
device = torch.device("cuda", 1)
self.assertEqual(base.get_current_memory_usage(device), 456.0)
mock_max_memory_allocated.assert_called_once_with(device)
@patch("torch.cuda.set_device")
def test_default_set_device_uses_cuda(self, mock_set_device):
base = CudaSRTPlatform()
device = torch.device("cuda", 1)
base.set_device(device)
mock_set_device.assert_called_once_with(device)
@patch("torch.cuda.get_device_name", return_value="NVIDIA H100")
def test_default_get_device_name_uses_cuda(self, mock_get_device_name):
base = CudaSRTPlatform()
self.assertEqual(base.get_device_name(1), "NVIDIA H100")
mock_get_device_name.assert_called_once_with(1)
@patch("torch.cuda.get_device_properties")
def test_default_get_device_uuid_uses_cuda(self, mock_get_device_properties):
mock_get_device_properties.return_value.uuid = "1234"
base = CudaSRTPlatform()
self.assertEqual(base.get_device_uuid(1), "1234")
mock_get_device_properties.assert_called_once_with(1)
@patch("torch.cuda.get_device_capability", return_value=(9, 0))
def test_default_get_device_capability_uses_cuda(self, mock_get_device_capability):
base = CudaSRTPlatform()
self.assertEqual(base.get_device_capability(1), DeviceCapability(9, 0))
mock_get_device_capability.assert_called_once_with(1)
@patch("torch.cuda.empty_cache")
def test_default_empty_cache_uses_cuda(self, mock_empty_cache):
base = CudaSRTPlatform()
base.empty_cache()
mock_empty_cache.assert_called_once_with()
@patch("torch.cuda.synchronize")
def test_default_synchronize_uses_cuda(self, mock_synchronize):
base = CudaSRTPlatform()
base.synchronize()
mock_synchronize.assert_called_once_with()
@patch("torch.cuda.mem_get_info", return_value=(123, 456), create=True)
def test_default_get_available_memory_uses_cuda(self, mock_mem_get_info):
base = CudaSRTPlatform()
self.assertEqual(base.get_available_memory(1), (123, 456))
mock_mem_get_info.assert_called_once_with(1)
def test_default_distributed_backend_is_nccl(self):
base = CudaSRTPlatform()
self.assertEqual(base.get_torch_distributed_backend_str(), "nccl")
@patch("torch.cuda.manual_seed_all")
@patch("torch.manual_seed")
@patch("sglang.srt.platforms.device_mixin.np.random.seed")
@@ -335,32 +234,12 @@ class TestCudaDeviceMixin(CustomTestCase):
class TestCpuDeviceMixin(CustomTestCase):
"""Tests for CPU device operation defaults (covers both x86 and ARM)."""
def test_cpu_platform_identity(self):
base = CpuSRTPlatform()
self.assertTrue(base.is_cpu())
self.assertFalse(base.is_cuda())
self.assertFalse(base.is_cuda_alike())
self.assertIsInstance(base, CpuDeviceMixin)
def test_default_get_device_returns_cpu_device(self):
base = CpuSRTPlatform()
# ``local_rank`` is ignored — CPU has no per-rank device.
self.assertEqual(base.get_device(0), torch.device("cpu"))
self.assertEqual(base.get_device(7), torch.device("cpu"))
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
def test_default_get_device_total_memory_uses_psutil(self, mock_vm):
mock_vm.return_value.total = 12345
base = CpuSRTPlatform()
self.assertEqual(base.get_device_total_memory(), 12345)
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
def test_default_get_available_memory_uses_psutil(self, mock_vm):
mock_vm.return_value.available = 100
mock_vm.return_value.total = 200
base = CpuSRTPlatform()
self.assertEqual(base.get_available_memory(), (100, 200))
@patch("sglang.srt.platforms.cpu.psutil.virtual_memory")
def test_default_get_current_memory_usage_is_system_used(self, mock_vm):
mock_vm.return_value.total = 1000
@@ -378,14 +257,6 @@ class TestCpuDeviceMixin(CustomTestCase):
free = base.get_device_total_memory() - base.get_current_memory_usage()
self.assertEqual(free, 300)
@patch("torch.cpu.set_device")
def test_default_set_device_uses_torch_cpu(self, mock_set_device):
base = CpuSRTPlatform()
device = torch.device("cpu")
base.set_device(device)
# Documented CPU no-op, but called for symmetry with CudaDeviceMixin.
mock_set_device.assert_called_once_with(device)
def test_default_set_device_does_not_flip_default(self):
base = CpuSRTPlatform()
# Must not call torch.set_default_device — process-wide default stays put.
@@ -394,22 +265,6 @@ class TestCpuDeviceMixin(CustomTestCase):
after = torch.empty(0).device
self.assertEqual(before, after)
@patch("sglang.srt.platforms.cpu.gc.collect")
def test_default_empty_cache_calls_gc_collect(self, mock_collect):
base = CpuSRTPlatform()
base.empty_cache()
mock_collect.assert_called_once_with()
@patch("torch.cpu.synchronize")
def test_default_synchronize_uses_torch_cpu(self, mock_synchronize):
base = CpuSRTPlatform()
base.synchronize()
mock_synchronize.assert_called_once_with()
def test_default_distributed_backend_is_gloo(self):
base = CpuSRTPlatform()
self.assertEqual(base.get_torch_distributed_backend_str(), "gloo")
@patch("platform.machine", return_value="aarch64")
def test_cpu_arch_property_resolves_and_caches(self, mock_machine):
base = CpuSRTPlatform()
@@ -431,15 +286,6 @@ class TestCpuDeviceMixin(CustomTestCase):
name = base.get_device_name()
self.assertIn("x86_64", name)
@patch("platform.machine", return_value="aarch64")
def test_get_device_uuid_returns_machine(self, _mock_machine):
base = CpuSRTPlatform()
self.assertEqual(base.get_device_uuid(), "aarch64")
def test_get_device_capability_returns_none(self):
base = CpuSRTPlatform()
self.assertIsNone(base.get_device_capability())
def test_cpu_srt_platform_capabilities(self):
base = CpuSRTPlatform()
self.assertFalse(base.supports_fp8())
@@ -449,32 +295,6 @@ class TestCpuDeviceMixin(CustomTestCase):
self.assertFalse(base.is_pin_memory_available())
class TestSRTPlatformOverrides(CustomTestCase):
"""Tests for SRTPlatform method overrides via plugins."""
def test_custom_get_dispatch_key_name(self):
class P(_StubPlatform):
_enum = PlatformEnum.NPU
device_name = "npu"
device_type = "npu"
def get_dispatch_key_name(self):
return "npu"
self.assertEqual(P().get_dispatch_key_name(), "npu")
def test_custom_get_compile_backend(self):
class P(_StubPlatform):
_enum = PlatformEnum.NPU
device_name = "npu"
device_type = "npu"
def get_compile_backend(self, mode=None):
return "inductor"
self.assertEqual(P().get_compile_backend(mode="npugraph_ex"), "inductor")
# ---------------------------------------------------------------------------
# Platform Discovery: _resolve_platform
# ---------------------------------------------------------------------------
test/registered/unit/server_args/test_server_args.py
-154
View File
@@ -33,21 +33,6 @@ _mock_device.start()
class TestPrepareServerArgs(CustomTestCase):
def test_prepare_server_args(self):
server_args = prepare_server_args(
[
"--model-path",
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
"--json-model-override-args",
'{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}}',
]
)
self.assertEqual(server_args.model_path, DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN)
self.assertEqual(
json.loads(server_args.json_model_override_args),
{"rope_scaling": {"factor": 2.0, "rope_type": "linear"}},
)
def test_config_nested_dict_args_are_json(self):
with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
f.write("mm-process-config:\n image:\n resize: 128\n")
@@ -145,15 +130,6 @@ class TestLoadBalanceMethod(unittest.TestCase):
str(context.exception),
)
def test_pd_decode_radix_cache_allows_mooncake(self):
server_args = self._load_balance_args(
disaggregation_mode="decode",
disaggregation_decode_enable_radix_cache=True,
disaggregation_transfer_backend="mooncake",
)
self.assertFalse(server_args.disable_radix_cache)
def test_pd_decode_radix_cache_rejects_fake_backend(self):
server_args = ServerArgs(
model_path="dummy",
@@ -170,15 +146,6 @@ class TestLoadBalanceMethod(unittest.TestCase):
str(context.exception),
)
def test_pd_decode_radix_cache_allows_ascend(self):
server_args = self._load_balance_args(
disaggregation_mode="decode",
disaggregation_decode_enable_radix_cache=True,
disaggregation_transfer_backend="ascend",
)
self.assertFalse(server_args.disable_radix_cache)
def test_pd_decode_radix_cache_allows_mooncake_tcp(self):
server_args = self._load_balance_args(
disaggregation_mode="decode",
@@ -403,14 +370,6 @@ class TestContextParallelServerArgs(CustomTestCase):
setattr(server_args, key, value)
return server_args
def test_canonical_prefill_cp_cli_sets_unified_fields(self):
args = self.parser.parse_args(
["--model", "dummy", "--enable-prefill-cp", "--cp-strategy", "interleave"]
)
self.assertTrue(args.enable_prefill_cp)
self.assertEqual(args.cp_strategy, "interleave")
def test_canonical_prefill_cp_requires_strategy(self):
args = self.parser.parse_args(["--model", "dummy", "--enable-prefill-cp"])
@@ -558,27 +517,6 @@ class TestContextParallelServerArgs(CustomTestCase):
class TestPortArgs(unittest.TestCase):
@patch("sglang.srt.server_args.get_free_port")
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
def test_init_new_with_nccl_port_none(self, mock_temp_file, mock_get_free_port):
"""Test that get_free_port() is called when nccl_port is None"""
mock_temp_file.return_value.name = "temp_file"
mock_get_free_port.return_value = 45678 # Mock ephemeral port
# Use MagicMock here to verify get_free_port is called
server_args = MagicMock()
server_args.nccl_port = None
server_args.enable_dp_attention = False
server_args.tokenizer_worker_num = 1
port_args = PortArgs.init_new(server_args)
# Verify get_free_port was called
mock_get_free_port.assert_called_once()
# Verify the returned port is used
self.assertEqual(port_args.nccl_port, 45678)
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
def test_init_new_standard_case(self, mock_temp_file):
mock_temp_file.return_value.name = "temp_file"
@@ -727,13 +665,6 @@ class TestSSLArgs(unittest.TestCase):
server_args._handle_ssl_validation()
return server_args
def test_default_ssl_fields_are_none(self):
server_args = ServerArgs(model_path="dummy")
self.assertIsNone(server_args.ssl_keyfile)
self.assertIsNone(server_args.ssl_certfile)
self.assertIsNone(server_args.ssl_ca_certs)
self.assertIsNone(server_args.ssl_keyfile_password)
def test_ssl_keyfile_without_certfile_raises(self):
with self.assertRaises(ValueError) as context:
self._validate_ssl(ssl_keyfile="key.pem")
@@ -744,12 +675,6 @@ class TestSSLArgs(unittest.TestCase):
self._validate_ssl(ssl_certfile="cert.pem")
self.assertIn("--ssl-keyfile", str(context.exception))
@patch("os.path.isfile", return_value=True)
def test_ssl_both_keyfile_and_certfile_accepted(self, _mock_isfile):
server_args = self._validate_ssl(ssl_keyfile="key.pem", ssl_certfile="cert.pem")
self.assertEqual(server_args.ssl_keyfile, "key.pem")
self.assertEqual(server_args.ssl_certfile, "cert.pem")
def test_url_returns_http_without_ssl(self):
server_args = ServerArgs(model_path="dummy")
self.assertTrue(server_args.url().startswith("http://"))
@@ -767,27 +692,6 @@ class TestSSLArgs(unittest.TestCase):
server_args = self._validate_ssl(ssl_keyfile="key.pem", ssl_certfile="cert.pem")
self.assertTrue(server_args.url().startswith("https://"))
@patch("os.path.isfile", return_value=True)
def test_ssl_cli_args_parsed(self, _mock_isfile):
server_args = prepare_server_args(
[
"--model-path",
"dummy",
"--ssl-keyfile",
"key.pem",
"--ssl-certfile",
"cert.pem",
"--ssl-ca-certs",
"ca.pem",
"--ssl-keyfile-password",
"secret",
]
)
self.assertEqual(server_args.ssl_keyfile, "key.pem")
self.assertEqual(server_args.ssl_certfile, "cert.pem")
self.assertEqual(server_args.ssl_ca_certs, "ca.pem")
self.assertEqual(server_args.ssl_keyfile_password, "secret")
def test_ssl_verify_without_ssl(self):
server_args = ServerArgs(model_path="dummy")
self.assertIs(server_args.ssl_verify(), True)
@@ -846,10 +750,6 @@ class TestSSLArgs(unittest.TestCase):
"SSL CA certificates file not found", str(context.exception)
)
def test_enable_ssl_refresh_default_false(self):
server_args = ServerArgs(model_path="dummy")
self.assertFalse(server_args.enable_ssl_refresh)
def test_enable_ssl_refresh_without_ssl_raises(self):
with self.assertRaises(ValueError) as context:
self._validate_ssl(enable_ssl_refresh=True)
@@ -865,21 +765,6 @@ class TestSSLArgs(unittest.TestCase):
)
self.assertTrue(server_args.enable_ssl_refresh)
@patch("os.path.isfile", return_value=True)
def test_enable_ssl_refresh_cli_flag(self, _mock_isfile):
server_args = prepare_server_args(
[
"--model-path",
"dummy",
"--ssl-keyfile",
"key.pem",
"--ssl-certfile",
"cert.pem",
"--enable-ssl-refresh",
]
)
self.assertTrue(server_args.enable_ssl_refresh)
class TestHiCacheArgs(unittest.TestCase):
def _make_args(self, **overrides) -> ServerArgs:
@@ -984,28 +869,6 @@ class TestHiCacheArgs(unittest.TestCase):
class TestNgramExternalSamArgs(CustomTestCase):
def test_prepare_server_args_parses_external_sam_args(self):
server_args = prepare_server_args(
[
"--model-path",
"dummy",
"--speculative-algorithm",
"NGRAM",
"--speculative-ngram-external-corpus-path",
"/tmp/ngram-corpus.jsonl",
"--speculative-ngram-external-sam-budget",
"4",
"--speculative-ngram-external-corpus-max-tokens",
"128",
]
)
self.assertEqual(
server_args.speculative_ngram_external_corpus_path,
"/tmp/ngram-corpus.jsonl",
)
self.assertEqual(server_args.speculative_ngram_external_sam_budget, 4)
self.assertEqual(server_args.speculative_ngram_external_corpus_max_tokens, 128)
def _make_dummy_ngram_args(self, **overrides):
args = ServerArgs(model_path="dummy")
args.speculative_algorithm = "NGRAM"
@@ -1070,13 +933,6 @@ class TestDecoupledSpecArgs(CustomTestCase):
self.assertEqual(server_args.decoupled_spec_rank, 0)
self.assertEqual(server_args.spec_trace_dir, "/tmp/tr")
def test_decoupled_spec_role_defaults_to_null(self):
server_args = prepare_server_args(["--model-path", "dummy"])
self.assertEqual(server_args.decoupled_spec_role, "null")
self.assertIsNone(server_args.decoupled_spec_bind_endpoint)
self.assertIsNone(server_args.decoupled_spec_connect_endpoints)
self.assertIsNone(server_args.decoupled_spec_rank)
def test_decoupled_spec_role_rejects_invalid_choice(self):
with self.assertRaises(SystemExit):
prepare_server_args(
@@ -1453,12 +1309,6 @@ class TestGrpcServerArgs(CustomTestCase):
def _args(**kwargs):
return ServerArgs(model_path="dummy", **kwargs)
def test_defaults_native_grpc_off_legacy_off(self):
sa = self._args()
sa._handle_deprecated_args()
self.assertIsNone(sa.grpc_port)
self.assertFalse(sa.smg_grpc_mode)
def test_http_only_high_port_does_not_derive_grpc_port(self):
sa = self._args(port=56000)
sa._handle_deprecated_args()
@@ -1587,10 +1437,6 @@ class TestTwoBatchOverlapBackend(CustomTestCase):
args = self._args(moe_a2a_backend="deepep", enable_dp_attention=False)
args._check_two_batch_overlap()
def test_tbo_disabled_is_noop(self):
args = self._args(enable_two_batch_overlap=False, enable_dp_attention=False)
args._check_two_batch_overlap()
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_ngram_corpus.py
-182
View File
@@ -177,12 +177,6 @@ class TestNgramCorpusBFS(CustomTestCase):
def test_masks(self):
np.testing.assert_array_equal(self.masks.tolist(), EXPECTED_BFS_MASKS)
def test_output_shapes(self):
n_queries = len(QUERY_SEQUENCES)
draft = 8
self.assertEqual(self.ids.shape, (n_queries, draft))
self.assertEqual(self.masks.shape, (n_queries, draft, draft))
class TestNgramCorpusProb(CustomTestCase):
"""Golden-output tests for Prob matching mode."""
@@ -202,35 +196,6 @@ class TestNgramCorpusProb(CustomTestCase):
def test_masks(self):
np.testing.assert_array_equal(self.masks.tolist(), EXPECTED_PROB_MASKS)
def test_output_shapes(self):
n_queries = len(QUERY_SEQUENCES)
self.assertEqual(self.ids.shape, (n_queries, 8))
self.assertEqual(self.masks.shape, (n_queries, 8, 8))
class TestNgramCorpusReset(CustomTestCase):
"""Verify reset clears all cached state."""
def test_reset_produces_empty_results(self):
corpus = _make_corpus("BFS")
corpus.batch_put(SEED_SEQUENCES)
corpus.synchronize()
ids_before, _ = _batch_get(corpus, [[1, 2, 3]])
self.assertTrue(
any(t != 0 for t in ids_before.tolist()[1:]),
"Expected non-trivial draft tokens before reset",
)
corpus.reset()
ids_after, _ = _batch_get(corpus, [[1, 2, 3]])
self.assertEqual(
ids_after.tolist(),
[3, 0, 0, 0, 0, 0, 0, 0],
"After reset, only last_token should be present (rest zero-padded)",
)
class TestNgramCorpusNoMatch(CustomTestCase):
"""Verify behavior when query has no match in the corpus."""
@@ -277,15 +242,6 @@ class TestNgramCorpusMultipleInserts(CustomTestCase):
class TestNgramCorpusSqueeze(CustomTestCase):
"""Verify cache eviction under memory pressure."""
def test_small_capacity_does_not_crash(self):
corpus = _make_corpus("BFS", capacity=200)
long_seq = list(range(1, 101))
corpus.batch_put([long_seq])
corpus.synchronize()
ids, masks = _batch_get(corpus, [[50, 51, 52]])
self.assertEqual(len(ids), 8, "Should still produce draft_token_num outputs")
def test_eviction_preserves_recent(self):
corpus = _make_corpus("BFS", capacity=500, max_trie_depth=6)
@@ -366,34 +322,6 @@ class TestNgramCorpusBatchConsistency(CustomTestCase):
)
class TestMaskValidity(CustomTestCase):
"""Verify structural invariants of the output mask for any draft tree."""
def _check_mask(self, masks_2d):
n = len(masks_2d)
for i in range(n):
self.assertEqual(masks_2d[i][i], 1, f"Diagonal must be 1 at row {i}")
self.assertEqual(masks_2d[0], [1] + [0] * (n - 1))
def test_bfs_mask_invariants(self):
corpus = _make_corpus("BFS")
corpus.batch_put(SEED_SEQUENCES)
corpus.synchronize()
_, masks = _batch_get(corpus, QUERY_SEQUENCES)
masks = masks.reshape(-1, 8, 8)
for i in range(masks.shape[0]):
self._check_mask(masks[i].tolist())
def test_prob_mask_invariants(self):
corpus = _make_corpus("PROB")
corpus.batch_put(SEED_SEQUENCES)
corpus.synchronize()
_, masks = _batch_get(corpus, QUERY_SEQUENCES)
masks = masks.reshape(-1, 8, 8)
for i in range(masks.shape[0]):
self._check_mask(masks[i].tolist())
class TestFrequencyBoosting(CustomTestCase):
"""Verify that repeated insertions change Prob-mode selection."""
@@ -509,64 +437,6 @@ class TestLongContext(CustomTestCase):
)
class TestDraftBudgetSaturation(CustomTestCase):
"""Verify the draft tree uses exactly draft_token_num slots."""
def test_full_budget_used(self):
corpus = _make_corpus("BFS", draft_token_num=8)
seq = list(range(1, 30))
corpus.batch_put([seq])
corpus.synchronize()
ids, _ = _batch_get(corpus, [[1, 2, 3]])
ids_list = ids.tolist()
self.assertEqual(len(ids_list), 8)
non_zero = [t for t in ids_list[1:] if t != 0]
self.assertGreater(
len(non_zero),
0,
"Draft budget should have non-zero tokens when cache has long chains",
)
class TestTruncate(CustomTestCase):
"""Verify truncation logic on batch_get output."""
def test_truncate_reduces_output(self):
corpus = _make_corpus("BFS", draft_token_num=8)
corpus.batch_put(SEED_SEQUENCES)
corpus.synchronize()
ids, _ = _batch_get(corpus, [[1, 2, 3]])
ids = ids.reshape(8)
self.assertEqual(len(ids), 8)
# Simulate truncate to 4
trunc_n = 4
trunc_ids = ids[:trunc_n]
self.assertEqual(len(trunc_ids), trunc_n)
def test_truncate_preserves_mask_structure(self):
corpus = _make_corpus("BFS", draft_token_num=8)
corpus.batch_put(SEED_SEQUENCES)
corpus.synchronize()
_, masks = _batch_get(corpus, [[1, 2, 3]])
n = 8
full_mask = masks.reshape(n, n)
trunc_n = 4
trunc_mask = full_mask[:trunc_n, :trunc_n]
for i in range(trunc_n):
for j in range(trunc_n):
self.assertEqual(
trunc_mask[i, j],
full_mask[i, j],
f"Mask mismatch at ({i},{j})",
)
class TestResetAndReinsert(CustomTestCase):
"""Verify that reset followed by new inserts works correctly."""
@@ -754,19 +624,6 @@ class TestNgramCorpusExternalSam(CustomTestCase):
iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=4),
)
def test_external_sam_only_chain(self):
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_documents=[[1, 2, 3, 4, 5]],
)
ids, masks = _batch_get(corpus, [[1, 2, 3]])
ids_list = ids.tolist()
self.assertEqual(ids_list[0], 3)
self.assertEqual(ids_list[1:3], [4, 5])
def test_external_sam_respects_document_boundaries(self):
corpus = _make_corpus(
"BFS",
@@ -797,23 +654,6 @@ class TestNgramCorpusExternalSam(CustomTestCase):
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 20, 21], leaf_paths)
def test_shared_prefix_keeps_both_branches(self):
corpus = _make_corpus(
"BFS",
draft_token_num=5,
external_sam_budget=2,
external_corpus_documents=[[1, 2, 3, 10, 99]],
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
ids, masks = _batch_get(corpus, [[1, 2, 3]])
leaf_paths = corpus.leaf_paths_from_mask(
ids.tolist(), masks.reshape(5, 5).tolist()
)
self.assertIn([3, 10, 11], leaf_paths)
self.assertIn([3, 10, 99], leaf_paths)
def test_shared_prefix_merge_can_underfill_budget(self):
corpus = _make_corpus(
"BFS",
@@ -930,17 +770,6 @@ class TestNgramCorpusMultiSam(CustomTestCase):
self.assertEqual(token_counts["a"], 5)
self.assertEqual(token_counts["b"], 5)
def test_remove(self):
corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
loaded_token_count = corpus.load_external_corpus_named("a", [[1, 2, 3, 4, 5]])
corpus.commit_external_corpus_load("a", loaded_token_count)
loaded_token_count = corpus.load_external_corpus_named(
"b", [[10, 20, 30, 40, 50]]
)
corpus.commit_external_corpus_load("b", loaded_token_count)
corpus.remove_external_corpus("a")
self.assertEqual(list(corpus.list_external_corpora().keys()), ["b"])
def test_remove_nonexistent_is_noop(self):
corpus = _make_corpus("BFS", draft_token_num=4, external_sam_budget=3)
corpus.remove_external_corpus("nonexistent")
@@ -977,17 +806,6 @@ class TestNgramCorpusMultiSam(CustomTestCase):
self.assertIn([3, 10, 11], leaf_paths)
self.assertNotIn([3, 20, 21], leaf_paths)
def test_make_corpus_with_documents(self):
"""_make_corpus helper loads documents as a named corpus."""
corpus = _make_corpus(
"BFS",
draft_token_num=4,
external_sam_budget=3,
external_corpus_documents=[[1, 2, 3, 4, 5]],
)
token_counts = corpus.list_external_corpora()
self.assertIn("test_corpus", token_counts)
def test_remove_frees_token_budget(self):
"""Removing a corpus should free its tokens from the total budget."""
corpus = _make_corpus(
test/registered/unit/test_model_overrides.py
-55
View File
@@ -39,9 +39,6 @@ class _FakeArgs:
class TestModelOverridableWhitelist(CustomTestCase):
def test_arg_defaults_to_not_overridable(self):
self.assertFalse(Arg().resolvable)
def test_whitelist_derivation_from_annotated_metadata(self):
self.assertEqual(
resolvable_fields(_FakeArgs),
@@ -91,9 +88,6 @@ class TestModelOverridableWhitelist(CustomTestCase):
),
)
def test_non_dataclass_yields_empty_whitelist(self):
self.assertEqual(resolvable_fields(SimpleNamespace), frozenset())
class _IsolatedRegistry(CustomTestCase):
"""Run each test against empty registries (they are process-global)."""
@@ -244,11 +238,6 @@ class _IsolatedPublish(CustomTestCase):
super().tearDown()
@dataclasses.dataclass
class _NoOverridableArgs:
x: int = 1
class TestPublishInstallsSlot(_IsolatedPublish):
"""Publish wiring: set_server_args installs the already-resolved object
into the context-owned slot (no transformation at publish time)."""
@@ -265,12 +254,6 @@ class TestPublishInstallsSlot(_IsolatedPublish):
set_global_server_args_for_scheduler(sa)
self.assertIs(get_server_args(), sa)
def test_empty_stash_publish_runs_gate_as_noop(self):
sa = _NoOverridableArgs()
sa._resolved_overrides = []
get_context().set_server_args(sa)
self.assertIs(get_server_args(), sa)
class TestGoldenModelOverrides(_IsolatedPublish):
"""Per-arch golden diff for migrated families: the declarative path must
@@ -327,11 +310,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
)
self.assertEqual(self._publish(sa).dtype, "bfloat16")
def test_pixtral_forces_bfloat16(self):
sa = self._construct("PixtralForConditionalGeneration", "pixtral")
self.assertEqual(sa.dtype, "bfloat16") # materialized
self.assertEqual(self._publish(sa).dtype, "bfloat16")
def test_user_requested_dtype_is_still_overridden(self):
# Legacy fidelity: the arch branch overwrote dtype unconditionally,
# so the declaration must too. The pristine request survives on
@@ -728,30 +706,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
)
self.assertEqual(_dllm_page_size(_view(dllm_algorithm=None)), {})
def test_declaration_overlay_mechanics(self):
from sglang.srt.arg_groups.overrides import run_post_process_pass
live = SimpleNamespace(x="user", y=None, _resolved_overrides=[])
def _resolve_x(view):
return {"x": "resolved"} if view.x == "user" else {}
def _read_x(view):
return {"y": view.x}
run_post_process_pass(live, _resolve_x)
# declaration recorded, but server_args stays pristine
self.assertEqual(
live._resolved_overrides, [(_resolve_x.__qualname__, {"x": "resolved"})]
)
self.assertEqual(live.x, "user")
# a later pass sees the resolved value through the view overlay
run_post_process_pass(live, _read_x)
self.assertEqual(
live._resolved_overrides[-1], (_read_x.__qualname__, {"y": "resolved"})
)
self.assertIsNone(live.y) # never applied in place
def test_overlap_disable_passes(self):
from sglang.srt.arg_groups.overrides import (
ResolvedView,
@@ -1479,15 +1433,6 @@ class TestGoldenModelOverrides(_IsolatedPublish):
)
)
def test_page_size_leaf_materializes_end_state(self):
sa = self._construct("LlamaForCausalLM", "llama")
declared_values = [
d["page_size"] for _s, d in sa._resolved_overrides if "page_size" in d
]
self.assertTrue(declared_values) # default fill declared
self.assertEqual(sa.page_size, declared_values[-1]) # materialized
self.assertEqual(self._publish(sa).page_size, declared_values[-1])
def test_qwen3_5_hybrid_coupled_declaration(self):
from sglang.srt.arg_groups.overrides import _qwen3_5_hybrid_overrides