[AMD] GLM-5.2 NextN: cast draft fused MoE to per-channel FP8 (#39155)

Co-authored-by: fanxingran <xingran.fan@amd.com>

test/registered/unit/models/test_glm_nextn_moe_ptpc.py

@@ -0,0 +1,160 @@
+"""CI for SGLANG_GLM_NEXTN_MOE_PTPC=1 (GLM-5.2 NextN per-channel FP8 draft MoE).
+
+The feature is off by default. Without a case that turns the flag on, CI never
+touches the Quark scheme rewrite and cannot claim the path works. These tests
+exercise that ON wiring on CPU without loading a 70B MXFP4 checkpoint.
+
+A full serve+generate job still needs the MXFP4 weights in the runner cache;
+register that separately as nightly if the checkpoint is present.
+"""
+
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.models.glm4_moe import (
+    GlmMoeDsaForCausalLMNextN,
+    enable_glm_nextn_moe_ptpc,
+    should_apply_glm_nextn_moe_ptpc,
+)
+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")
+
+LAYER = 78
+PREFIX = f"model.layers.{LAYER}"
+EXPERT_LEAF = f"{PREFIX}.mlp.experts.0.w1"
+ATTN_LEAF = f"{PREFIX}.self_attn.q_proj"
+EXPERT_WEIGHT = f"{PREFIX}.mlp.experts.0.gate_proj.weight"
+_PTPC_ENV = "sglang.srt.models.glm4_moe.envs.SGLANG_GLM_NEXTN_MOE_PTPC.get"
+
+
+def _quark_cfg(*, exclude=None, layer_quant=None):
+    return SimpleNamespace(
+        get_name=lambda: "quark",
+        quant_config={"layer_quant_config": dict(layer_quant or {}), "exclude": []},
+        exclude_layers=list(
+            exclude if exclude is not None else [EXPERT_LEAF, ATTN_LEAF]
+        ),
+    )
+
+
+class TestEnableGlmNextnMoePtpc(CustomTestCase):
+    def test_off_by_default(self):
+        self.assertFalse(enable_glm_nextn_moe_ptpc(_quark_cfg()))
+
+    def test_on_requires_quark(self):
+        with patch(
+            _PTPC_ENV,
+            return_value=True,
+        ):
+            self.assertTrue(enable_glm_nextn_moe_ptpc(_quark_cfg()))
+            self.assertFalse(
+                enable_glm_nextn_moe_ptpc(SimpleNamespace(get_name=lambda: "fp8"))
+            )
+            self.assertFalse(enable_glm_nextn_moe_ptpc(None))
+
+    def test_apply_requires_excluded_mtp_experts(self):
+        with patch(
+            _PTPC_ENV,
+            return_value=True,
+        ):
+            self.assertTrue(should_apply_glm_nextn_moe_ptpc(_quark_cfg(), LAYER))
+            self.assertFalse(
+                should_apply_glm_nextn_moe_ptpc(_quark_cfg(exclude=[ATTN_LEAF]), LAYER)
+            )
+
+
+class TestResolveNextnQuantConfigPtpcOn(CustomTestCase):
+    def _resolve(self, cfg, flag: bool):
+        model = GlmMoeDsaForCausalLMNextN.__new__(GlmMoeDsaForCausalLMNextN)
+        hf = SimpleNamespace(num_hidden_layers=LAYER)
+        with patch(
+            _PTPC_ENV,
+            return_value=flag,
+        ):
+            return model._resolve_nextn_quant_config(hf, cfg)
+
+    def test_flag_off_excludes_fused_experts(self):
+        src = _quark_cfg()
+        out = self._resolve(src, flag=False)
+        self.assertIn("model.decoder.mlp.experts", out.exclude_layers)
+        self.assertNotIn(
+            "model.decoder.mlp.experts",
+            out.quant_config.get("layer_quant_config", {}),
+        )
+
+    def test_flag_on_assigns_ptpc_scheme_instead_of_bf16_exclude(self):
+        src = _quark_cfg()
+        out = self._resolve(src, flag=True)
+        self.assertNotIn("model.decoder.mlp.experts", out.exclude_layers)
+        scheme = out.quant_config["layer_quant_config"]["model.decoder.mlp.experts"]
+        self.assertEqual(scheme["weight"]["dtype"], "fp8_e4m3")
+        self.assertEqual(scheme["weight"]["qscheme"], "per_channel")
+        self.assertFalse(scheme["weight"]["is_dynamic"])
+        self.assertEqual(scheme["input_tensors"]["dtype"], "fp8_e4m3")
+        self.assertEqual(scheme["input_tensors"]["qscheme"], "per_channel")
+        self.assertTrue(scheme["input_tensors"]["is_dynamic"])
+
+    def test_flag_on_does_not_mutate_caller_config(self):
+        src = _quark_cfg()
+        orig_exclude = list(src.exclude_layers)
+        orig_layer = dict(src.quant_config.get("layer_quant_config") or {})
+        self._resolve(src, flag=True)
+        self.assertEqual(src.exclude_layers, orig_exclude)
+        self.assertEqual(src.quant_config.get("layer_quant_config") or {}, orig_layer)
+
+    def test_flag_on_skips_ptpc_when_experts_not_excluded(self):
+        src = _quark_cfg(exclude=[ATTN_LEAF])
+        out = self._resolve(src, flag=True)
+        self.assertNotIn("model.decoder.mlp.experts", out.exclude_layers)
+        self.assertNotIn(
+            "model.decoder.mlp.experts",
+            out.quant_config.get("layer_quant_config", {}),
+        )
+
+    def test_flag_on_skips_ptpc_when_mtp_not_in_exclude(self):
+        src = _quark_cfg(exclude=[])
+        out = self._resolve(src, flag=True)
+        self.assertNotIn(
+            "model.decoder.mlp.experts",
+            out.quant_config.get("layer_quant_config", {}),
+        )
+
+
+class TestMaybeQuantGlmNextnMoeToPtpc(CustomTestCase):
+    def _cast(self, cfg, flag: bool):
+        loader = GlmMoeDsaForCausalLMNextN.__new__(GlmMoeDsaForCausalLMNextN)
+        loader.quant_config = cfg
+        loader.config = SimpleNamespace(num_hidden_layers=LAYER)
+        weights = [(EXPERT_WEIGHT, torch.ones(4, 8, dtype=torch.bfloat16))]
+        with patch(
+            _PTPC_ENV,
+            return_value=flag,
+        ):
+            return list(loader._maybe_quant_glm_nextn_moe_to_ptpc(weights))
+
+    def test_flag_on_casts_excluded_bf16_experts(self):
+        out = self._cast(_quark_cfg(), flag=True)
+        names = [name for name, _ in out]
+        self.assertIn(EXPERT_WEIGHT, names)
+        self.assertIn(EXPERT_WEIGHT[: -len("weight")] + "weight_scale", names)
+        weight = dict(out)[EXPERT_WEIGHT]
+        self.assertEqual(weight.dtype, torch.float8_e4m3fn)
+
+    def test_flag_on_does_not_cast_when_experts_not_excluded(self):
+        src = _quark_cfg(exclude=[ATTN_LEAF])
+        out = self._cast(src, flag=True)
+        self.assertEqual(len(out), 1)
+        name, tensor = out[0]
+        self.assertEqual(name, EXPERT_WEIGHT)
+        self.assertEqual(tensor.dtype, torch.bfloat16)
+
+
+if __name__ == "__main__":
+    unittest.main()