Refactor device timer, clean up metrics collector, and add fwd occupancy metric (#24197)

test/registered/observability/test_metrics.py

@@ -62,12 +62,12 @@ class TestEnableMetrics(CustomTestCase):
                     {"mode": "decode"},
                 ),
                 (
-                    "sglang:dp_cooperation_gpu_execution_seconds_total",
-                    {"category": "forward_extend"},
+                    "sglang:dp_cooperation_forward_execution_seconds_total",
+                    {"category": "extend"},
                 ),
                 (
-                    "sglang:dp_cooperation_gpu_execution_seconds_total",
-                    {"category": "forward_decode"},
+                    "sglang:dp_cooperation_forward_execution_seconds_total",
+                    {"category": "decode"},
                 ),
             ]
             _check_metrics_positive(self, metrics, metrics_to_check)
@@ -129,15 +129,17 @@ class TestEnableMetrics(CustomTestCase):
             for _ in response.iter_lines(decode_unicode=False):
                 pass
 
-            response = requests.post(
-                f"{DEFAULT_URL_FOR_TEST}/generate",
-                json={
-                    "text": "Hello",
-                    "sampling_params": {"temperature": 0, "max_new_tokens": 5},
-                },
-                headers={"x-smg-routing-key": "test-key"},
-            )
-            self.assertEqual(response.status_code, 200)
+            for i in range(2):
+                # Send the request twice to trigger cached token metrics
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/generate",
+                    json={
+                        "text": "Hello, " * 100,
+                        "sampling_params": {"temperature": 0, "max_new_tokens": 5},
+                    },
+                    headers={"x-smg-routing-key": "test-key"},
+                )
+                self.assertEqual(response.status_code, 200)
 
             # Get metrics
             metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
@@ -209,8 +211,8 @@ class TestEnableMetrics(CustomTestCase):
         metrics_to_check = [
             ("sglang:realtime_tokens_total", {"mode": "prefill_compute"}),
             ("sglang:realtime_tokens_total", {"mode": "decode"}),
-            ("sglang:gpu_execution_seconds_total", {"category": "forward_extend"}),
-            ("sglang:gpu_execution_seconds_total", {"category": "forward_decode"}),
+            ("sglang:forward_execution_seconds_total", {"category": "extend"}),
+            ("sglang:forward_execution_seconds_total", {"category": "decode"}),
             ("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
         ]
         _check_metrics_positive(self, metrics, metrics_to_check)

test/registered/perf/test_bench_one_batch_1gpu.py

@@ -1,12 +1,17 @@
+import os
+import re
+import subprocess
 import unittest
 
+import numpy as np
+
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST,
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     CustomTestCase,
     is_in_ci,
-    run_bench_offline_throughput,
+    kill_process_tree,
     run_bench_one_batch,
     write_github_step_summary,
 )
@@ -24,10 +29,46 @@ class TestBenchOneBatch1GPU(CustomTestCase):
         self.assertGreater(output_throughput, 50)
 
     def test_bs1_default(self):
-        output_throughput = run_bench_offline_throughput(
-            DEFAULT_MODEL_NAME_FOR_TEST, ["--cuda-graph-max-bs", "2"]
+        env = os.environ.copy()
+        env["SGLANG_ENABLE_METRICS_DEVICE_TIMER"] = "1"
+
+        command = [
+            "python3",
+            "-m",
+            "sglang.bench_offline_throughput",
+            "--num-prompts",
+            "1",
+            "--dataset-name",
+            "random",
+            "--random-input-len",
+            "256",
+            "--random-output-len",
+            "1024",
+            "--model-path",
+            DEFAULT_MODEL_NAME_FOR_TEST,
+            "--cuda-graph-max-bs",
+            "2",
+        ]
+
+        print(f"command={' '.join(command)}")
+        process = subprocess.Popen(
+            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
         )
 
+        try:
+            stdout, stderr = process.communicate()
+            output = stdout.decode(errors="backslashreplace")
+            error = stderr.decode(errors="backslashreplace")
+            print(f"Output: {output}", flush=True)
+            print(f"Error: {error}", flush=True)
+
+            output_throughput = -1
+            for line in output.split("\n"):
+                if "Last generation throughput (tok/s):" in line:
+                    output_throughput = float(line.split(":")[-1])
+        finally:
+            kill_process_tree(process.pid)
+
         if is_in_ci():
             write_github_step_summary(
                 f"### test_bs1_default (llama-3.1-8b)\n"
@@ -35,6 +76,23 @@ class TestBenchOneBatch1GPU(CustomTestCase):
             )
             self.assertGreater(output_throughput, 135)
 
+        fwd_occupancy_values = []
+        for line in error.split("\n"):
+            match = re.search(r"fwd occupancy:\s*([\d.]+|nan)%", line)
+            if match:
+                val = match.group(1)
+                if val != "nan":
+                    fwd_occupancy_values.append(float(val))
+
+        print(f"{fwd_occupancy_values=}", flush=True)
+        self.assertGreater(
+            len(fwd_occupancy_values), 0, "No fwd occupancy values found in logs"
+        )
+
+        fwd_occupancy_p90 = float(np.percentile(fwd_occupancy_values, 90))
+        print(f"{fwd_occupancy_p90=}", flush=True)
+        self.assertGreater(fwd_occupancy_p90, 97.5)
+
 
 if __name__ == "__main__":
     unittest.main()