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

[Model] Support Ling-3.0-flash (BailingMoeV3) (#33561)

Browse Source 
Signed-off-by: JustinTong <justintong0323@gmail.com>
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: 得泽 <zhangkaihong.zkh@antgroup.com>
Co-authored-by: 翎悦 <vito.yy@antgroup.com>
Co-authored-by: 羽癫 <yudian.zy@antgroup.com>
Co-authored-by: tiwei.btw <tiwei.btw@antgroup.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: 文赋 <zibin.zb@antgroup.com>
Co-authored-by: JustinTong <justintong0323@gmail.com>
This commit is contained in:
Xinyuan Tong
2026-08-26 17:27:23 -07:00
committed by GitHub
co-authored by luoyuan.luo 得泽 翎悦 羽癫 tiwei.btw Liangsheng Yin 文赋 JustinTong
parent 8739d56a31
commit 20621aa14b
76 changed files with 5184 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_kda_kernels.py
+66 -6
View File
@@ -409,7 +409,18 @@ class TestKDAPackedDecode(unittest.TestCase):
@staticmethod
def _run_baseline(
mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, H, HV, K, V
mixed_qkv,
a,
b,
A_log,
dt_bias,
ssm_states,
cache_indices,
H,
HV,
K,
V,
lower_bound=None,
):
B = mixed_qkv.shape[0]
q_flat, k_flat, v_flat = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1)
@@ -436,11 +447,22 @@ class TestKDAPackedDecode(unittest.TestCase):
scale=K**-0.5,
use_qk_l2norm_in_kernel=True,
is_kda=True,
lower_bound=lower_bound,
)
@staticmethod
def _run_packed(
mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, HV, K, V
mixed_qkv,
a,
b,
A_log,
dt_bias,
ssm_states,
cache_indices,
HV,
K,
V,
lower_bound=None,
):
B = mixed_qkv.shape[0]
out = mixed_qkv.new_empty(B, 1, HV, V)
@@ -455,10 +477,11 @@ class TestKDAPackedDecode(unittest.TestCase):
out=out,
ssm_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
lower_bound=lower_bound,
)
return out.transpose(0, 1)
def _check(self, B, H, HV, K, V):
def _check(self, B, H, HV, K, V, lower_bound=None):
device = get_device()
dtype = torch.bfloat16
pool_size = B + 4
@@ -469,10 +492,31 @@ class TestKDAPackedDecode(unittest.TestCase):
s_baseline = ssm_states.clone()
o_packed = self._run_packed(
mixed_qkv, a, b, A_log, dt_bias, s_packed, cache_indices, HV, K, V
mixed_qkv,
a,
b,
A_log,
dt_bias,
s_packed,
cache_indices,
HV,
K,
V,
lower_bound=lower_bound,
)
o_baseline = self._run_baseline(
mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V
mixed_qkv,
a,
b,
A_log,
dt_bias,
s_baseline,
cache_indices,
H,
HV,
K,
V,
lower_bound=lower_bound,
)
torch.testing.assert_close(
@@ -501,6 +545,9 @@ class TestKDAPackedDecode(unittest.TestCase):
# Common KDA config with HV > H (grouped query).
self._check(B=8, H=8, HV=16, K=128, V=128)
def test_safe_gate_lower_bound(self):
self._check(B=8, H=16, HV=16, K=128, V=128, lower_bound=-5.0)
def test_pad_slot(self):
"""Entries with state_idx == -1 must produce zero output and skip state writeback."""
device = get_device()
@@ -544,6 +591,7 @@ class TestKDAPackedDecode(unittest.TestCase):
device = get_device()
dtype = torch.bfloat16
B, H, HV, K, V = 4, 16, 16, 128, 128
lower_bound = -5.0
pool_size = B + 4
mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices = self._make_inputs(
B, H, HV, K, V, pool_size, dtype, device
@@ -568,11 +616,23 @@ class TestKDAPackedDecode(unittest.TestCase):
cache_indices=cache_indices,
num_v_heads=HV,
head_v_dim=V,
lower_bound=lower_bound,
)
s_baseline = ssm_states.clone()
o_baseline = self._run_baseline(
mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V
mixed_qkv,
a,
b,
A_log,
dt_bias,
s_baseline,
cache_indices,
H,
HV,
K,
V,
lower_bound=lower_bound,
)
# Dispatcher returns [1, B, HV, V], same layout as the baseline.
test/registered/attention/test_kda_prefill_flashkda.py
+5 -4
View File
@@ -80,7 +80,7 @@ def _chunk_kda_ref(d, lower_bound):
"""Triton chunk_kda reference. chunk_kda mutates g/v and the state in place,
so feed clones; returns (output, updated_state_slots)."""
st = d["pool"].clone()
out = chunk_kda(
out, _ = chunk_kda(
q=d["q"].clone(),
k=d["k"].clone(),
v=d["v"].clone(),
@@ -105,7 +105,7 @@ def test_flashkda_matches_triton_safe_gate(seq_lens):
ref_out, ref_state = _chunk_kda_ref(d, LOWER_BOUND)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
out, h = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
@@ -121,6 +121,7 @@ def test_flashkda_matches_triton_safe_gate(seq_lens):
)
torch.cuda.synchronize()
assert h is None
assert torch.isfinite(out).all(), "FlashKDA output has non-finite values"
assert torch.isfinite(st_fk).all(), "FlashKDA final state has non-finite values"
# bf16 cross-implementation noise (chunk=16 CUTLASS vs chunk=64 Triton);
@@ -140,7 +141,7 @@ def test_flashkda_falls_back_without_lower_bound():
ref_out, _ = _chunk_kda_ref(d, None)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
out, _ = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
@@ -170,7 +171,7 @@ def test_flashkda_spec_verify_falls_back():
ref_out, _ = _chunk_kda_ref(d, LOWER_BOUND)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
out, _ = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
test/registered/kernels/test_fused_kda_conv_recurrent_verify.py
+187
View File
@@ -0,0 +1,187 @@
import sys
import pytest
import torch
from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import (
fused_kda_conv_gating_verify,
)
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
from sglang.kernels.ops.mamba.causal_conv1d_triton import (
causal_conv1d_update,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
_DEVICE = "cuda"
_CASES = [
(1, 4, 4, 4, 128, 128, 4, False, None, False, 1),
(1, 4, 4, 4, 128, 128, 4, True, None, False, 2),
(1, 4, 4, 4, 128, 128, 4, True, 2.0, False, 3),
(3, 4, 4, 4, 128, 128, 4, True, None, False, 4),
(3, 4, 4, 4, 128, 128, 4, True, None, True, 5),
(2, 3, 4, 4, 128, 128, 4, True, None, False, 6),
(2, 8, 2, 2, 128, 128, 4, True, 1.5, False, 7),
(1, 4, 8, 8, 64, 64, 4, True, None, False, 8),
]
def _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed):
torch.manual_seed(seed)
dim = 2 * H * K + HV * V
seq_len = B * T
lines = slots = 8
inputs = {
"mixed": torch.randn(seq_len, dim, device=_DEVICE, dtype=torch.bfloat16) * 0.5,
"w": torch.randn(dim, W, device=_DEVICE, dtype=torch.bfloat16) * 0.3,
"bias": (
torch.randn(dim, device=_DEVICE, dtype=torch.bfloat16) * 0.1
if has_bias
else None
),
"a": torch.randn(seq_len, HV * K, device=_DEVICE, dtype=torch.bfloat16) * 0.5,
"b": torch.randn(seq_len, HV, device=_DEVICE, dtype=torch.bfloat16),
"A_log": torch.randn(HV, device=_DEVICE, dtype=torch.float32) * 0.5,
"dt_bias": torch.randn(HV * K, device=_DEVICE, dtype=torch.float32) * 0.5,
# Pool layouts mirroring MambaPool: conv [lines, state_len, dim] (then
# transposed), ssm [slots, HV, V, K] fp32, window [lines, T, W-1, dim],
# intermediate ssm cache [lines, T, HV, V, K] fp32.
"conv_pool": torch.randn(
lines, W - 1, dim, device=_DEVICE, dtype=torch.bfloat16
),
"ssm": torch.randn(slots, HV, V, K, device=_DEVICE, dtype=torch.float32) * 0.2,
"win_pool": torch.zeros(
lines, T, W - 1, dim, device=_DEVICE, dtype=torch.bfloat16
),
"inter_ssm": torch.zeros(
lines, T, HV, V, K, device=_DEVICE, dtype=torch.float32
),
}
idx_vals = list(range(2, 2 + B))
if neg_slot and B >= 2:
idx_vals[1] = -1
inputs["idx_vals"] = idx_vals
inputs["cache_indices"] = torch.tensor(idx_vals, device=_DEVICE, dtype=torch.int32)
inputs["inter_indices"] = torch.arange(B, device=_DEVICE, dtype=torch.int32)
return inputs
def _run_reference(inp, B, T, H, HV, K, V, lower_bound):
dim = 2 * H * K + HV * V
seq_len = B * T
conv = inp["conv_pool"].clone()
ssm = inp["ssm"].clone()
win = inp["win_pool"].clone()
ic = inp["inter_ssm"].clone()
x3 = inp["mixed"].reshape(B, T, dim).transpose(1, 2)
out3 = causal_conv1d_update(
x3,
conv.transpose(-1, -2),
inp["w"],
inp["bias"],
activation="silu",
conv_state_indices=inp["cache_indices"],
intermediate_conv_window=win.transpose(-1, -2),
intermediate_state_indices=inp["inter_indices"],
)
mixed_out = out3.transpose(1, 2).reshape(seq_len, dim)
q, k, v = mixed_out.split([H * K, H * K, HV * V], dim=-1)
q = q.unflatten(-1, (H, K)).unsqueeze(0)
k = k.unflatten(-1, (H, K)).unsqueeze(0)
v = v.unflatten(-1, (HV, V)).unsqueeze(0)
cu = torch.arange(0, B + 1, device=_DEVICE, dtype=torch.int32) * T
o = fused_sigmoid_gating_delta_rule_update(
A_log=inp["A_log"],
a=inp["a"],
dt_bias=inp["dt_bias"],
softplus_beta=1.0,
softplus_threshold=20.0,
q=q,
k=k,
v=v,
b=inp["b"],
initial_state_source=ssm,
initial_state_indices=inp["cache_indices"],
use_qk_l2norm_in_kernel=True,
cu_seqlens=cu,
is_kda=True,
disable_state_update=True,
intermediate_states_buffer=ic,
intermediate_state_indices=inp["inter_indices"],
cache_steps=T,
retrieve_parent_token=None,
lower_bound=lower_bound,
)
return o, conv, win, ic
def _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps):
conv = inp["conv_pool"].clone()
ssm = inp["ssm"].clone()
win = inp["win_pool"].clone()
ic = inp["inter_ssm"].clone()
o = fused_kda_conv_gating_verify(
mixed_qkv=inp["mixed"],
conv_weight=inp["w"],
conv_bias=inp["bias"],
conv_state=conv.transpose(-1, -2),
conv_state_indices=inp["cache_indices"],
intermediate_conv_window=win.transpose(-1, -2),
intermediate_state_indices=inp["inter_indices"],
a=inp["a"],
b=inp["b"],
A_log=inp["A_log"],
dt_bias=inp["dt_bias"],
ssm_states=ssm,
cache_indices=inp["cache_indices"],
intermediate_states_buffer=ic,
scale=K**-0.5,
T=T,
num_q_heads=H,
num_v_heads=HV,
head_k_dim=K,
head_v_dim=V,
lower_bound=lower_bound,
num_warps=num_warps,
)
return o, conv, win, ic
def _compare_case(case, num_warps):
B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case
inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed)
o_ref, conv_ref, win_ref, ic_ref = _run_reference(
inp, B, T, H, HV, K, V, lower_bound
)
o_fus, conv_fus, win_fus, ic_fus = _run_fused(
inp, B, T, H, HV, K, V, lower_bound, num_warps
)
idx_vals = inp["idx_vals"]
valid_rows = [i for i, slot in enumerate(idx_vals) if slot >= 0]
touched_slots = [slot for slot in idx_vals if slot >= 0]
o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows]
o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows]
assert torch.equal(o_ref_v, o_fus_v)
assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots])
assert torch.equal(win_ref[valid_rows], win_fus[valid_rows])
torch.testing.assert_close(
ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0
)
@pytest.mark.parametrize("case", _CASES)
def test_matches_unfused_reference(case):
_compare_case(case, num_warps=4)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/unit/configs/test_model_config_shapes.py
+62 -3
View File
@@ -1,9 +1,12 @@
"""Unit tests for ModelConfig shape normalization."""
import math
import unittest
from types import SimpleNamespace
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.configs.model_config import (
AttentionArch,
ModelConfig,
_quant_config_to_dict,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -66,6 +69,62 @@ class TestModelConfigShapes(CustomTestCase):
self.assertEqual(model_config.swa_head_dim, 64)
self.assertEqual(model_config.swa_v_head_dim, 48)
def test_ling_mla_nope_shapes(self):
text_config = _make_text_config(
architectures=["BailingMoeV3ForCausalLM"],
kv_lora_rank=512,
qk_nope_head_dim=128,
qk_rope_head_dim=64,
use_mla_nope=True,
v_head_dim=128,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.attention_arch, AttentionArch.MLA)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.qk_rope_head_dim, 0)
self.assertEqual(model_config.scaling, 1 / math.sqrt(128))
def test_ling_mla_rope_shapes(self):
text_config = _make_text_config(
architectures=["BailingMoeV3ForCausalLM"],
kv_lora_rank=512,
qk_nope_head_dim=128,
qk_rope_head_dim=64,
use_mla_nope=False,
v_head_dim=128,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.attention_arch, AttentionArch.MLA)
self.assertEqual(model_config.head_dim, 128)
self.assertEqual(model_config.qk_rope_head_dim, 64)
self.assertEqual(model_config.scaling, 1 / math.sqrt(192))
def test_sarvam_mla_shapes(self):
text_config = _make_text_config(
architectures=["SarvamMLAForCausalLM"],
kv_lora_rank=512,
qk_nope_head_dim=128,
qk_rope_head_dim=64,
rope_scaling=None,
v_head_dim=128,
)
model_config = self._derive_shapes(text_config)
self.assertEqual(model_config.attention_arch, AttentionArch.MLA)
self.assertEqual(model_config.head_dim, 192)
self.assertEqual(model_config.qk_rope_head_dim, 64)
self.assertEqual(model_config.scaling, 1 / math.sqrt(192))
def test_quant_config_objects_are_normalized(self):
quant_config = SimpleNamespace(to_dict=lambda: {"quant_method": "test"})
self.assertEqual(_quant_config_to_dict(quant_config), {"quant_method": "test"})
if __name__ == "__main__":
unittest.main()
test/registered/unit/entrypoints/openai/test_serving_chat.py
+23
View File
@@ -3236,6 +3236,29 @@ class ServingChatTestCase(unittest.TestCase):
)
self.assertTrue(self.chat._get_reasoning_from_request(req_enabled))
def test_fallback_ling3_default_on(self):
"""Ling3 public checkpoints default `thinking_option='on'` in the chat
template when `enable_thinking` is omitted, and the template detector
cannot infer that indirect assignment. The parser fallback must mirror
the template default: omitted kwargs enable reasoning, only an explicit
`enable_thinking=False` disables it. Regression: the detector shipped
with `explicit_enable_thinking`, which left `reasoning_content` null on
default requests while the model was in fact thinking."""
self._setup_fallback("ling3")
req = ChatCompletionRequest(
model="x", messages=[{"role": "user", "content": "hi"}]
)
cases = [
(None, True), # no chat_template_kwargs → thinking (template default)
({}, True), # empty kwargs → thinking
({"enable_thinking": True}, True), # explicit on
({"enable_thinking": False}, False), # explicit off
]
for kwargs, expected in cases:
with self.subTest(kwargs=kwargs):
req.chat_template_kwargs = kwargs
self.assertEqual(self.chat._get_reasoning_from_request(req), expected)
def test_fallback_no_detector_returns_false(self):
self.chat.reasoning_parser = "qwen3"
self.chat._reasoning_detector = None
test/registered/unit/function_call/test_function_call_parser.py
+192
View File
@@ -28,6 +28,7 @@ from sglang.srt.function_call.inkling_detector import InklingDetector
from sglang.srt.function_call.json_array_parser import JsonArrayParser
from sglang.srt.function_call.kimik2_detector import KimiK2Detector
from sglang.srt.function_call.lfm2_detector import Lfm2Detector
from sglang.srt.function_call.ling3_detector import Ling3Detector
from sglang.srt.function_call.llama32_detector import Llama32Detector
from sglang.srt.function_call.mistral_detector import MistralDetector
from sglang.srt.function_call.pythonic_detector import PythonicDetector
@@ -2944,6 +2945,84 @@ class TestGlm4MoeDetector(unittest.TestCase):
)
self.assertEqual(result.normal_text, "")
def test_streaming_tool_call(self):
chunks = [
"<tool_call>get_weather\n",
"<arg_key>city</arg_key>\n<arg_value>Beijing</arg_value>\n",
"<arg_key>date</arg_key>\n<arg_value>2024-06-27</arg_value>\n",
"</tool_call>",
]
tool_calls = []
for chunk in chunks:
result = self.detector.parse_streaming_increment(chunk, self.tools)
for tool_call_chunk in result.calls:
if (
hasattr(tool_call_chunk, "tool_index")
and tool_call_chunk.tool_index is not None
):
while len(tool_calls) <= tool_call_chunk.tool_index:
tool_calls.append({"name": "", "parameters": ""})
tc = tool_calls[tool_call_chunk.tool_index]
if tool_call_chunk.name:
tc["name"] = tool_call_chunk.name
if tool_call_chunk.parameters:
tc["parameters"] += tool_call_chunk.parameters
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(
tool_calls[0]["parameters"], '{"city": "Beijing", "date": "2024-06-27"}'
)
def test_streaming_tool_call_without_arguments(self):
chunks = [
"<tool_call>get_weather\n",
"</tool_call>",
]
tool_calls = []
for chunk in chunks:
result = self.detector.parse_streaming_increment(chunk, self.tools)
for tool_call_chunk in result.calls:
if (
hasattr(tool_call_chunk, "tool_index")
and tool_call_chunk.tool_index is not None
):
while len(tool_calls) <= tool_call_chunk.tool_index:
tool_calls.append({"name": "", "parameters": ""})
tc = tool_calls[tool_call_chunk.tool_index]
if tool_call_chunk.name:
tc["name"] = tool_call_chunk.name
if tool_call_chunk.parameters:
tc["parameters"] += tool_call_chunk.parameters
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(tool_calls[0]["parameters"], "{}")
self.assertEqual(self.detector.streamed_args_for_tool[0], "{}")
def test_streaming_tool_call_without_arguments_single_chunk(self):
"""Test no-argument tool call when name and end token arrive together."""
chunks = ["<tool_call>get_weather\n</tool_call>"]
tool_calls = []
for chunk in chunks:
result = self.detector.parse_streaming_increment(chunk, self.tools)
for tool_call_chunk in result.calls:
if (
hasattr(tool_call_chunk, "tool_index")
and tool_call_chunk.tool_index is not None
):
while len(tool_calls) <= tool_call_chunk.tool_index:
tool_calls.append({"name": "", "parameters": ""})
tc = tool_calls[tool_call_chunk.tool_index]
if tool_call_chunk.name:
tc["name"] = tool_call_chunk.name
if tool_call_chunk.parameters:
tc["parameters"] += tool_call_chunk.parameters
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(tool_calls[0]["parameters"], "{}")
self.assertEqual(self.detector.streamed_args_for_tool[0], "{}")
def test_streaming_multiple_tool_calls(self):
"""Test streaming incremental parsing of multiple tool calls."""
chunks = [
@@ -3602,6 +3681,119 @@ class TestGlm47MoeDetector(unittest.TestCase):
_glm47_native_structural_tag_available.cache_clear()
class TestLing3Detector(unittest.TestCase):
def setUp(self):
self.tools = [
Tool(
type="function",
function=Function(
name="get_weather",
description="Get weather information",
parameters={
"type": "object",
"properties": {
"city": {"type": "string"},
"date": {"type": "string"},
},
},
),
),
Tool(
type="function",
function=Function(
name="get_date",
description="Get current date",
parameters={"type": "object", "properties": {}},
),
),
]
self.detector = Ling3Detector()
def _collect_streaming_tool_calls(self, chunks):
tool_calls = []
for chunk in chunks:
result = self.detector.parse_streaming_increment(chunk, self.tools)
for tool_call_chunk in result.calls:
while len(tool_calls) <= tool_call_chunk.tool_index:
tool_calls.append({"name": "", "parameters": ""})
tc = tool_calls[tool_call_chunk.tool_index]
if tool_call_chunk.name:
tc["name"] = tool_call_chunk.name
if tool_call_chunk.parameters:
tc["parameters"] += tool_call_chunk.parameters
return tool_calls
def test_detect_and_parse_newline_and_compact_tool_call(self):
cases = {
"newline": (
"<tool_call>get_weather\n"
"<arg_key>city</arg_key><arg_value>Beijing</arg_value>"
"<arg_key>date</arg_key><arg_value>2024-06-27</arg_value>"
"</tool_call>",
'{"city": "Beijing", "date": "2024-06-27"}',
),
"compact": (
"<tool_call>get_weather"
"<arg_key>city</arg_key><arg_value>Shanghai</arg_value>"
"<arg_key>date</arg_key><arg_value>2024-06-28</arg_value>"
"</tool_call>",
'{"city": "Shanghai", "date": "2024-06-28"}',
),
}
for layout, (text, expected) in cases.items():
with self.subTest(layout=layout):
result = self.detector.detect_and_parse(text, self.tools)
self.assertEqual(len(result.calls), 1)
self.assertEqual(result.calls[0].name, "get_weather")
self.assertEqual(result.calls[0].parameters, expected)
def test_detect_and_parse_empty_args(self):
result = self.detector.detect_and_parse(
"<tool_call>get_date</tool_call>", self.tools
)
self.assertEqual(len(result.calls), 1)
self.assertEqual(result.calls[0].name, "get_date")
self.assertEqual(json.loads(result.calls[0].parameters), {})
def test_streaming_empty_args_emits_single_empty_object(self):
tool_calls = self._collect_streaming_tool_calls(
["<tool_call>get_date", "</tool_call>"]
)
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0]["name"], "get_date")
self.assertEqual(tool_calls[0]["parameters"], "{}")
self.assertEqual(self.detector.streamed_args_for_tool[0], "{}")
def test_streaming_newline_and_compact_tool_call(self):
cases = {
"newline": (
[
"<tool_call>get_weather\n",
"<arg_key>city</arg_key><arg_value>Beijing</arg_value>",
"<arg_key>date</arg_key><arg_value>2024-06-27</arg_value>",
"</tool_call>",
],
'{"city": "Beijing", "date": "2024-06-27"}',
),
"compact": (
[
"<tool_call>get_weather",
"<arg_key>city</arg_key><arg_value>Shanghai</arg_value>",
"<arg_key>date</arg_key><arg_value>2024-06-28</arg_value>",
"</tool_call>",
],
'{"city": "Shanghai", "date": "2024-06-28"}',
),
}
for layout, (chunks, expected) in cases.items():
with self.subTest(layout=layout):
self.setUp()
tool_calls = self._collect_streaming_tool_calls(chunks)
self.assertEqual(len(tool_calls), 1)
self.assertEqual(tool_calls[0]["name"], "get_weather")
self.assertEqual(tool_calls[0]["parameters"], expected)
class TestJsonArrayParser(unittest.TestCase):
def setUp(self):
# Create sample tools for testing
test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py
+54 -8
View File
@@ -1,13 +1,16 @@
"""Unit tests for fused shared-expert weight scaling on per-rank shared slots.
"""Unit tests for fused shared-expert weight scaling.
These tests pin the contract of ``remap_topk_for_per_rank_shared_slots`` for
the fused shared expert's topk weight on the two paths this fix covers:
These tests pin the fused shared expert's topk weight contract on three paths:
* aiter (HIP) path: routed_scaling_factor is folded into the routed weights and
the post-MoE multiply is skipped, so the shared weight must be 1.0
for a net 1.0x contribution.
* post-MoE scaling path (default): the whole MoE output is multiplied by
routed_scaling_factor afterward, so the shared weight must be 1/rsf.
* aiter (HIP) per-rank-slot path: routed_scaling_factor is folded into the
routed weights and the post-MoE multiply is skipped, so the shared weight
must be 1.0 for a net 1.0x contribution.
* post-MoE scaling per-rank-slot path (default): the whole MoE output is
multiplied by routed_scaling_factor afterward, so the shared weight must
be 1/rsf.
* standard EP path (no per-rank slots): every rank computes the fused shared
expert and the outputs are all-reduced, so the model-supplied 1/ep_size
factor must be applied to the shared weight.
"""
from sglang.test.ci.ci_register import register_cpu_ci
@@ -80,6 +83,49 @@ class TestFusedSharedExpertScaling(CustomTestCase):
shared_weight = self._run_remap(use_aiter=False)
self.assertAlmostEqual(shared_weight, 1.0 / self.ROUTED_SCALING_FACTOR)
def _run_post_process_standard_path(self, *, scaling_factor):
topk_ids = torch.tensor([[5, 40, 100, 256]], dtype=torch.int32)
topk_weights = torch.tensor([[1.0, 0.5, 0.25, 1.0]], dtype=torch.float32)
topk_config = TopKConfig(
top_k=4,
num_fused_shared_experts=1,
fused_shared_experts_scaling_factor=scaling_factor,
allow_routed_experts_capture=False,
)
router_logits = torch.zeros((1, 256), dtype=torch.float32)
with (
patch.object(topk_module, "_is_cuda", False),
patch.object(topk_module, "_is_hip", False),
patch.object(topk_module, "_use_aiter", False),
patch.object(
topk_module, "has_per_rank_fused_shared_slots", return_value=False
),
):
_out_ids, out_weights, _recorder_ids = topk_module._post_process_topk_ids(
topk_ids.clone(),
topk_weights.clone(),
topk_config,
router_logits,
layer_id=0,
)
self.assertTrue(torch.equal(out_weights[0, :-1], topk_weights[0, :-1]))
return out_weights[0, -1].item()
def test_standard_ep_path_applies_shared_scaling_factor(self):
# Regression: models pass 1/ep_size under standard EP (every rank
# computes the fused shared expert and outputs are all-reduced), but
# the standard CUDA post-process dropped the factor, so the shared
# contribution was summed ep_size times (corrupt EP4 output on
# BailingMoeV3, BF16 and FP8 alike).
shared_weight = self._run_post_process_standard_path(scaling_factor=0.25)
self.assertAlmostEqual(shared_weight, 0.25)
def test_standard_path_without_factor_keeps_shared_weight(self):
# TP mode passes no factor; the shared weight must pass through
# unscaled (guards the predicate against degrading to always-scale).
shared_weight = self._run_post_process_standard_path(scaling_factor=None)
self.assertAlmostEqual(shared_weight, 1.0)
def test_shared_expert_ids_route_to_home_rank(self):
# Sanity: the shared slot id is placed at this rank's interleaved
# position (ep_rank * num_local_experts + num_local_routed).
test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py
+132 -56
View File
@@ -1,32 +1,10 @@
"""CPU regression test for WNA16 compressed-tensors MoE with no "Linear" group.
CompressedTensorsWNA16MoE used to read ``target_scheme_map["Linear"]`` in its
constructor. That raised ``KeyError: 'Linear'`` for compressed-tensors MoE
checkpoints whose ``config_groups`` only target the expert projections through a
regex or per-layer FQN target and therefore have no group literally named
"Linear" (e.g. mixed-precision INT4/INT8 MoE quant configs). ``get_moe_scheme``
already resolves the per-layer weight scheme by matching the layer against the
config_groups targets, so it now threads that ``weight_quant`` into the scheme
constructor instead of assuming a "Linear" group.
These tests pin that contract: building a MoE compressed-tensors config with no
"Linear" group and calling ``get_moe_scheme`` must return the correct WNA16 MoE
scheme rather than raising ``KeyError``. This is pure config-parsing logic (no
weights are created and no kernels run), so it runs on CPU.
The configs mirror real Laguna-style MoE quant configs: WNA16 int4/int8, group
strategy, group_size 128, symmetric, expert projections targeted by regex or by
per-layer FQN, with attention / router layers ignored.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
from unittest import mock
import torch
from sglang.srt.layers.moe import MoeRunnerBackend
from sglang.srt.layers.quantization.compressed_tensors import compressed_tensors
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsConfig,
)
@@ -34,18 +12,13 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsWNA16MoE,
CompressedTensorsWNA16TritonMoE,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
# WNA16 MoE Marlin (default) and Triton backends are both valid resolutions for
# this config; only the "no KeyError, correct WNA16 int-N scheme" contract matters.
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
_WNA16_MOE_SCHEMES = (CompressedTensorsWNA16MoE, CompressedTensorsWNA16TritonMoE)
# Layer whose experts we resolve a scheme for. get_moe_scheme() expands this into
# ".0.gate_proj" / ".0.up_proj" / ".0.down_proj" and matches each against targets.
EXPERTS_LAYER = "model.layers.0.mlp.experts"
# Per-layer FQN targets: the three expert projections of layer 0, named
# explicitly rather than via regex. Still no "Linear" group.
PER_LAYER_EXPERT_TARGETS = [
f"{EXPERTS_LAYER}.0.gate_proj",
f"{EXPERTS_LAYER}.0.up_proj",
@@ -53,28 +26,22 @@ PER_LAYER_EXPERT_TARGETS = [
]
def _make_wna16_moe_config(targets, num_bits):
"""A WNA16 compressed-tensors MoE quant config with NO "Linear" group.
Only the expert projections are quantized, targeted via ``targets`` (regex or
per-layer FQN). Attention / router / lm_head are ignored, exactly as a real
mixed-precision MoE checkpoint would express it.
"""
def _make_wna16_moe_config(targets, num_bits, **weight_overrides):
weights = {
"num_bits": num_bits,
"type": "int",
"symmetric": True,
"strategy": "group",
"group_size": 128,
}
weights.update(weight_overrides)
return {
"quant_method": "compressed-tensors",
# pack-quantized => WNA16 (weight-only, int, no input activations).
"format": "pack-quantized",
"config_groups": {
"group_0": {
"targets": targets,
"weights": {
"num_bits": num_bits,
"type": "int",
"symmetric": True,
"strategy": "group",
"group_size": 128,
},
# Weight-only: no activation quantization.
"weights": weights,
"input_activations": None,
}
},
@@ -83,18 +50,11 @@ def _make_wna16_moe_config(targets, num_bits):
class TestWNA16MoENoLinearGroup(CustomTestCase):
"""Regression: get_moe_scheme() must not assume a "Linear" config group."""
def _assert_wna16_moe(self, config_dict, expected_bits):
quant_config = CompressedTensorsConfig.from_config(config_dict)
# Precondition that reproduces the original bug: the parsed scheme map
# has no "Linear" group, so the old target_scheme_map["Linear"] lookup
# would KeyError.
self.assertNotIn("Linear", quant_config.target_scheme_map)
layer = torch.nn.Module()
# Would raise KeyError: 'Linear' before the fix.
scheme = quant_config.get_moe_scheme(layer, layer_name=EXPERTS_LAYER)
self.assertIsInstance(scheme, _WNA16_MOE_SCHEMES)
@@ -113,6 +73,122 @@ class TestWNA16MoENoLinearGroup(CustomTestCase):
config = _make_wna16_moe_config(PER_LAYER_EXPERT_TARGETS, num_bits=4)
self._assert_wna16_moe(config, expected_bits=4)
def test_blackwell_int4_auto_uses_triton(self):
for group_size in (32, 128):
with self.subTest(group_size=group_size):
quant_config = CompressedTensorsConfig.from_config(
_make_wna16_moe_config(
["re:.*mlp.experts.*"],
num_bits=4,
group_size=group_size,
)
)
with (
mock.patch.object(
compressed_tensors,
"get_moe_runner_backend",
return_value=MoeRunnerBackend.AUTO,
),
mock.patch.object(
compressed_tensors, "is_sm100_supported", return_value=True
),
):
scheme = quant_config.get_moe_scheme(
torch.nn.Module(), layer_name=EXPERTS_LAYER
)
self.assertIsInstance(scheme, CompressedTensorsWNA16TritonMoE)
def test_blackwell_auto_rejects_unvalidated_triton_layouts(self):
cases = {
"asymmetric": {"symmetric": False},
"channel": {"strategy": "channel", "group_size": None},
"group64": {"group_size": 64},
"actorder": {"actorder": "group"},
}
for name, overrides in cases.items():
with self.subTest(name=name):
quant_config = CompressedTensorsConfig.from_config(
_make_wna16_moe_config(
["re:.*mlp.experts.*"], num_bits=4, **overrides
)
)
with (
mock.patch.object(
compressed_tensors,
"get_moe_runner_backend",
return_value=MoeRunnerBackend.AUTO,
),
mock.patch.object(
compressed_tensors, "is_sm100_supported", return_value=True
),
):
scheme = quant_config.get_moe_scheme(
torch.nn.Module(), layer_name=EXPERTS_LAYER
)
self.assertIsInstance(scheme, CompressedTensorsWNA16MoE)
self.assertNotIsInstance(scheme, CompressedTensorsWNA16TritonMoE)
def test_explicit_triton_rejects_unvalidated_layout(self):
quant_config = CompressedTensorsConfig.from_config(
_make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=4, symmetric=False)
)
with (
mock.patch.object(
compressed_tensors,
"get_moe_runner_backend",
return_value=MoeRunnerBackend.TRITON,
),
self.assertRaisesRegex(ValueError, "only supports symmetric INT4"),
):
quant_config.get_moe_scheme(torch.nn.Module(), layer_name=EXPERTS_LAYER)
def test_blackwell_explicit_marlin_is_preserved(self):
quant_config = CompressedTensorsConfig.from_config(
_make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=4)
)
with (
mock.patch.object(
compressed_tensors,
"get_moe_runner_backend",
return_value=MoeRunnerBackend.MARLIN,
),
mock.patch.object(
compressed_tensors, "is_sm100_supported", return_value=True
),
):
scheme = quant_config.get_moe_scheme(
torch.nn.Module(), layer_name=EXPERTS_LAYER
)
self.assertIsInstance(scheme, CompressedTensorsWNA16MoE)
def test_blackwell_int8_auto_keeps_marlin(self):
quant_config = CompressedTensorsConfig.from_config(
_make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=8)
)
with (
mock.patch.object(
compressed_tensors,
"get_moe_runner_backend",
return_value=MoeRunnerBackend.AUTO,
),
mock.patch.object(
compressed_tensors, "is_sm100_supported", return_value=True
),
):
scheme = quant_config.get_moe_scheme(
torch.nn.Module(), layer_name=EXPERTS_LAYER
)
self.assertIsInstance(scheme, CompressedTensorsWNA16MoE)
self.assertNotIsInstance(scheme, CompressedTensorsWNA16TritonMoE)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_mamba_checkpoint_depth.py
+3
View File
@@ -50,6 +50,9 @@ def _track_seqlen(*, tree_page: int, prefix_len: int, extend_len: int) -> int:
req.mamba_branching_seqlen = None
batch = ScheduleBatch(reqs=[req])
batch.model_config = SimpleNamespace(
hf_text_config=SimpleNamespace(mamba_chunk_size=CHUNK)
)
batch.tree_cache = SimpleNamespace(page_size=tree_page)
batch.req_to_token_pool = MagicMock()
batch.req_to_token_pool.get_mamba_ping_pong_other_idx.return_value = 1
test/registered/unit/mem_cache/test_flashkda_strided_state_access.py
+2 -1
View File
@@ -178,11 +178,12 @@ class TestFlashKDAStridedStateAccess(unittest.TestCase):
conv_before = [cv.clone() for cv in conv_views]
cache_indices = torch.tensor([5, 2], dtype=torch.int32)
out = self._run_extend(ssm_states, cache_indices)
out, intermediate_states = self._run_extend(ssm_states, cache_indices)
# Routing: the fused path ran exactly once (a silent re-route to the
# triton fallback would make every assertion below vacuous).
self.assertEqual(self.fake.calls, 1)
self.assertIsNone(intermediate_states)
self.assertEqual(tuple(out.shape), (1, 2 * _SEQ_LEN, _H, _V))
# Gather: the external kernel must receive a CONTIGUOUS copy whose rows
test/registered/unit/models/test_shared_experts_fusion_gates.py
+132 -2
View File
@@ -14,9 +14,13 @@ through `get_parallel().override(...)`; the ones that are pure config /
quantization are exercised directly.
"""
import importlib.util
import sys
import unittest
import unittest.mock
from types import SimpleNamespace
from types import ModuleType, SimpleNamespace
import pytest
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
@@ -29,6 +33,25 @@ def _quant(name: str):
return SimpleNamespace(get_name=lambda: name)
def _import_bailing_modules():
if importlib.util.find_spec("vllm") is not None:
from sglang.srt.models import bailing_moe_nextn, bailing_moe_v3
return bailing_moe_v3, bailing_moe_nextn
# CPU CI omits vLLM; these fusion gates never execute the imported AWQ kernel.
vllm = ModuleType("vllm")
vllm.__path__ = []
custom_ops = ModuleType("vllm._custom_ops")
custom_ops.awq_dequantize = unittest.mock.Mock()
with unittest.mock.patch.dict(
sys.modules, {"vllm": vllm, "vllm._custom_ops": custom_ops}
):
from sglang.srt.models import bailing_moe_nextn, bailing_moe_v3
return bailing_moe_v3, bailing_moe_nextn
class _FusionGateCase(CustomTestCase):
def _seed(self, **fields):
override = get_context().override_server_args(**fields)
@@ -228,6 +251,113 @@ class TestMiniMaxGates(_FusionGateCase):
)
class TestBailingMoeV3Gate(_FusionGateCase):
def _config(self):
return SimpleNamespace(
architectures=["BailingMoeV3ForCausalLM"],
num_shared_experts=1,
moe_intermediate_size=1024,
)
def _compressed_tensors(self, ignore):
return SimpleNamespace(
get_name=lambda: "compressed_tensors",
ignore=ignore,
packed_modules_mapping={},
)
def _reason_on_cuda(self, quant_config):
bailing_moe_v3, _ = _import_bailing_modules()
self._seed()
with (
unittest.mock.patch.object(bailing_moe_v3, "_is_cuda", True),
unittest.mock.patch.object(
bailing_moe_v3.torch.cuda,
"get_device_capability",
return_value=(9, 0),
),
):
return self._reason(
bailing_moe_v3.BailingMoeV3ForCausalLM,
self._config(),
quant_config,
)
def test_compressed_tensors_mixed_expert_layout_cannot_fuse(self):
reason = self._reason_on_cuda(
self._compressed_tensors(
["re:.*(mlp|shared_experts)\\.(gate|up|gate_up|down|eh)_proj.*"]
)
)
self.assertIn("different quant methods", reason)
def test_compressed_tensors_uniform_expert_layout_can_fuse(self):
self.assertIsNone(self._reason_on_cuda(self._compressed_tensors([])))
def test_nextn_uses_its_rewritten_architecture(self):
bailing_moe_v3, bailing_moe_nextn = _import_bailing_modules()
config = self._config()
config.architectures = ["BailingMoeForCausalLMNextN"]
config.model_type = "bailing_hybrid"
config.use_kda = True
self._seed()
with (
unittest.mock.patch.object(bailing_moe_v3, "_is_cuda", True),
unittest.mock.patch.object(
bailing_moe_v3.torch.cuda,
"get_device_capability",
return_value=(9, 0),
),
):
reason = self._reason(
bailing_moe_nextn.BailingMoeForCausalLMNextN,
config,
self._compressed_tensors(
["re:.*(mlp|shared_experts)\\.(gate|up|gate_up|down|eh)_proj.*"]
),
)
self.assertIn("different quant methods", reason)
def test_nextn_constructor_calls_v3_fusion_setup(self):
bailing_moe_v3, bailing_moe_nextn = _import_bailing_modules()
config = SimpleNamespace(
architectures=["BailingMoeForCausalLMNextN"],
model_type="bailing_hybrid",
use_kda=True,
num_shared_experts=1,
vocab_size=32000,
hidden_size=4096,
)
parallel = SimpleNamespace(
tp_size=1,
moe_ep_size=1,
enable_dp_lm_head=False,
)
with (
unittest.mock.patch.object(
bailing_moe_nextn, "get_parallel", return_value=parallel
),
unittest.mock.patch.object(
bailing_moe_v3, "get_parallel", return_value=parallel
),
unittest.mock.patch.object(
bailing_moe_v3,
"is_shared_experts_fusion_disabled",
return_value=False,
),
unittest.mock.patch.object(bailing_moe_nextn, "BailingMoEModelNextN"),
unittest.mock.patch.object(bailing_moe_nextn, "ParallelLMHead"),
unittest.mock.patch.object(bailing_moe_nextn, "LogitsProcessor"),
):
model = bailing_moe_nextn.BailingMoeForCausalLMNextN(config)
self.assertEqual(model.num_fused_shared_experts, 1)
class TestQwen3_5Gate(_FusionGateCase):
def test_every_entry_class_answers(self):
import sglang.srt.models.qwen3_5 as qwen3_5
@@ -569,4 +699,4 @@ class TestFamiliesWithoutAGate(_FusionGateCase):
if __name__ == "__main__":
unittest.main()
sys.exit(pytest.main([__file__]))
test/registered/unit/parser/test_reasoning_parser.py
+67
View File
@@ -14,6 +14,7 @@ from sglang.srt.parser.reasoning_parser import (
InklingDetector,
KimiDetector,
KimiK2Detector,
Ling3Detector,
Nemotron3Detector,
Qwen3Detector,
ReasoningParser,
@@ -466,6 +467,69 @@ class TestGlm45Detector(CustomTestCase):
self.assertEqual(result.normal_text, "<tool_call>tool call")
class TestLing3Detector(CustomTestCase):
def setUp(self):
self.detector = Ling3Detector()
def test_init(self):
self.assertEqual(self.detector.tool_start_token, "<tool_call>")
self.assertEqual(self.detector.reasoning_default, "enable_thinking")
self.assertTrue(self.detector.thinks_internally)
self.assertTrue(self.detector._force_nonempty_content)
self.assertFalse(self.detector._in_reasoning)
def test_tool_interrupt(self):
text = "<think>I need a tool<tool_call>get_weather</tool_call>"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "I need a tool")
self.assertEqual(result.normal_text, "<tool_call>get_weather</tool_call>")
def test_reasoning_only_swaps_to_normal_text(self):
text = "<think>Final answer without a closing think tag"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, "Final answer without a closing think tag")
def test_reasoning_only_with_end_token_swaps_to_normal_text(self):
text = "<think>Final answer accidentally wrapped as reasoning</think>"
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "")
self.assertEqual(
result.normal_text, "Final answer accidentally wrapped as reasoning"
)
def test_force_nonempty_content_false_disables_swap(self):
detector = Ling3Detector(force_nonempty_content=False)
text = "<think>Reasoning only</think>"
result = detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Reasoning only")
self.assertEqual(result.normal_text, "")
def test_does_not_swap_when_normal_text_exists(self):
text = "<think>Reasoning here</think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "Reasoning here")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_empty_reasoning_with_normal_text(self):
text = "<think></think>The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, "The answer is 42.")
def test_plain_text_without_thinking(self):
text = "The answer is 42."
result = self.detector.detect_and_parse(text)
self.assertEqual(result.reasoning_text, "")
self.assertEqual(result.normal_text, text)
def test_streaming_reasoning_only_currently_streams_reasoning(self):
self.detector.parse_streaming_increment("<think>")
result = self.detector.parse_streaming_increment("The answer is 42.")
self.assertEqual(result.reasoning_text, "The answer is 42.")
self.assertEqual(result.normal_text, "")
class TestHunyuanDetector(CustomTestCase):
"""Test cases for Hunyuan detector with tool interruption support."""
@@ -678,6 +742,9 @@ class TestReasoningParser(CustomTestCase):
parser = ReasoningParser("glm45")
self.assertIsInstance(parser.detector, Glm45Detector)
parser = ReasoningParser("ling3")
self.assertIsInstance(parser.detector, Ling3Detector)
parser = ReasoningParser("hunyuan")
self.assertIsInstance(parser.detector, HunyuanDetector)
test/registered/unit/test_model_overrides.py
+17
View File
@@ -2117,6 +2117,7 @@ class TestGoldenModelOverrides(_IsolatedPublish):
disable_overlap_schedule=False,
page_size=None,
linear_attn_backend="triton",
linear_attn_prefill_backend=None,
)
defaults.update(kw)
return ResolvedView(
@@ -2142,6 +2143,13 @@ class TestGoldenModelOverrides(_IsolatedPublish):
"mamba_radix_cache_strategy": "extra_buffer",
},
)
self.assertEqual(
_mamba_radix_cache_resolution(_view("BailingMoeV3ForCausalLM")),
{
"uses_mamba_radix_cache": True,
"mamba_radix_cache_strategy": "extra_buffer",
},
)
# auto + no extra-buffer support (Lfm2) -> no_buffer + overlap disable
self.assertEqual(
_mamba_radix_cache_resolution(_view("Lfm2ForCausalLM")),
@@ -2204,6 +2212,15 @@ class TestGoldenModelOverrides(_IsolatedPublish):
SimpleNamespace(linear_attn_backend="fla"), "Qwen3NextForCausalLM"
)
)
self.assertTrue(
supports_mamba_cache_extra_buffer(
SimpleNamespace(
linear_attn_backend="triton",
linear_attn_prefill_backend="flashinfer",
),
"Qwen3_5MoeForConditionalGeneration",
)
)
def test_qwen3_5_hybrid_coupled_declaration(self):
from sglang.srt.arg_groups.overrides import _qwen3_5_hybrid_overrides