Cap diagnostic detail computation for failing tensors (#30656)

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

@@ -41,6 +41,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
     _split_logs,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
+    FailureDisplayBudget,
     compare_tensor_pair,
     compute_tensor_info,
 )
@@ -132,6 +133,7 @@ def compare_bundle_pair(
     token_aligner_mode: Optional[str],
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
@@ -147,6 +149,7 @@ def compare_bundle_pair(
             token_aligner_mode=token_aligner_mode,
             token_aligner_plan=token_aligner_plan,
             diff_threshold_rules=diff_threshold_rules,
+            failure_display_budget=failure_display_budget,
             thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
             viz_output_dir=viz_output_dir,
             compute_per_token=compute_per_token,
@@ -165,6 +168,7 @@ def _compare_bundle_pair_inner(
     token_aligner_mode: Optional[str],
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
@@ -221,6 +225,7 @@ def _compare_bundle_pair_inner(
         token_aligner_mode=token_aligner_mode,
         token_aligner_plan=token_aligner_plan,
         diff_threshold_rules=diff_threshold_rules,
+        failure_display_budget=failure_display_budget,
         thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
         viz_output_dir=viz_output_dir,
         compute_per_token=compute_per_token,
@@ -244,6 +249,7 @@ def _compare_bundle_pair_tensor_type(
     token_aligner_mode: Optional[str],
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
@@ -309,6 +315,7 @@ def _compare_bundle_pair_tensor_type(
         x_target=aligned_target,
         name=name,
         diff_threshold_rules=diff_threshold_rules,
+        failure_display_budget=failure_display_budget,
         seq_dim=seq_dim,
     )
     record = ComparisonTensorRecord(

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

@@ -41,6 +41,7 @@ from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
 from sglang.srt.debug_utils.comparator.report_sink import report_sink
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
     DEFAULT_PREDICATE,
+    FailureDisplayBudget,
 )
 from sglang.srt.debug_utils.comparator.threshold_dsl import (
     DiffThresholdRule,
@@ -150,6 +151,7 @@ def run(args: argparse.Namespace) -> int:
             diff_threshold_rules=parse_diff_threshold_rules(
                 args.diff_threshold, default_predicate=DEFAULT_PREDICATE
             ),
+            failure_display_budget=FailureDisplayBudget(),
             thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair,
             viz_output_dir=viz_output_dir,
             compute_per_token=visualize_per_token is not None,
@@ -230,6 +232,7 @@ def _compare_bundle_pairs(
     token_aligner_mode: Optional[str],
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
@@ -265,6 +268,7 @@ def _compare_bundle_pairs(
                 token_aligner_mode=token_aligner_mode,
                 token_aligner_plan=token_aligner_plan,
                 diff_threshold_rules=diff_threshold_rules,
+                failure_display_budget=failure_display_budget,
                 thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
                 viz_output_dir=viz_output_dir,
                 compute_per_token=compute_per_token,

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

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import Optional
 
 import torch
@@ -30,11 +31,30 @@ SAMPLE_DIFF_THRESHOLD = 1e-3
 DEFAULT_PREDICATE: str = "rel <= 0.001"
 
 
+@dataclass
+class FailureDisplayBudget:
+    max_detail: int = 50
+    num_emitted: int = 0
+
+    def take(self) -> bool:
+        if self.max_detail < 0:
+            return True
+        if self.num_emitted >= self.max_detail:
+            return False
+        self.num_emitted += 1
+        return True
+
+
 def compute_tensor_info(
-    tensor: torch.Tensor, *, include_sample: bool = False
+    tensor: torch.Tensor,
+    *,
+    include_sample: bool = False,
+    include_percentiles: bool = True,
 ) -> TensorInfo:
     """Compute TensorInfo (shape, dtype, stats, optional sample) for a single tensor."""
-    stats: TensorStats = _compute_tensor_stats(tensor.float())
+    stats: TensorStats = _compute_tensor_stats(
+        tensor.float(), include_percentiles=include_percentiles
+    )
     sample: Optional[str] = (
         str(get_truncated_value(tensor.float())) if include_sample else None
     )
@@ -52,14 +72,13 @@ def compare_tensor_pair(
     name: str = "",
     diff_threshold_rules: Optional[list[DiffThresholdRule]] = None,
     seq_dim: Optional[int] = None,
+    failure_display_budget: Optional[FailureDisplayBudget] = None,
 ) -> TensorComparisonInfo:
     predicate = resolve_predicate(
         name, diff_threshold_rules, default_predicate=DEFAULT_PREDICATE
     )
 
-    baseline_info: TensorInfo = compute_tensor_info(x_baseline)
-    target_info: TensorInfo = compute_tensor_info(x_target)
-
+    x_baseline_original = x_baseline
     x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
     unified_shape = list(x_baseline.shape)
 
@@ -81,8 +100,31 @@ def compare_tensor_pair(
             x_target=x_target_f,
             predicate=predicate,
             seq_dim=seq_dim,
+            include_percentiles=False,
         )
 
+    is_failure = shape_mismatch or (diff is not None and not diff.passed)
+    needs_detail = is_failure and (
+        failure_display_budget is None or failure_display_budget.take()
+    )
+
+    baseline_info: TensorInfo = compute_tensor_info(
+        x_baseline_original, include_percentiles=needs_detail
+    )
+    target_info: TensorInfo = compute_tensor_info(
+        x_target, include_percentiles=needs_detail
+    )
+
+    if not shape_mismatch and needs_detail:
+        diff = compute_diff(
+            x_baseline=x_baseline_f,
+            x_target=x_target_f,
+            predicate=predicate,
+            seq_dim=seq_dim,
+            include_percentiles=True,
+        )
+
+    if diff is not None:
         needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD
         if needs_sample:
             baseline_info.sample = str(get_truncated_value(x_baseline_f))
@@ -97,6 +139,7 @@ def compare_tensor_pair(
                     x_baseline=x_baseline_f.to(downcast_dtype),
                     x_target=x_target_f.to(downcast_dtype),
                     predicate=predicate,
+                    include_percentiles=needs_detail,
                 )
 
     return TensorComparisonInfo(
@@ -111,7 +154,9 @@ def compare_tensor_pair(
     )
 
 
-def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
+def _compute_tensor_stats(
+    x: torch.Tensor, *, include_percentiles: bool = True
+) -> TensorStats:
     if x.numel() == 0:
         return TensorStats(
             mean=0.0,
@@ -122,7 +167,9 @@ def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
             percentiles={},
         )
 
-    include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD
+    include_quantiles: bool = (
+        include_percentiles and x.numel() < QUANTILE_NUMEL_THRESHOLD
+    )
     return TensorStats(
         mean=torch.mean(x).item(),
         abs_mean=torch.mean(x.abs()).item(),
@@ -148,6 +195,7 @@ def compute_diff(
     x_target: torch.Tensor,
     predicate: str = DEFAULT_PREDICATE,
     seq_dim: Optional[int] = None,
+    include_percentiles: bool = True,
 ) -> DiffInfo:
     if x_baseline.numel() == 0:
         return DiffInfo(
@@ -171,7 +219,9 @@ def compute_diff(
     )
     mean_abs_diff = raw_abs_diff.mean().item()
 
-    include_quantiles: bool = raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD
+    include_quantiles: bool = (
+        include_percentiles and raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD
+    )
 
     per_token_rel_diff: Optional[list[float]] = None
     if seq_dim is not None and x_baseline.dim() > seq_dim:

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

@@ -6,6 +6,7 @@ import torch
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
     QUANTILE_NUMEL_THRESHOLD,
     SAMPLE_DIFF_THRESHOLD,
+    FailureDisplayBudget,
     _compute_tensor_stats,
     compare_tensor_pair,
     compute_diff,
@@ -542,5 +543,136 @@ class TestCompareTensorPairPredicate:
         assert info.diff_downcast.predicate == "rel <= 0.0085 or max_abs <= 1e-4"
 
 
+class TestFailureDisplayBudget:
+    @staticmethod
+    def _failing_pair() -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.zeros(4, 4), torch.ones(4, 4)
+
+    def test_take_grants_exactly_max_detail_units(self) -> None:
+        """take() grants exactly max_detail units, then denies."""
+        budget = FailureDisplayBudget(max_detail=2)
+        assert budget.take() is True
+        assert budget.take() is True
+        assert budget.take() is False
+        assert budget.num_emitted == 2
+
+    def test_negative_max_detail_disables_limit(self) -> None:
+        """A negative max_detail always grants detail and never counts."""
+        budget = FailureDisplayBudget(max_detail=-1)
+        assert all(budget.take() for _ in range(100))
+        assert budget.num_emitted == 0
+
+    def test_no_budget_always_emits_detail(self) -> None:
+        """Without a budget object, every failing comparison gets full detail."""
+        x, y = self._failing_pair()
+        for _ in range(3):
+            info = compare_tensor_pair(x_baseline=x, x_target=y, name="fail")
+            assert info.diff is not None and 50 in info.diff.abs_diff_percentiles
+
+    def test_passing_tensor_skips_percentiles_and_budget(self) -> None:
+        """A passing comparison emits no percentile detail and consumes no budget."""
+        budget = FailureDisplayBudget()
+        x = torch.randn(4, 4)
+        info = compare_tensor_pair(
+            x_baseline=x, x_target=x.clone(), name="pass", failure_display_budget=budget
+        )
+        assert info.diff is not None and info.diff.passed is True
+        assert info.diff.abs_diff_percentiles == {}
+        assert info.baseline.stats.percentiles == {}
+        assert info.target.stats.percentiles == {}
+        assert budget.num_emitted == 0
+
+    def test_failing_tensor_within_budget_has_percentiles(self) -> None:
+        """A failing comparison within budget emits stats and diff percentiles."""
+        budget = FailureDisplayBudget()
+        x, y = self._failing_pair()
+        info = compare_tensor_pair(
+            x_baseline=x, x_target=y, name="fail", failure_display_budget=budget
+        )
+        assert info.diff is not None and info.diff.passed is False
+        assert 50 in info.diff.abs_diff_percentiles
+        assert 50 in info.baseline.stats.percentiles
+        assert 50 in info.target.stats.percentiles
+        assert budget.num_emitted == 1
+
+    def test_failing_tensor_beyond_budget_keeps_verdict_drops_detail(self) -> None:
+        """A failing comparison over budget keeps its verdict and metrics but drops percentiles."""
+        budget = FailureDisplayBudget(max_detail=1)
+        x, y = self._failing_pair()
+        first = compare_tensor_pair(
+            x_baseline=x, x_target=y, name="first", failure_display_budget=budget
+        )
+        second = compare_tensor_pair(
+            x_baseline=x, x_target=y, name="second", failure_display_budget=budget
+        )
+        assert first.diff is not None and 50 in first.diff.abs_diff_percentiles
+        assert second.diff is not None and second.diff.passed is False
+        assert second.diff.max_abs_diff == pytest.approx(1.0)
+        assert second.diff.abs_diff_percentiles == {}
+        assert second.baseline.stats.percentiles == {}
+        assert second.target.stats.percentiles == {}
+
+    def test_shape_mismatch_consumes_budget_and_emits_stats_detail(self) -> None:
+        """A shape mismatch counts as a failure and gets full stats percentile detail."""
+        budget = FailureDisplayBudget(max_detail=1)
+        info = compare_tensor_pair(
+            x_baseline=torch.randn(3, 4),
+            x_target=torch.randn(5, 6),
+            name="mismatch",
+            failure_display_budget=budget,
+        )
+        assert info.shape_mismatch is True and info.diff is None
+        assert 50 in info.baseline.stats.percentiles
+        assert 50 in info.target.stats.percentiles
+        assert budget.num_emitted == 1
+
+    def test_sample_still_emitted_beyond_budget(self) -> None:
+        """Sample emission for large diffs is independent of the detail budget."""
+        budget = FailureDisplayBudget(max_detail=0)
+        x, y = self._failing_pair()
+        info = compare_tensor_pair(
+            x_baseline=x, x_target=y, name="big", failure_display_budget=budget
+        )
+        assert info.baseline.sample is not None
+        assert info.target.sample is not None
+
+    def test_downcast_diff_detail_follows_budget(self) -> None:
+        """The downcast diff carries percentiles only when the failure is within budget."""
+        x = torch.zeros(4, 4, dtype=torch.float32)
+        y = torch.ones(4, 4, dtype=torch.bfloat16)
+        within = compare_tensor_pair(
+            x_baseline=x,
+            x_target=y,
+            name="within",
+            failure_display_budget=FailureDisplayBudget(),
+        )
+        assert within.diff_downcast is not None
+        assert 50 in within.diff_downcast.abs_diff_percentiles
+        beyond = compare_tensor_pair(
+            x_baseline=x,
+            x_target=y,
+            name="beyond",
+            failure_display_budget=FailureDisplayBudget(max_detail=0),
+        )
+        assert beyond.diff_downcast is not None
+        assert beyond.diff_downcast.abs_diff_percentiles == {}
+
+    def test_compute_diff_include_percentiles_flag(self) -> None:
+        """compute_diff with include_percentiles=False omits abs_diff_percentiles."""
+        x, y = self._failing_pair()
+        without = compute_diff(x_baseline=x, x_target=y, include_percentiles=False)
+        with_detail = compute_diff(x_baseline=x, x_target=y, include_percentiles=True)
+        assert without.abs_diff_percentiles == {}
+        assert 50 in with_detail.abs_diff_percentiles
+        assert without.passed == with_detail.passed
+        assert without.rel_diff == with_detail.rel_diff
+
+    def test_compute_tensor_info_include_percentiles_flag(self) -> None:
+        """compute_tensor_info with include_percentiles=False omits stats percentiles."""
+        t = torch.randn(16)
+        assert compute_tensor_info(t, include_percentiles=False).stats.percentiles == {}
+        assert 50 in compute_tensor_info(t, include_percentiles=True).stats.percentiles
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -5241,5 +5241,25 @@ class TestDiffThresholdPredicateExitCode:
         assert exit_code == 1
 
 
+class TestFailureDisplayBudgetWiring:
+    def test_run_emits_full_detail_for_failing_tensor(self, tmp_path, capsys) -> None:
+        """run() builds a fresh per-run budget, so failing tensors carry percentile detail."""
+        ones = torch.ones(4, 4)
+        baseline = _create_rank_dump(
+            tmp_path / "baseline", rank=0, name="g", tensor=ones
+        )
+        target = _create_rank_dump(
+            tmp_path / "target", rank=0, name="g", tensor=ones * 2
+        )
+
+        records, exit_code = _run_and_parse(_make_argv(baseline, target), capsys)
+
+        tensors = [r for r in records if isinstance(r, ComparisonTensorRecord)]
+        assert len(tensors) == 1
+        assert tensors[0].diff is not None and tensors[0].diff.passed is False
+        assert len(tensors[0].diff.abs_diff_percentiles) > 0
+        assert exit_code == 1
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))