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/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()