Enhance metrics in dump comparator (#19560)

python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py

@@ -3,6 +3,7 @@ from typing import Optional
 import torch
 
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
+    DEFAULT_PERCENTILES,
     DiffInfo,
     TensorComparisonInfo,
     TensorInfo,
@@ -89,21 +90,22 @@ def compare_tensor_pair(
 
 
 def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
-    include_quantiles = x.numel() < QUANTILE_NUMEL_THRESHOLD
+    include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD
     return TensorStats(
         mean=torch.mean(x).item(),
+        abs_mean=torch.mean(x.abs()).item(),
         std=torch.std(x).item(),
         min=torch.min(x).item(),
         max=torch.max(x).item(),
-        p1=_quantile_or_none(x, q=0.01, include=include_quantiles),
+        percentiles=_compute_percentiles(x, include=include_quantiles),
-        p5=_quantile_or_none(x, q=0.05, include=include_quantiles),
-        p95=_quantile_or_none(x, q=0.95, include=include_quantiles),
-        p99=_quantile_or_none(x, q=0.99, include=include_quantiles),
     )
 
 
-def _quantile_or_none(x: torch.Tensor, *, q: float, include: bool) -> Optional[float]:
+def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
-    return torch.quantile(x, q).item() if include else None
+    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}
 
 
 def _compute_diff(
@@ -118,10 +120,15 @@ def _compute_diff(
     max_abs_diff = raw_abs_diff.max().item()
     mean_abs_diff = raw_abs_diff.mean().item()
 
+    include_quantiles: bool = raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD
+
     return DiffInfo(
         rel_diff=rel_diff,
         max_abs_diff=max_abs_diff,
         mean_abs_diff=mean_abs_diff,
+        abs_diff_percentiles=_compute_percentiles(
+            raw_abs_diff, include=include_quantiles
+        ),
         max_diff_coord=list(max_diff_coord),
         baseline_at_max=x_baseline[max_diff_coord].item(),
         target_at_max=x_target[max_diff_coord].item(),

python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py

@@ -56,20 +56,30 @@ def format_comparison(info: TensorComparisonInfo) -> str:
 
 def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]:
     lines: list[str] = []
+
     for stat_name in TensorStats.model_fields:
-        value_baseline = getattr(baseline, stat_name)
+        if stat_name == "percentiles":
-        value_target = getattr(target, stat_name)
-        if value_baseline is None or value_target is None:
             continue
+        value_baseline: float = getattr(baseline, stat_name)
+        value_target: float = getattr(target, stat_name)
         lines.append(
             f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
             f"(diff: {value_target - value_baseline:.4f})"
         )
+
+    for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
+        value_baseline = baseline.percentiles[p]
+        value_target = target.percentiles[p]
+        lines.append(
+            f"[p{p}] {value_baseline:.4f} vs {value_target:.4f} "
+            f"(diff: {value_target - value_baseline:.4f})"
+        )
+
     return lines
 
 
 def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
-    return [
+    lines: list[str] = [
         prefix_text
         + "\t".join(
             f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
@@ -83,3 +93,12 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
         f"baseline={diff.baseline_at_max} "
         f"target={diff.target_at_max}",
     ]
+
+    if diff.abs_diff_percentiles:
+        quantile_parts: list[str] = [
+            f"p{p}={value:.4f}"
+            for p, value in sorted(diff.abs_diff_percentiles.items())
+        ]
+        lines.append("[abs_diff] " + " ".join(quantile_parts))
+
+    return lines

python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py

@@ -2,16 +2,16 @@ from typing import Optional
 
 from sglang.srt.debug_utils.comparator.utils import _StrictBase
 
+DEFAULT_PERCENTILES: tuple[int, ...] = (1, 5, 50, 95, 99)
+
 
 class TensorStats(_StrictBase):
     mean: float
+    abs_mean: float
     std: float
     min: float
     max: float
-    p1: Optional[float] = None
+    percentiles: dict[int, float] = {}
-    p5: Optional[float] = None
-    p95: Optional[float] = None
-    p99: Optional[float] = None
 
 
 class TensorInfo(_StrictBase):
@@ -25,6 +25,7 @@ class DiffInfo(_StrictBase):
     rel_diff: float
     max_abs_diff: float
     mean_abs_diff: float
+    abs_diff_percentiles: dict[int, float] = {}
     max_diff_coord: list[int]
     baseline_at_max: float
     target_at_max: float

test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py

@@ -21,28 +21,34 @@ class TestComputeTensorStats:
         stats = _compute_tensor_stats(x)
 
         assert stats.mean == pytest.approx(3.0, abs=1e-4)
+        assert stats.abs_mean == pytest.approx(3.0, abs=1e-4)
         assert stats.std == pytest.approx(1.5811, abs=1e-3)
         assert stats.min == pytest.approx(1.0, abs=1e-4)
         assert stats.max == pytest.approx(5.0, abs=1e-4)
 
+    def test_abs_mean_with_negative_values(self):
+        x = torch.tensor([-3.0, -1.0, 1.0, 3.0])
+        stats = _compute_tensor_stats(x)
+
+        assert stats.mean == pytest.approx(0.0, abs=1e-4)
+        assert stats.abs_mean == pytest.approx(2.0, abs=1e-4)
+
     def test_quantile_values(self):
         x = torch.linspace(0.0, 100.0, steps=1000)
         stats = _compute_tensor_stats(x)
 
-        assert stats.p1 == pytest.approx(1.0, abs=0.5)
+        assert stats.percentiles[1] == pytest.approx(1.0, abs=0.5)
-        assert stats.p5 == pytest.approx(5.0, abs=0.5)
+        assert stats.percentiles[5] == pytest.approx(5.0, abs=0.5)
-        assert stats.p95 == pytest.approx(95.0, abs=0.5)
+        assert stats.percentiles[50] == pytest.approx(50.0, abs=0.5)
-        assert stats.p99 == pytest.approx(99.0, abs=0.5)
+        assert stats.percentiles[95] == pytest.approx(95.0, abs=0.5)
+        assert stats.percentiles[99] == pytest.approx(99.0, abs=0.5)
 
     def test_large_tensor_skips_quantiles(self):
         x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
         stats = _compute_tensor_stats(x)
 
         assert stats.mean is not None
-        assert stats.p1 is None
+        assert stats.percentiles == {}
-        assert stats.p5 is None
-        assert stats.p95 is None
-        assert stats.p99 is None
 
 
 class TestComputeDiff:
@@ -53,6 +59,9 @@ class TestComputeDiff:
         assert diff.rel_diff == pytest.approx(0.0, abs=1e-5)
         assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
         assert diff.mean_abs_diff == pytest.approx(0.0, abs=1e-5)
+        assert diff.abs_diff_percentiles[50] == pytest.approx(0.0, abs=1e-5)
+        assert diff.abs_diff_percentiles[95] == pytest.approx(0.0, abs=1e-5)
+        assert diff.abs_diff_percentiles[99] == pytest.approx(0.0, abs=1e-5)
         assert diff.passed is True
 
     def test_known_offset(self):
@@ -67,8 +76,18 @@ class TestComputeDiff:
         assert diff.baseline_at_max == pytest.approx(1.0, abs=1e-4)
         assert diff.target_at_max == pytest.approx(1.5, abs=1e-4)
         assert diff.mean_abs_diff == pytest.approx(0.5 / 100, abs=1e-4)
+        assert diff.abs_diff_percentiles[1] == pytest.approx(0.0, abs=1e-4)
+        assert diff.abs_diff_percentiles[50] == pytest.approx(0.0, abs=1e-4)
+        assert diff.abs_diff_percentiles[99] > 0
         assert diff.passed is False
 
+    def test_large_tensor_skips_diff_quantiles(self):
+        x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
+        y = x + 0.001
+        diff = _compute_diff(x_baseline=x, x_target=y)
+
+        assert diff.abs_diff_percentiles == {}
+
     def test_rel_diff_value(self):
         x = torch.tensor([1.0, 0.0])
         y = torch.tensor([0.0, 1.0])

test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py

@@ -16,25 +16,47 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 
+_DEFAULT_PERCENTILES: dict[int, float] = {
+    1: -1.8,
+    5: -1.5,
+    50: 0.0,
+    95: 1.5,
+    99: 1.8,
+}
+
+
 def _make_stats(
     mean: float = 0.0,
+    abs_mean: float = 0.8,
     std: float = 1.0,
     min: float = -2.0,
     max: float = 2.0,
-    p1: float | None = -1.8,
+    percentiles: dict[int, float] | None = None,
-    p5: float | None = -1.5,
-    p95: float | None = 1.5,
-    p99: float | None = 1.8,
 ) -> TensorStats:
     return TensorStats(
-        mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
+        mean=mean,
+        abs_mean=abs_mean,
+        std=std,
+        min=min,
+        max=max,
+        percentiles=percentiles if percentiles is not None else _DEFAULT_PERCENTILES,
     )
 
 
+_DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = {
+    1: 0.0001,
+    5: 0.0001,
+    50: 0.0002,
+    95: 0.0004,
+    99: 0.0005,
+}
+
+
 def _make_diff(
     rel_diff: float = 0.0001,
     max_abs_diff: float = 0.0005,
     mean_abs_diff: float = 0.0002,
+    abs_diff_percentiles: dict[int, float] | None = None,
     diff_threshold: float = 1e-3,
     passed: bool = True,
 ) -> DiffInfo:
@@ -42,6 +64,11 @@ def _make_diff(
         rel_diff=rel_diff,
         max_abs_diff=max_abs_diff,
         mean_abs_diff=mean_abs_diff,
+        abs_diff_percentiles=(
+            abs_diff_percentiles
+            if abs_diff_percentiles is not None
+            else _DEFAULT_ABS_DIFF_PERCENTILES
+        ),
         max_diff_coord=[2, 3],
         baseline_at_max=1.0,
         target_at_max=1.0005,
@@ -89,16 +116,19 @@ class TestFormatComparison:
             "After unify [shape] [4, 8] vs [4, 8]\t"
             "[dtype] torch.float32 vs torch.float32\n"
             "[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0001 (diff: 0.0001)\n"
             "[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
             "[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
             "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
+            "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
             "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005"
+            "baseline=1.0 target=1.0005\n"
+            "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
         )
 
     def test_shape_mismatch(self):
@@ -116,11 +146,13 @@ class TestFormatComparison:
             "After unify [shape] [3, 4] vs [5, 6]\t"
             "[dtype] torch.float32 vs torch.float32\n"
             "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
             "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
+            "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
             "⚠️ Shape mismatch"
@@ -148,20 +180,24 @@ class TestFormatComparison:
             "After unify [shape] [4, 8] vs [4, 8]\t"
             "[dtype] torch.float32 vs torch.bfloat16\n"
             "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
             "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
+            "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
             "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005\n"
+            "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
             "When downcast to torch.bfloat16: "
             "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005"
+            "baseline=1.0 target=1.0005\n"
+            "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
         )
 
     def test_with_shape_unification(self):
@@ -182,16 +218,19 @@ class TestFormatComparison:
             "After unify [shape] [4, 8] vs [4, 8]\t"
             "[dtype] torch.float32 vs torch.float32\n"
             "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
             "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
+            "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
             "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
-            "baseline=1.0 target=1.0005"
+            "baseline=1.0 target=1.0005\n"
+            "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
         )
 
     def test_with_samples(self):
@@ -210,22 +249,25 @@ class TestFormatComparison:
             "After unify [shape] [4, 8] vs [4, 8]\t"
             "[dtype] torch.float32 vs torch.float32\n"
             "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
             "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
             "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
+            "[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
             "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
             "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
             "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n"
             "max_abs_diff happens at coord=[2, 3] with "
             "baseline=1.0 target=1.0005\n"
+            "[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
             "x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
             "x_target(sample)=tensor([0.1, 0.3, ...])"
         )
 
-    def test_none_quantiles(self):
+    def test_empty_percentiles(self):
-        stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
+        stats_no_quantiles = _make_stats(percentiles={})
 
         info = TensorComparisonInfo(
             name="no_quantiles",
@@ -233,7 +275,7 @@ class TestFormatComparison:
             target=_make_tensor_info(stats=stats_no_quantiles),
             unified_shape=[4, 8],
             shape_mismatch=False,
-            diff=_make_diff(),
+            diff=_make_diff(abs_diff_percentiles={}),
         )
 
         assert format_comparison(info) == (
@@ -242,6 +284,7 @@ class TestFormatComparison:
             "After unify [shape] [4, 8] vs [4, 8]\t"
             "[dtype] torch.float32 vs torch.float32\n"
             "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
+            "[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
             "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
             "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
             "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"

test/registered/debug_utils/comparator/tensor_comparator/test_types.py

@@ -23,16 +23,14 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 
-def _make_stats(**overrides: float) -> TensorStats:
+def _make_stats(**overrides) -> TensorStats:
-    defaults = dict(
+    defaults: dict = dict(
         mean=0.5,
+        abs_mean=1.2,
         std=1.0,
         min=-2.0,
         max=3.0,
-        p1=-1.8,
+        percentiles={1: -1.8, 5: -1.5, 50: 0.0, 95: 2.5, 99: 2.8},
-        p5=-1.5,
-        p95=2.5,
-        p99=2.8,
     )
     defaults.update(overrides)
     return TensorStats(**defaults)
@@ -66,7 +64,7 @@ def _make_tensor_info(**overrides) -> TensorInfo:
 class TestStrictBase:
     def test_rejects_extra_fields(self):
         with pytest.raises(Exception):
-            TensorStats(mean=0.0, std=1.0, min=-1.0, max=1.0, bogus=42)
+            TensorStats(mean=0.0, abs_mean=0.5, std=1.0, min=-1.0, max=1.0, bogus=42)
 
     def test_rejects_extra_fields_on_diff(self):
         with pytest.raises(Exception):

test/registered/debug_utils/comparator/test_model_validation.py

@@ -187,7 +187,7 @@ def _make_tensor_info() -> TensorInfo:
     return TensorInfo(
         shape=[4, 4],
         dtype="float32",
-        stats=TensorStats(mean=0.0, std=1.0, min=-2.0, max=2.0),
+        stats=TensorStats(mean=0.0, abs_mean=0.8, std=1.0, min=-2.0, max=2.0),
     )