add dflash gemma4 support (#27471)

Co-authored-by: kpham-sgl <khoa.pham@radixark.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-17 16:39:28 -07:00
parent cd60c4edd0
commit 5ea0d1d093
4 changed files with 212 additions and 17 deletions
@@ -0,0 +1,177 @@
+import unittest
+from types import SimpleNamespace
+from typing import Optional
+
+import requests
+
+from sglang.srt.utils import 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,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_cuda_ci(est_time=720, stage="extra-a", runner_config="2-gpu-large")
+
+MODEL_NAME = "31B"
+TARGET_PATH = "google/gemma-4-31B-it"
+DRAFT_PATH = "z-lab/gemma-4-31B-it-DFlash"
+TENSOR_PARALLEL_SIZE = 2
+
+DRAFT_ATTENTION_BACKEND = "flashinfer"
+SPECULATIVE_NUM_DRAFT_TOKENS = 16
+GSM8K_NUM_EXAMPLES = 200
+GSM8K_NUM_THREADS = 128
+SERVER_LAUNCH_TIMEOUT = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3
+
+# Match the existing Gemma4 31B MTP accuracy floor.
+GSM8K_SCORE_THRESHOLD = 0.75
+ACCEPT_LENGTH_THRESHOLD = 5.4
+
+
+def get_server_info(base_url: str) -> dict:
+    response = requests.get(base_url + "/server_info", timeout=10)
+    response.raise_for_status()
+    return response.json()
+
+
+def get_avg_spec_accept_length(base_url: str) -> Optional[float]:
+    try:
+        info = get_server_info(base_url)
+    except Exception:
+        return None
+    internal_states = info.get("internal_states") or []
+    if not internal_states:
+        return None
+    value = internal_states[0].get("avg_spec_accept_length")
+    if value is None:
+        return None
+    return float(value)
+
+
+class TestGemma4DFlash31B(CustomTestCase):
+    base_url = DEFAULT_URL_FOR_TEST
+
+    @classmethod
+    def _common_server_args(cls) -> list[str]:
+        args = [
+            "--trust-remote-code",
+            "--attention-backend",
+            "triton",
+            "--dtype",
+            "bfloat16",
+            "--mem-fraction-static",
+            "0.55",
+            "--max-running-requests",
+            "16",
+            "--context-length",
+            "2048",
+            "--max-total-tokens",
+            "32768",
+            "--skip-server-warmup",
+        ]
+        if TENSOR_PARALLEL_SIZE > 1:
+            args += ["--tp-size", str(TENSOR_PARALLEL_SIZE)]
+        return args
+
+    @classmethod
+    def _server_args(cls) -> list[str]:
+        return [
+            "--speculative-algorithm",
+            "DFLASH",
+            "--speculative-draft-model-path",
+            DRAFT_PATH,
+            "--speculative-num-draft-tokens",
+            str(SPECULATIVE_NUM_DRAFT_TOKENS),
+            "--speculative-draft-attention-backend",
+            DRAFT_ATTENTION_BACKEND,
+        ] + cls._common_server_args()
+
+    @classmethod
+    def _gsm8k_args(cls) -> SimpleNamespace:
+        return SimpleNamespace(
+            base_url=cls.base_url,
+            model=TARGET_PATH,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=GSM8K_NUM_EXAMPLES,
+            num_threads=GSM8K_NUM_THREADS,
+            num_shots=5,
+        )
+
+    @staticmethod
+    def _stop_process(process) -> None:
+        try:
+            kill_process_tree(process.pid)
+        except Exception:
+            pass
+
+    def test_gsm8k_dflash(self) -> None:
+        process = None
+        try:
+            process = popen_launch_server(
+                TARGET_PATH,
+                self.base_url,
+                timeout=SERVER_LAUNCH_TIMEOUT,
+                other_args=self._server_args(),
+            )
+            requests.get(self.base_url + "/flush_cache", timeout=30)
+
+            server_info = get_server_info(self.base_url)
+            self.assertEqual(
+                server_info.get("speculative_algorithm"),
+                "DFLASH",
+                f"{MODEL_NAME}: server did not start with DFLASH",
+            )
+            self.assertEqual(
+                server_info.get("speculative_draft_attention_backend"),
+                DRAFT_ATTENTION_BACKEND,
+                f"{MODEL_NAME}: unexpected DFLASH draft attention backend",
+            )
+            self.assertEqual(
+                server_info.get("speculative_num_draft_tokens"),
+                SPECULATIVE_NUM_DRAFT_TOKENS,
+                f"{MODEL_NAME}: unexpected DFLASH block size",
+            )
+            self.assertFalse(
+                bool(server_info.get("disable_cuda_graph")),
+                f"{MODEL_NAME}: CUDA graph is disabled",
+            )
+
+            metrics = run_eval(self._gsm8k_args())
+            dflash_score = float(metrics["score"])
+            avg_accept = get_avg_spec_accept_length(self.base_url)
+        finally:
+            if process is not None:
+                self._stop_process(process)
+
+        print(
+            f"[Gemma4 {MODEL_NAME} DFlash] "
+            f"score={dflash_score:.4f} threshold={GSM8K_SCORE_THRESHOLD:.4f} "
+            f"avg_spec_accept_length={avg_accept}"
+        )
+        if is_in_ci():
+            write_github_step_summary(
+                f"### Gemma4 {MODEL_NAME} DFlash\n"
+                f"score={dflash_score:.4f}\n"
+                f"threshold={GSM8K_SCORE_THRESHOLD:.4f}\n"
+                f"avg_spec_accept_length={avg_accept}\n"
+            )
+
+        self.assertGreaterEqual(dflash_score, GSM8K_SCORE_THRESHOLD)
+        self.assertIsNotNone(avg_accept)
+        self.assertGreaterEqual(
+            avg_accept,
+            ACCEPT_LENGTH_THRESHOLD,
+            f"{MODEL_NAME}: accept length too low",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()