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