From eb75d990f7fbf6aa0f4ae77dc94246933d7a11e3 Mon Sep 17 00:00:00 2001 From: Joe Rowell Date: Wed, 1 Jul 2026 20:08:09 +0200 Subject: [PATCH] [Bugfix] compressed-tensors WNA16 MoE: don't assume a "Linear" config group (#29761) Signed-off-by: Joe Rowell Co-authored-by: Jiminator Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> --- .../compressed_tensors/compressed_tensors.py | 12 +- .../compressed_tensors_w4a4_mxint4_moe.py | 10 +- .../schemes/compressed_tensors_wNa16_moe.py | 13 +- ..._compressed_tensors_wna16_moe_no_linear.py | 118 ++++++++++++++++++ 4 files changed, 145 insertions(+), 8 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 431ffadda..0495e951f 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -682,10 +682,12 @@ class CompressedTensorsConfig(QuantizationConfig): logger.info_once( "Using CompressedTensorsMxInt4MoE with flashinfer_trtllm backend" ) - return CompressedTensorsMxInt4MoE(self) + return CompressedTensorsMxInt4MoE(self, weight_quant=weight_quant) elif _is_hip: logger.info_once("Using CompressedTensorsWNA16TritonMoE (ROCm)") - return CompressedTensorsWNA16TritonMoE(self) + return CompressedTensorsWNA16TritonMoE( + self, weight_quant=weight_quant + ) else: moe_backend = get_moe_runner_backend() if moe_backend.is_triton(): @@ -693,9 +695,11 @@ class CompressedTensorsConfig(QuantizationConfig): "Using CompressedTensorsWNA16TritonMoE " "(moe_runner_backend=triton)" ) - return CompressedTensorsWNA16TritonMoE(self) + return CompressedTensorsWNA16TritonMoE( + self, weight_quant=weight_quant + ) logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod") - return CompressedTensorsWNA16MoE(self) + return CompressedTensorsWNA16MoE(self, weight_quant=weight_quant) else: if ( self._is_dynamic_token_w4(weight_quant, input_quant) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py index 7dc430e9c..cc3624a7b 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py @@ -25,6 +25,8 @@ logger = logging.getLogger(__name__) __all__ = ["CompressedTensorsMxInt4MoE"] if TYPE_CHECKING: + from compressed_tensors.quantization import QuantizationArgs + from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, StandardDispatchOutput, @@ -46,9 +48,13 @@ if is_flashinfer_available(): class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme): - def __init__(self, quant_config: CompressedTensorsConfig): + def __init__( + self, quant_config: CompressedTensorsConfig, weight_quant: QuantizationArgs + ): self.quant_config = quant_config - config = self.quant_config.target_scheme_map["Linear"].get("weights") + # Per-layer scheme already resolved by get_moe_scheme(); reuse it directly + # (mixed-precision MoE has no "Linear" config group to fall back on). + config = weight_quant self.num_bits = config.num_bits self.packed_factor = 32 // config.num_bits self.strategy = config.strategy diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index 942e3b8f0..236420847 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -28,6 +28,8 @@ from sglang.srt.layers.quantization.utils import replace_parameter from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, set_weight_attrs if TYPE_CHECKING: + from compressed_tensors.quantization import QuantizationArgs + from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, StandardDispatchOutput, @@ -62,9 +64,16 @@ class GPTQMarlinState(Enum): class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): - def __init__(self, quant_config: CompressedTensorsConfig, num_gpu_experts=-1): + def __init__( + self, + quant_config: CompressedTensorsConfig, + weight_quant: QuantizationArgs, + num_gpu_experts: int = -1, + ): self.quant_config = quant_config - config = self.quant_config.target_scheme_map["Linear"].get("weights") + # Per-layer scheme already resolved by get_moe_scheme(); reuse it directly + # (mixed-precision MoE has no "Linear" config group to fall back on). + config = weight_quant self.num_bits = config.num_bits self.packed_factor = 32 // config.num_bits self.strategy = config.strategy diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py b/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py new file mode 100644 index 000000000..f00ad7eff --- /dev/null +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py @@ -0,0 +1,118 @@ +"""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 + +import torch + +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from sglang.srt.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsWNA16MoE, + CompressedTensorsWNA16TritonMoE, +) +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. +_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", + f"{EXPERTS_LAYER}.0.down_proj", +] + + +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. + """ + 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. + "input_activations": None, + } + }, + "ignore": ["lm_head", "re:.*self_attn.*", "re:.*mlp.gate$"], + } + + +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) + self.assertEqual(scheme.num_bits, expected_bits) + self.assertEqual(scheme.group_size, 128) + + def test_regex_expert_targets_int4(self): + config = _make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=4) + self._assert_wna16_moe(config, expected_bits=4) + + def test_regex_expert_targets_int8(self): + config = _make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=8) + self._assert_wna16_moe(config, expected_bits=8) + + def test_per_layer_fqn_expert_targets_int4(self): + config = _make_wna16_moe_config(PER_LAYER_EXPERT_TARGETS, num_bits=4) + self._assert_wna16_moe(config, expected_bits=4) + + +if __name__ == "__main__": + unittest.main()