[AMD][GLM-5.2] Keep GlmMoeDsa MoE e_score_correction_bias in fp32 (#37133)

Co-authored-by: JohnQinAMD <yanyuan.qin@amd.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Zhang, Jiejing <jiejing.zhang@amd.com>

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",

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,

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",

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