diff --git a/python/sglang/srt/layers/quantization/quark/utils.py b/python/sglang/srt/layers/quantization/quark/utils.py index dcd3aeb59..6a22495ba 100644 --- a/python/sglang/srt/layers/quantization/quark/utils.py +++ b/python/sglang/srt/layers/quantization/quark/utils.py @@ -55,6 +55,19 @@ def should_ignore_layer( # proj_name = qkv_proj proj_name = layer_name.split(".")[-1] + # a fused module can be excluded under its fused name, so match it before expanding + if check_equal_or_regex_match(layer_name=layer_name, targets=ignore): + return True + + # excludes may name experts individually, so an excluded expert excludes the module + if layer_name.endswith(".experts"): + expert_prefix = layer_name + "." + if any( + isinstance(target, str) and target.startswith(expert_prefix) + for target in ignore + ): + return True + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that @@ -87,12 +100,9 @@ def should_ignore_layer( "requires all to use the same scheme." ) - # Unfused layers like down_proj and o_proj will match - # the safetensors checkpoint already. + # an unfused name was already tried by the direct check above else: - should_ignore_layer = check_equal_or_regex_match( - layer_name=layer_name, targets=ignore - ) + should_ignore_layer = False assert should_ignore_layer is not None diff --git a/test/registered/unit/layers/quantization/test_quark_utils.py b/test/registered/unit/layers/quantization/test_quark_utils.py index 5ed8240b2..b8dab0d4c 100644 --- a/test/registered/unit/layers/quantization/test_quark_utils.py +++ b/test/registered/unit/layers/quantization/test_quark_utils.py @@ -115,5 +115,46 @@ class TestE8M0ToF32(CustomTestCase): self.assertTrue(torch.isnan(out[2]).item()) +QKV_MAPPING = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + + +class TestShouldIgnoreLayerFusedNames(CustomTestCase): + """An `exclude` entry naming an already-fused module, or naming experts + individually, must exclude the fused module SGLang builds; otherwise an + MXFP4-packed parameter is allocated for a BF16 tensor and loading aborts.""" + + # ---- Bug-catchers: must FAIL on unfixed code --------------------------- + + def test_directly_excluded_fused_qkv_is_ignored(self): + name = "visual.blocks.0.attn.qkv_proj" + self.assertTrue( + should_ignore_layer(name, ignore=[name], fused_mapping=QKV_MAPPING) + ) + + def test_per_expert_excludes_ignore_the_fused_moe_module(self): + layer = "model.layers.6.mlp.experts" + ignore = [ + f"{layer}.{i}.{proj}" + for i in range(3) + for proj in ("down_proj", "gate_proj", "up_proj") + ] + self.assertTrue( + should_ignore_layer(layer, ignore=ignore, fused_mapping=QKV_MAPPING) + ) + + # ---- Guards: behavior that must NOT change ----------------------------- + + def test_unrelated_moe_layer_is_not_ignored(self): + # a prefix match must not bleed into a neighboring layer index + ignore = ["model.layers.6.mlp.experts.0.down_proj"] + self.assertFalse( + should_ignore_layer( + "model.layers.7.mlp.experts", + ignore=ignore, + fused_mapping=QKV_MAPPING, + ) + ) + + if __name__ == "__main__": unittest.main()