diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 43c6be020..ff3b6d340 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -151,6 +151,14 @@ def is_deepseek_dsa(config) -> bool: ) +def is_glm_moe_dsa(config) -> bool: + """True for GLM-5.2, both the main arch and the NextN draft head.""" + return _hf_arch(config) in ( + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", + ) + + def is_kimi_k3(config) -> bool: return _hf_arch(config) in ( "KimiK3ForConditionalGeneration", @@ -682,7 +690,6 @@ class ModelConfig: context_length: Optional[int] = None, **kwargs, ): - cfg = resolving_view(server_args) quantization = ( cfg.speculative_draft_model_quantization @@ -2239,7 +2246,6 @@ def is_hybrid_swa_model( model_architectures: List[str], hf_text_config: Optional[PretrainedConfig] = None, ): - hybrid_swa_archs = { "Llama4ForConditionalGeneration", "DeepseekV4ForCausalLM", diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index f079b99c9..f2f17c1ef 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1757,9 +1757,20 @@ def biased_grouped_topk_gpu( topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) + # Don't re-downcast an fp32 correction bias at the aiter boundary: an + # offset bias loses too many levels in bf16 and reorders top-k. Cast the + # gating logits up instead. Gated on the bias dtype rather than the + # architecture, so a bias that arrives as bf16 is byte-identical to + # before, as is the radix4 path above. + if correction_bias.dtype == torch.float32: + aiter_gating_output = gating_output.to(torch.float32) + aiter_bias = correction_bias + else: + aiter_gating_output = gating_output + aiter_bias = bias aiter_biased_grouped_topk( - gating_output, - bias, + aiter_gating_output, + aiter_bias, topk_weights, topk_ids, num_expert_group, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 180dff034..31ee7d675 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -49,6 +49,7 @@ from sglang.srt.configs.model_config import ( get_dsa_index_n_heads, get_dsa_index_topk, is_deepseek_dsa, + is_glm_moe_dsa, ) from sglang.srt.distributed import ( divide, @@ -473,7 +474,9 @@ class MoEGate(nn.Module): ) if config.topk_method == "noaux_tc" and not is_hash_moe: correction_bias_dtype = torch.float32 - if quant_config is not None: + # GLM-5.2's bias sits at an offset where its spread is only a few bf16 ULPs + # wide, so bf16 collapses it and reorders top-k routing. HF stores it fp32. + if quant_config is not None and not is_glm_moe_dsa(config): if _use_aiter and quant_config.get_name() in ( "fp8", "compressed_tensors", diff --git a/test/registered/unit/models/test_glmmoedsa_correction_bias_fp32.py b/test/registered/unit/models/test_glmmoedsa_correction_bias_fp32.py new file mode 100644 index 000000000..bc958e48b --- /dev/null +++ b/test/registered/unit/models/test_glmmoedsa_correction_bias_fp32.py @@ -0,0 +1,141 @@ +"""Unit tests for the GLM-5.2 (GlmMoeDsa) fp32 MoE correction-bias fix. + +GlmMoeDsa's MoE ``e_score_correction_bias`` values are ~34. bf16 has ULP 0.25 at +that magnitude, so downcasting collapses the ~174 distinct biases to ~3 levels, +which scrambles top-k expert routing (noaux_tc picks experts by sigmoid-score + +bias, and the ~34 bias dominates the [0, 1] sigmoid term). The fix keeps the bias +in fp32 for GlmMoeDsa at both the parameter-construction site (MoEGate) and the +aiter routing boundary (layers/moe/topk.py). These tests are pure dtype / CPU +logic -- no server, no weight loading, no GPU required. +""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.configs.model_config import is_glm_moe_dsa +from sglang.srt.models.deepseek_v2 import MoEGate +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +GLM_MAIN_ARCH = "GlmMoeDsaForCausalLM" +GLM_NEXTN_ARCH = "GlmMoeDsaForCausalLMNextN" # draft head, rewritten in model_config.py +NON_GLM_ARCH = "DeepseekV3ForCausalLM" + +# Biases spanning a ~0.5-wide window around 34: distinct in fp32 (fp32 ULP ~4e-6 +# there), but within ~2-3 bf16 bins (bf16 ULP 0.25 at magnitude 34). +NUM_EXPERTS = 174 +BIAS_BASE = 34.0 + + +def _glm_bias_values() -> torch.Tensor: + step = 0.5 / NUM_EXPERTS + return torch.tensor( + [BIAS_BASE + i * step for i in range(NUM_EXPERTS)], dtype=torch.float32 + ) + + +class TestMoEGateCorrectionBiasDtype(CustomTestCase): + """Behavioral guard on the production dtype block in MoEGate.__init__. + + Fails on pre-fix code (GlmMoeDsa was downcast to bf16 like every other aiter + fp8 model); passes once the arch-gate skips the downcast for GlmMoeDsa. + """ + + def _build_gate(self, arch: str) -> MoEGate: + config = SimpleNamespace( + n_routed_experts=NUM_EXPERTS, + hidden_size=16, + topk_method="noaux_tc", + architectures=[arch], + ) + quant_config = SimpleNamespace(get_name=lambda: "fp8") + # Force the aiter fp8 path (the branch that downcasts to bf16 for non-GLM); + # _is_cpu=False avoids the AMX PackWeightMethod branch so the test is + # deterministic across CPU and GPU CI runners. + import sglang.srt.models.deepseek_v2 as dv2 + + with patch.object(dv2, "_use_aiter", True), patch.object(dv2, "_is_cpu", False): + return MoEGate(config=config, quant_config=quant_config) + + def test_glm_main_keeps_fp32(self): + gate = self._build_gate(GLM_MAIN_ARCH) + self.assertEqual(gate.e_score_correction_bias.dtype, torch.float32) + + def test_glm_nextn_keeps_fp32(self): + # Guards the NextN draft head: the "GlmMoeDsa" substring must cover it too. + gate = self._build_gate(GLM_NEXTN_ARCH) + self.assertEqual(gate.e_score_correction_bias.dtype, torch.float32) + + def test_non_glm_still_downcasts_bf16(self): + # Blast-radius guard: the fix must not widen dtype for other aiter models. + # Fails if the gate is accidentally made too broad (e.g. always-skip). + gate = self._build_gate(NON_GLM_ARCH) + self.assertEqual(gate.e_score_correction_bias.dtype, torch.bfloat16) + + +class TestIsGlmMoeDsaHelper(CustomTestCase): + """The arch-gate predicate used at both fix sites.""" + + def test_matches_main_and_nextn(self): + self.assertTrue(is_glm_moe_dsa(SimpleNamespace(architectures=[GLM_MAIN_ARCH]))) + self.assertTrue(is_glm_moe_dsa(SimpleNamespace(architectures=[GLM_NEXTN_ARCH]))) + + def test_reads_the_first_architecture_like_its_neighbours(self): + # is_deepseek_dsa and is_kimi_k3 next to it both decide on + # architectures[0]; a HF config carries the model's own arch there. + self.assertFalse( + is_glm_moe_dsa(SimpleNamespace(architectures=["Foo", GLM_MAIN_ARCH])) + ) + + def test_rejects_non_glm(self): + self.assertFalse(is_glm_moe_dsa(SimpleNamespace(architectures=[NON_GLM_ARCH]))) + + def test_returns_false_for_none_or_empty_architectures(self): + # A config with architectures=None or [] must return False, never + # mis-gating a non-GLM model. _hf_arch() returns None for both. + self.assertFalse(is_glm_moe_dsa(SimpleNamespace(architectures=None))) + self.assertFalse(is_glm_moe_dsa(SimpleNamespace(architectures=[]))) + + +class TestCorrectionBiasBf16Collapse(CustomTestCase): + """Pins the numeric mechanism the fp32 fix protects against.""" + + def test_bf16_collapses_distinct_biases(self): + biases = _glm_bias_values() + fp32_distinct = torch.unique(biases).numel() + bf16_distinct = torch.unique(biases.to(torch.bfloat16)).numel() + # fp32 keeps every distinct bias; bf16 collapses the 174 values to a + # handful (the documented ~3), i.e. an order-of-magnitude information loss. + self.assertEqual(fp32_distinct, NUM_EXPERTS) + self.assertLessEqual(bf16_distinct, 4) + self.assertLess(bf16_distinct * 20, fp32_distinct) + + def test_bf16_bias_scrambles_topk_routing(self): + # noaux_tc selects top-k experts by (sigmoid(logits) + correction_bias). + # With the bias collapsed to ~3 levels, selection within a level is decided + # by the tiny sigmoid term instead of the intended bias order -> the chosen + # expert set diverges from the fp32 (correct) selection for most tokens. + torch.manual_seed(0) + topk = 8 + num_tokens = 64 + biases_fp32 = _glm_bias_values()[torch.randperm(NUM_EXPERTS)] + biases_bf16 = biases_fp32.to(torch.bfloat16).to(torch.float32) + + scores = torch.randn(num_tokens, NUM_EXPERTS).sigmoid() + top_fp32 = (scores + biases_fp32).topk(topk, dim=-1).indices + top_bf16 = (scores + biases_bf16).topk(topk, dim=-1).indices + + differ = sum( + set(top_fp32[t].tolist()) != set(top_bf16[t].tolist()) + for t in range(num_tokens) + ) + self.assertGreater(differ / num_tokens, 0.5) + + +if __name__ == "__main__": + unittest.main()