From e4696ed62d6f7b4d2531e5fe0e406dfe4dd88163 Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Thu, 25 Jun 2026 15:08:33 -0700
Subject: [PATCH] [test] Report token tps in fwd occupancy kit and force
 ignore_eos (#29332)

---
 python/sglang/test/kits/fwd_occupancy_kit.py | 62 +++++++++++++-------
 1 file changed, 42 insertions(+), 20 deletions(-)

diff --git a/python/sglang/test/kits/fwd_occupancy_kit.py b/python/sglang/test/kits/fwd_occupancy_kit.py
index 4a34e7f63..6385a61e3 100644
--- a/python/sglang/test/kits/fwd_occupancy_kit.py
+++ b/python/sglang/test/kits/fwd_occupancy_kit.py
@@ -82,27 +82,32 @@ class FwdOccupancyMixin:
         resp = requests.get(
             self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
         )
-        assert resp.status_code == 200, (
-            f"/metrics returned {resp.status_code}; the test class's server "
-            "must be launched with --enable-metrics"
-        )
-        assert "sglang:fwd_occupancy" in resp.text, (
-            "sglang:fwd_occupancy gauge not exposed; set "
-            "SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env and "
-            "pass --enable-metrics"
-        )
+        if resp.status_code != 200:
+            raise AssertionError(
+                f"/metrics returned {resp.status_code}; the test class's "
+                "server must be launched with --enable-metrics"
+            )
+        if "sglang:fwd_occupancy" not in resp.text:
+            raise AssertionError(
+                "sglang:fwd_occupancy gauge not exposed; set "
+                "SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env "
+                "and pass --enable-metrics"
+            )
 
     def _fwd_occupancy_fire(self, prompt: str, max_new_tokens: int):
-        """Fire one /generate. Must not be called concurrently -- that
-        would break the single-batch invariant."""
+        """Fire one /generate, return (completion_tokens, wall_time).
+        Must not be called concurrently -- that would break the
+        single-batch invariant."""
+        t0 = time.perf_counter()
         try:
-            requests.post(
+            resp = requests.post(
                 self.base_url + "/generate",
                 json={
                     "text": prompt,
                     "sampling_params": {
                         "temperature": 0.0,
                         "max_new_tokens": max_new_tokens,
+                        "ignore_eos": True,
                     },
                 },
                 timeout=_GENERATE_REQUEST_TIMEOUT,
@@ -110,7 +115,13 @@ class FwdOccupancyMixin:
         except requests.RequestException:
             # Final stats-vs-threshold is the signal; individual fire
             # failure isn't.
-            pass
+            return 0, 0.0
+        elapsed = time.perf_counter() - t0
+        try:
+            tokens = resp.json().get("meta_info", {}).get("completion_tokens", 0)
+        except ValueError:  # non-JSON body
+            tokens = 0
+        return tokens, elapsed
 
     def _fwd_occupancy_warmup(self):
         """Fill cuda graphs + step the device-timer past its first NaN
@@ -123,15 +134,19 @@ class FwdOccupancyMixin:
 
     def _fwd_occupancy_measure(self):
         """Background-fire one long single-batch request, scrape
-        /metrics on the foreground; return non-NaN samples."""
+        /metrics on the foreground; return (non-NaN samples,
+        token_tps)."""
         samples = []
         request_done = threading.Event()
+        result = {"completion_tokens": 0, "elapsed": 0.0}
 
         def fire_one():
             try:
-                self._fwd_occupancy_fire(
-                    self.fwd_occupancy_prompt,
-                    self.fwd_occupancy_max_new_tokens,
+                result["completion_tokens"], result["elapsed"] = (
+                    self._fwd_occupancy_fire(
+                        self.fwd_occupancy_prompt,
+                        self.fwd_occupancy_max_new_tokens,
+                    )
                 )
             finally:
                 request_done.set()
@@ -146,12 +161,17 @@ class FwdOccupancyMixin:
             time.sleep(self.fwd_occupancy_scrape_interval)
 
         firer.join(timeout=_GENERATE_REQUEST_TIMEOUT)
-        return samples
+        token_tps = (
+            result["completion_tokens"] / result["elapsed"]
+            if result["elapsed"] > 0
+            else 0.0
+        )
+        return samples, token_tps
 
     def test_fwd_occupancy(self):
         self._assert_metrics_device_timer_enabled()
         self._fwd_occupancy_warmup()
-        samples = self._fwd_occupancy_measure()
+        samples, token_tps = self._fwd_occupancy_measure()
 
         self.assertGreaterEqual(
             len(samples),
@@ -166,7 +186,8 @@ class FwdOccupancyMixin:
         samples_sorted = sorted(samples)
         median = statistics.median(samples_sorted)
         peak = samples_sorted[-1]
-        p10 = samples_sorted[max(0, len(samples_sorted) // 10 - 1)]
+        p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10))
+        p10 = samples_sorted[p10_idx]
         print(
             "\n"
             + tabulate.tabulate(
@@ -176,6 +197,7 @@ class FwdOccupancyMixin:
                     ["peak", f"{peak:.2f}"],
                     ["p10", f"{p10:.2f}"],
                     ["threshold", f"{self.fwd_occupancy_threshold:.2f}"],
+                    ["token tps", f"{token_tps:.2f}"],
                 ],
                 headers=["fwd_occupancy", "value"],
                 tablefmt="github",