[Gemma4] Optimize Gemm4 with fused Q/K/V RMSNorm + per-expert FP8 ckpt loader (#24696)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
2026-05-10 00:24:12 -07:00
parent a87fb399de
commit d3fd91ed97
4 changed files with 317 additions and 15 deletions
@@ -0,0 +1,109 @@
+"""End-to-end test for compressed-tensors per-expert FP8 MoE checkpoint
+loading on Gemma4 (e.g. RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic).
+
+Regression coverage for the load_weights path that recognises
+`experts.<id>.{gate,up,down}_proj.{weight,weight_scale}` keys and folds
+them into SGLang's fused FusedMoE parameters. Without that path, all
+routed-expert weights are silently skipped at load time and the model
+emits only `<pad>` tokens at inference (GSM8K collapses to 0.0).
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import get_device_sm, kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+# Compressed-tensors per-expert FP8 MoE checkpoint that exercises the
+# loader path (gated repo + ~27 GB download + 4 GPUs at TP=4).
+register_cuda_ci(est_time=120, suite="stage-c-test-4-gpu-h100")
+
+
+@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
+class TestGemma4FP8PerExpertLoading(CustomTestCase):
+    """Three-stage check that catches the silent-skip failure mode:
+    1. server health
+    2. completion is not the all-`<pad>` garbage state
+    3. GSM8K accuracy matches the BF16 baseline
+    """
+
+    model = "RedHatAI/gemma-4-26B-A4B-it-FP8-Dynamic"
+    base_url = DEFAULT_URL_FOR_TEST
+
+    @classmethod
+    def setUpClass(cls):
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp",
+                "4",
+                "--trust-remote-code",
+                "--random-seed",
+                "42",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_health(self):
+        r = requests.get(self.base_url + "/health")
+        self.assertEqual(r.status_code, 200)
+
+    def test_basic_generation_not_garbage(self):
+        """Pre-fix the server starts but every routed expert is zero-init,
+        which leads chat completions to deterministic `<pad>` spam."""
+        r = requests.post(
+            self.base_url + "/v1/chat/completions",
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": "What is 7 + 5?"}],
+                "temperature": 0,
+                "max_tokens": 32,
+            },
+        )
+        self.assertEqual(r.status_code, 200)
+        text = r.json()["choices"][0]["message"]["content"]
+        self.assertNotIn(
+            "<pad>", text, f"Output looks like the pre-fix garbage state: {text!r}"
+        )
+        self.assertGreater(len(text.strip()), 0, "Empty completion")
+        self.assertIn("12", text, f"Expected the answer to mention '12': {text!r}")
+
+    def test_gsm8k_accuracy(self):
+        """Pre-fix this scores exactly 0.00 (zero routed-expert weights);
+        post-fix it matches the BF16 baseline (~0.95 on 20 samples)."""
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=20,
+            num_threads=16,
+        )
+        metrics = run_eval(args)
+        score = float(metrics["score"])
+        print(f"Gemma4 FP8 per-expert GSM8K-20 score: {score:.3f}")
+        # Threshold rules out the failure mode (0.00) while leaving ample
+        # margin under the BF16 baseline (~0.95).
+        self.assertGreaterEqual(
+            score,
+            0.80,
+            f"Per-expert FP8 ckpt accuracy collapsed: {score} "
+            "(pre-fix value is 0.00; BF16 baseline is ~0.95).",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()