From f5fdf9c5d812e952edf84adf6e1dcb701d682f07 Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Mon, 8 Jun 2026 14:50:31 +0800
Subject: [PATCH] Speed up dump comparator percentile computation using numpy
 (#26874)

---
 .../comparator/tensor_comparator/comparator.py |  7 +++++--
 .../tensor_comparator/test_comparator.py       | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
index 4549218c9..9becd78c7 100644
--- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
+++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
@@ -125,8 +125,11 @@ def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
 def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
     if not include:
         return {}
-    x_float: torch.Tensor = x.float()
-    return {p: torch.quantile(x_float, p / 100.0).item() for p in DEFAULT_PERCENTILES}
+    import numpy as np
+
+    arr = x.detach().float().numpy().ravel()
+    values = np.percentile(arr, list(DEFAULT_PERCENTILES))
+    return {p: float(v) for p, v in zip(DEFAULT_PERCENTILES, values)}
 
 
 def compute_diff(
diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
index 7e6c26212..3b1a4d382 100644
--- a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
+++ b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
@@ -192,6 +192,24 @@ class TestComputeTensorStats:
         assert stats.percentiles[95] == pytest.approx(95.0, abs=0.5)
         assert stats.percentiles[99] == pytest.approx(99.0, abs=0.5)
 
+    def test_percentiles_exact_for_uniform_range(self):
+        """Percentiles of arange(0, 101) equal the percentile index exactly (linear interp)."""
+        x = torch.arange(0, 101, dtype=torch.float32)
+        stats = _compute_tensor_stats(x)
+
+        for p in (1, 5, 50, 95, 99):
+            assert stats.percentiles[p] == pytest.approx(float(p), abs=1e-4)
+
+    def test_percentiles_match_torch_quantile_reference(self):
+        """numpy-based percentiles must match torch.quantile on the same data within tight tolerance."""
+        torch.manual_seed(0)
+        x = torch.randn(5000)
+        stats = _compute_tensor_stats(x)
+
+        for p in (1, 5, 50, 95, 99):
+            expected = torch.quantile(x.float(), p / 100.0).item()
+            assert stats.percentiles[p] == pytest.approx(expected, abs=1e-4)
+
     def test_large_tensor_skips_quantiles(self):
         x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
         stats = _compute_tensor_stats(x)