Enhance replicated tensor checker in dump comparator (#19597)

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

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional
 
 import torch
@@ -21,9 +21,11 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
     execute_token_aligner,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
+    UnsharderResult,
     execute_unsharder_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
 from sglang.srt.debug_utils.comparator.utils import Pair
 
 
@@ -31,6 +33,7 @@ from sglang.srt.debug_utils.comparator.utils import Pair
 class AlignerResult:
     tensors: Optional[Pair[torch.Tensor]]
     failed_side_xy: Optional[str]  # "x" or "y"; None if success
+    replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
 
 
 def execute_aligner_plan(
@@ -39,18 +42,26 @@ def execute_aligner_plan(
     plan: AlignerPlan,
 ) -> AlignerResult:
     """Execute unified unshard/reorder + token-align."""
+    all_checks: list[ReplicatedCheckResult] = []
 
     # Per-side: unshard + reorder -> dict[step, tensor]
-    step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans(
+    step_tensors_x, checks_x = _execute_step_plans(
         tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
     )
-    step_tensors_y: dict[int, torch.Tensor] = _execute_step_plans(
+    all_checks.extend(checks_x)
+
+    step_tensors_y, checks_y = _execute_step_plans(
         tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
     )
+    all_checks.extend(checks_y)
 
     if not step_tensors_x or not step_tensors_y:
         failed_side_xy: str = "x" if not step_tensors_x else "y"
-        return AlignerResult(tensors=None, failed_side_xy=failed_side_xy)
+        return AlignerResult(
+            tensors=None,
+            failed_side_xy=failed_side_xy,
+            replicated_checks=all_checks,
+        )
 
     # Cross-side: token alignment (or direct extraction for single-step)
     if plan.token_aligner_plan is not None:
@@ -72,55 +83,64 @@ def execute_aligner_plan(
             y=execute_axis_aligner_plan(tensor=combined.y, plan=aligner_plan, side="y"),
         )
 
-    return AlignerResult(tensors=combined, failed_side_xy=None)
+    return AlignerResult(
+        tensors=combined,
+        failed_side_xy=None,
+        replicated_checks=all_checks,
+    )
 
 
 def _execute_step_plans(
     tensors: list[torch.Tensor],
     step_plans: list[AlignerPerStepPlan],
-) -> dict[int, torch.Tensor]:
+) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]:
     result: dict[int, torch.Tensor] = {}
+    all_checks: list[ReplicatedCheckResult] = []
 
     for step_plan in step_plans:
         step_tensors: list[torch.Tensor] = [
             tensors[i] for i in step_plan.input_object_indices
         ]
-        tensor: Optional[torch.Tensor] = execute_sub_plans(
+        tensor, checks = execute_sub_plans(
             tensors=step_tensors, plans=step_plan.sub_plans
         )
+        all_checks.extend(checks)
         if tensor is not None:
             result[step_plan.step] = tensor
 
-    return result
+    return result, all_checks
 
 
 def execute_sub_plans(
     tensors: list[torch.Tensor],
     plans: list[AlignerPerStepSubPlan],
-) -> Optional[torch.Tensor]:
+) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]:
     if not tensors:
-        return None
+        return None, []
 
     if not plans:
         if len(tensors) != 1:
-            return None
+            return None, []
-        return tensors[0]
+        return tensors[0], []
 
-    current = tensors
+    current: list[torch.Tensor] = tensors
+    all_checks: list[ReplicatedCheckResult] = []
     for plan in plans:
-        current = execute_sub_plan(tensors=current, plan=plan)
+        current, checks = execute_sub_plan(tensors=current, plan=plan)
+        all_checks.extend(checks)
 
     assert len(current) == 1
-    return current[0]
+    return current[0], all_checks
 
 
 def execute_sub_plan(
     tensors: list[torch.Tensor],
     plan: AlignerPerStepSubPlan,
-) -> list[torch.Tensor]:
+) -> tuple[list[torch.Tensor], list[ReplicatedCheckResult]]:
     if isinstance(plan, UnsharderPlan):
-        return execute_unsharder_plan(plan, tensors)
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, tensors)
+        return unsharder_result.tensors, unsharder_result.replicated_checks
     elif isinstance(plan, ReordererPlan):
-        return execute_reorderer_plan(plan, tensors)
+        return execute_reorderer_plan(plan, tensors), []
     else:
         raise NotImplementedError(f"Unknown {plan=}")

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py

@@ -236,7 +236,7 @@ def _load_and_align_aux_tensor(
             dim_names: list[str] = resolve_dim_names(dims_str)
             tensors = [apply_dim_names(t, dim_names) for t in tensors]
 
-        result = execute_sub_plans(tensors=tensors, plans=sub_plans)
+        result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans)
         assert result is not None
         return result.rename(None)  # strip named dims before returning to plugin
 

python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass, field
 from typing import Optional
 
 import torch
@@ -14,27 +15,37 @@ from sglang.srt.debug_utils.comparator.dims import (
     ParallelAxis,
     resolve_dim_by_name,
 )
-from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
+from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
-from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import compute_diff
+
+_REPLICATED_ATOL: float = 1e-6
+
+
+@dataclass(frozen=True)
+class UnsharderResult:
+    tensors: list[torch.Tensor]
+    replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
 
 
 def execute_unsharder_plan(
     plan: UnsharderPlan,
     tensors: list[torch.Tensor],
-) -> list[torch.Tensor]:
+) -> UnsharderResult:
-    result: list[torch.Tensor] = []
+    result_tensors: list[torch.Tensor] = []
+    all_checks: list[ReplicatedCheckResult] = []
 
     for group_idx, group in enumerate(plan.groups):
         group_tensors = [tensors[i] for i in group]
-        tensor = _apply_unshard(
+        tensor, checks = _apply_unshard(
             plan.params,
             group_tensors,
             axis=plan.axis,
             group_index=group_idx,
         )
-        result.append(tensor)
+        result_tensors.append(tensor)
+        all_checks.extend(checks)
 
-    return result
+    return UnsharderResult(tensors=result_tensors, replicated_checks=all_checks)
 
 
 def _apply_unshard(
@@ -43,25 +54,28 @@ def _apply_unshard(
     *,
     axis: ParallelAxis,
     group_index: int,
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, list[ReplicatedCheckResult]]:
     if isinstance(params, PickParams):
-        _verify_replicated_group(
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
             ordered_tensors,
             axis=axis,
             group_index=group_index,
         )
-        return ordered_tensors[0]
+        return ordered_tensors[0], checks
 
     if isinstance(params, ConcatParams):
         dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
-        return torch.cat(ordered_tensors, dim=dim)
+        return torch.cat(ordered_tensors, dim=dim), []
 
     if isinstance(params, CpThdConcatParams):
         thd_dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
-        return _thd_concat(
+        return (
+            _thd_concat(
                 ordered_tensors,
                 dim=thd_dim,
                 seq_lens_per_rank=params.seq_lens_per_rank,
+            ),
+            [],
         )
 
     if isinstance(params, ReduceSumParams):
@@ -70,7 +84,7 @@ def _apply_unshard(
         names: tuple[Optional[str], ...] = ordered_tensors[0].names
         if names[0] is not None:
             result = result.refine_names(*names)
-        return result
+        return result, []
 
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
 
@@ -80,22 +94,32 @@ def _verify_replicated_group(
     *,
     axis: ParallelAxis,
     group_index: int,
-) -> None:
+) -> list[ReplicatedCheckResult]:
-    baseline = ordered_tensors[0].rename(None)
+    baseline: torch.Tensor = ordered_tensors[0].rename(None).float()
+    checks: list[ReplicatedCheckResult] = []
 
     for i in range(1, len(ordered_tensors)):
-        other = ordered_tensors[i].rename(None)
+        other: torch.Tensor = ordered_tensors[i].rename(None).float()
-        if not torch.allclose(baseline, other, atol=1e-6):
+        diff_info = compute_diff(
-            warning_sink.add(
+            x_baseline=baseline,
-                ReplicatedMismatchWarning(
+            x_target=other,
+            diff_threshold=_REPLICATED_ATOL,
+        )
+        passed: bool = diff_info.max_abs_diff <= _REPLICATED_ATOL
+        checks.append(
+            ReplicatedCheckResult(
                 axis=axis.value,
                 group_index=group_index,
-                    differing_index=i,
+                compared_index=i,
                 baseline_index=0,
-                    max_abs_diff=(baseline - other).abs().max().item(),
+                passed=passed,
+                atol=_REPLICATED_ATOL,
+                diff=diff_info,
             )
         )
 
+    return checks
+
 
 def _thd_concat(
     ordered_tensors: list[torch.Tensor],

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

@@ -160,6 +160,7 @@ def _compare_bundle_pair_tensor_type(
     aligner_result: AlignerResult = execute_aligner_plan(
         tensors_pair=tensors_pair, plan=plan
     )
+    replicated_checks = aligner_result.replicated_checks
 
     if aligner_result.tensors is None:
         assert aligner_result.failed_side_xy is not None
@@ -183,7 +184,11 @@ def _compare_bundle_pair_tensor_type(
         diff_threshold=diff_threshold,
         seq_dim=seq_dim,
     )
-    record = ComparisonRecord(**info.model_dump(), aligner_plan=plan)
+    record = ComparisonRecord(
+        **info.model_dump(),
+        aligner_plan=plan,
+        replicated_checks=replicated_checks,
+    )
 
     if viz_output_dir is not None:
         _try_generate_viz(

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

@@ -8,8 +8,10 @@ from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_valida
 
 from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
     format_comparison,
+    format_replicated_checks,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
+    DiffInfo,
     TensorComparisonInfo,
 )
 from sglang.srt.debug_utils.comparator.utils import _StrictBase
@@ -20,22 +22,6 @@ if TYPE_CHECKING:
     )
 
 
-class ReplicatedMismatchWarning(_StrictBase):
-    kind: Literal["replicated_mismatch"] = "replicated_mismatch"
-    axis: str
-    group_index: int
-    differing_index: int
-    baseline_index: int
-    max_abs_diff: float
-
-    def to_text(self) -> str:
-        return (
-            f"Replicated along {self.axis}: group {self.group_index}, "
-            f"index {self.differing_index} differs from {self.baseline_index} "
-            f"(max_abs_diff={self.max_abs_diff:.6e})"
-        )
-
-
 class GeneralWarning(_StrictBase):
     kind: Literal["general"] = "general"
     category: str
@@ -45,10 +31,19 @@ class GeneralWarning(_StrictBase):
         return self.message
 
 
-AnyWarning = Annotated[
+# Type alias — currently only GeneralWarning exists.
-    Union[ReplicatedMismatchWarning, GeneralWarning],
+# When adding new warning types, convert back to Union + Discriminator("kind").
-    Discriminator("kind"),
+AnyWarning = GeneralWarning
-]
+
+
+class ReplicatedCheckResult(_StrictBase):
+    axis: str
+    group_index: int
+    compared_index: int
+    baseline_index: int
+    passed: bool
+    atol: float
+    diff: DiffInfo
 
 
 class _OutputRecord(_StrictBase):
@@ -126,15 +121,20 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
 
     type: Literal["comparison"] = "comparison"
     aligner_plan: Optional[AlignerPlan] = None
+    replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
 
     @property
     def category(self) -> str:
         if self.warnings:
             return "failed"
+        if any(not check.passed for check in self.replicated_checks):
+            return "failed"
         return "passed" if self.diff is not None and self.diff.passed else "failed"
 
     def _format_body(self) -> str:
         body: str = format_comparison(self)
+        if self.replicated_checks:
+            body += "\n" + format_replicated_checks(self.replicated_checks)
         if self.aligner_plan is not None:
             body += "\n" + _format_aligner_plan(self.aligner_plan)
         return body

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

@@ -57,7 +57,7 @@ def compare_tensor_pair(
     downcast_dtype: Optional[torch.dtype] = None
 
     if not shape_mismatch:
-        diff = _compute_diff(
+        diff = compute_diff(
             x_baseline=x_baseline_f,
             x_target=x_target_f,
             diff_threshold=diff_threshold,
@@ -74,7 +74,7 @@ def compare_tensor_pair(
                 Pair(x=baseline_original_dtype, y=target_original_dtype)
             )
             if downcast_dtype is not None:
-                diff_downcast = _compute_diff(
+                diff_downcast = compute_diff(
                     x_baseline=x_baseline_f.to(downcast_dtype),
                     x_target=x_target_f.to(downcast_dtype),
                     diff_threshold=diff_threshold,
@@ -121,7 +121,7 @@ def _compute_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
     return {p: torch.quantile(x_float, p / 100.0).item() for p in DEFAULT_PERCENTILES}
 
 
-def _compute_diff(
+def compute_diff(
     x_baseline: torch.Tensor,
     x_target: torch.Tensor,
     diff_threshold: float = 1e-3,

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

@@ -1,9 +1,16 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
     DiffInfo,
     TensorComparisonInfo,
     TensorStats,
 )
 
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
+
 
 def format_comparison(info: TensorComparisonInfo) -> str:
     lines: list[str] = []
@@ -54,6 +61,22 @@ def format_comparison(info: TensorComparisonInfo) -> str:
     return "\n".join(lines)
 
 
+def format_replicated_checks(checks: list[ReplicatedCheckResult]) -> str:
+    lines: list[str] = ["Replicated checks:"]
+
+    for check in checks:
+        marker: str = "✅" if check.passed else "❌"
+        lines.append(
+            f"  {marker} axis={check.axis} group={check.group_index} "
+            f"idx={check.compared_index} vs {check.baseline_index}: "
+            f"rel_diff={check.diff.rel_diff:.6e} "
+            f"max_abs_diff={check.diff.max_abs_diff:.6e} "
+            f"mean_abs_diff={check.diff.mean_abs_diff:.6e}"
+        )
+
+    return "\n".join(lines)
+
+
 def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]:
     lines: list[str] = []
 

test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py

@@ -1,5 +1,4 @@
 import sys
-from typing import Optional
 
 import pytest
 import torch
@@ -32,22 +31,25 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
 
 class TestExecuteSubPlans:
     def test_empty_tensors_returns_none(self) -> None:
-        result: Optional[torch.Tensor] = execute_sub_plans(tensors=[], plans=[])
+        result, checks = execute_sub_plans(tensors=[], plans=[])
         assert result is None
+        assert checks == []
 
     def test_no_plans_single_tensor_passthrough(self) -> None:
         tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
-        result: Optional[torch.Tensor] = execute_sub_plans(tensors=[tensor], plans=[])
+        result, checks = execute_sub_plans(tensors=[tensor], plans=[])
         assert result is not None
         assert torch.equal(result, tensor)
+        assert checks == []
 
     def test_no_plans_multiple_tensors_returns_none(self) -> None:
         tensors: list[torch.Tensor] = [
             torch.tensor([1.0]),
             torch.tensor([2.0]),
         ]
-        result: Optional[torch.Tensor] = execute_sub_plans(tensors=tensors, plans=[])
+        result, checks = execute_sub_plans(tensors=tensors, plans=[])
         assert result is None
+        assert checks == []
 
     def test_with_unsharder_plan(self) -> None:
         t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h")
@@ -59,13 +61,12 @@ class TestExecuteSubPlans:
             groups=[[0, 1]],
         )
 
-        result: Optional[torch.Tensor] = execute_sub_plans(
+        result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan])
-            tensors=[t0, t1], plans=[plan]
-        )
 
         assert result is not None
         expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
         assert torch.equal(result.rename(None), expected)
+        assert checks == []
 
 
 class TestExecuteSubPlan:
@@ -74,7 +75,7 @@ class TestExecuteSubPlan:
             pass
 
         with pytest.raises(NotImplementedError, match="Unknown"):
-            execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan())  # type: ignore[arg-type]
+            execute_sub_plan(tensors=[torch.tensor([1.0])], plan=_FakePlan())
 
 
 class TestExecuteStepPlans:
@@ -90,11 +91,10 @@ class TestExecuteStepPlans:
             sub_plans=[],
         )
 
-        result: dict[int, torch.Tensor] = _execute_step_plans(
+        result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan])
-            tensors=tensors, step_plans=[step_plan]
-        )
 
         assert result == {}
+        assert checks == []
 
     def test_single_step_passthrough(self) -> None:
         tensor: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -105,12 +105,11 @@ class TestExecuteStepPlans:
             sub_plans=[],
         )
 
-        result: dict[int, torch.Tensor] = _execute_step_plans(
+        result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan])
-            tensors=[tensor], step_plans=[step_plan]
-        )
 
         assert 5 in result
         assert torch.equal(result[5], tensor)
+        assert checks == []
 
 
 class TestExecuteAlignerPlan:

test/registered/debug_utils/comparator/aligner/reorderer/test_executor.py

@@ -20,7 +20,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     UnsharderPlan,
 )
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
-from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -227,10 +226,8 @@ class TestThdCpZigzagE2E:
             params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=seq_lens_per_rank),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
-            unsharded: list[torch.Tensor] = execute_unsharder_plan(
+        unsharded: list[torch.Tensor] = unsharder_result.tensors
-                unshard_plan, rank_tensors
-            )
         assert len(unsharded) == 1
 
         # Step 2: THD reorder
@@ -266,10 +263,8 @@ class TestThdCpZigzagE2E:
             ),
             groups=[list(range(cp_size))],
         )
-        with warning_sink.context():
+        unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
-            unsharded: list[torch.Tensor] = execute_unsharder_plan(
+        unsharded: list[torch.Tensor] = unsharder_result.tensors
-                unshard_plan, rank_tensors
-            )
         assert len(unsharded) == 1
 
         # Step 2: THD reorder

test/registered/debug_utils/comparator/aligner/reorderer/test_planner.py

@@ -18,7 +18,6 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
 from sglang.srt.debug_utils.comparator.dims import DimSpec, ParallelAxis, parse_dims
-from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -157,12 +156,11 @@ class TestCpZigzagTpE2E:
         assert len(reorderer_plans) == 1
 
         current: list[torch.Tensor] = [t.refine_names(*dim_names) for t in tensors]
-        with warning_sink.context():
         for plan in all_plans:
             if isinstance(plan, ReordererPlan):
                 current = execute_reorderer_plan(plan, current)
             else:
-                    current = execute_unsharder_plan(plan, current)
+                current = execute_unsharder_plan(plan, current).tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)

test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py

@@ -4,6 +4,7 @@ import pytest
 import torch
 
 from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
+    UnsharderResult,
     _apply_unshard,
     _verify_replicated_group,
     execute_unsharder_plan,
@@ -23,7 +24,7 @@ from sglang.srt.debug_utils.comparator.dims import (
     ParallelAxis,
     parse_dims,
 )
-from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -49,11 +50,12 @@ class TestExecuteUnsharderPlan:
         assert len(plans) == 1
 
         named_shards: list[torch.Tensor] = _name_tensors(shards, dim_specs)
-        with warning_sink.context() as warnings:
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], named_shards)
+            plans[0], named_shards
-        assert len(result) == 1
+        )
-        assert torch.allclose(result[0].rename(None), full_tensor)
+        assert len(unsharder_result.tensors) == 1
-        assert warnings == []
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
+        assert unsharder_result.replicated_checks == []
 
     def test_scrambled_world_ranks_correct_result(self) -> None:
         full_tensor = torch.randn(4, 8)
@@ -79,11 +81,12 @@ class TestExecuteUnsharderPlan:
             dim_specs,
         )
 
-        with warning_sink.context() as warnings:
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], tensors_ordered_by_world_rank)
+            plans[0], tensors_ordered_by_world_rank
-        assert len(result) == 1
+        )
-        assert torch.allclose(result[0].rename(None), full_tensor)
+        assert len(unsharder_result.tensors) == 1
-        assert warnings == []
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
+        assert unsharder_result.replicated_checks == []
 
     def test_single_step_reduces_tensor_count(self) -> None:
         """8 tensors with 2 groups of 4 produce 2 output tensors."""
@@ -113,13 +116,15 @@ class TestExecuteUnsharderPlan:
                 tensors.append(source[tp_rank])
 
         named_tensors: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
+        intermediate_result: UnsharderResult = execute_unsharder_plan(
-            intermediate = execute_unsharder_plan(plans[0], named_tensors)
+            plans[0], named_tensors
-        assert len(intermediate) == 4
+        )
+        assert len(intermediate_result.tensors) == 4
 
-        with warning_sink.context():
+        final_result: UnsharderResult = execute_unsharder_plan(
-            final = execute_unsharder_plan(plans[1], intermediate)
+            plans[1], intermediate_result.tensors
-        assert len(final) == 1
+        )
+        assert len(final_result.tensors) == 1
 
     def test_cp_tp_concat(self) -> None:
         """CP=2 + TP=2: multi-step unshard reconstructs original tensor."""
@@ -145,9 +150,9 @@ class TestExecuteUnsharderPlan:
         assert len(plans) == 2
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -187,9 +192,9 @@ class TestExecuteUnsharderPlan:
         assert len(plans) == 2
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -241,9 +246,9 @@ class TestExecuteUnsharderPlan:
         assert len(plans) == 3
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -290,9 +295,9 @@ class TestExecuteUnsharderPlan:
         assert len(plans) == 3
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -312,11 +317,12 @@ class TestPickOperation:
         assert len(plans) == 1
         assert isinstance(plans[0].params, PickParams)
 
-        with warning_sink.context() as warnings:
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], [tensor, tensor.clone()])
+            plans[0], [tensor, tensor.clone()]
-        assert len(result) == 1
+        )
-        assert torch.allclose(result[0].rename(None), tensor)
+        assert len(unsharder_result.tensors) == 1
-        assert warnings == []
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor)
+        assert all(c.passed for c in unsharder_result.replicated_checks)
 
     def test_pick_multiple_groups(self) -> None:
         """PickParams with multiple groups picks one from each."""
@@ -348,10 +354,11 @@ class TestPickOperation:
         tensor = torch.randn(4)
         tensors = [tensor.clone() for _ in range(4)]
 
-        with warning_sink.context() as warnings:
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(pick_plans[0], tensors)
+            pick_plans[0], tensors
-        assert len(result) == 2
+        )
-        assert warnings == []
+        assert len(unsharder_result.tensors) == 2
+        assert all(c.passed for c in unsharder_result.replicated_checks)
 
     def test_replicated_tp_sharded_cp_e2e(self) -> None:
         """CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip."""
@@ -376,9 +383,9 @@ class TestPickOperation:
         assert len(plans) == 2
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -406,44 +413,44 @@ class TestPickOperation:
         assert all(isinstance(p.params, PickParams) for p in plans)
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
 
 
 class TestVerifyReplicatedGroup:
-    def test_warns_on_mismatch(self) -> None:
+    def test_fails_on_mismatch(self) -> None:
-        """_verify_replicated_group produces warning when replicas differ."""
+        """_verify_replicated_group returns failed check when replicas differ."""
         tensor_a = torch.ones(4)
         tensor_b = torch.ones(4) + 0.1
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [tensor_a, tensor_b],
             axis=ParallelAxis.TP,
             group_index=0,
         )
-        assert len(warnings) == 1
+        assert len(checks) == 1
-        assert warnings[0].axis == "tp"
+        assert checks[0].axis == "tp"
-        assert warnings[0].group_index == 0
+        assert checks[0].group_index == 0
-        assert warnings[0].differing_index == 1
+        assert checks[0].compared_index == 1
-        assert warnings[0].baseline_index == 0
+        assert checks[0].baseline_index == 0
-        assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
+        assert not checks[0].passed
+        assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
 
-    def test_no_warn_when_identical(self) -> None:
+    def test_passes_when_identical(self) -> None:
-        """_verify_replicated_group produces no warning for identical replicas."""
+        """_verify_replicated_group returns passed check for identical replicas."""
         tensor = torch.randn(4, 8)
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [tensor, tensor.clone()],
             axis=ParallelAxis.TP,
             group_index=0,
         )
-        assert warnings == []
+        assert len(checks) == 1
+        assert checks[0].passed
 
     def test_multiple_mismatches(self) -> None:
         """_verify_replicated_group reports each differing replica."""
@@ -451,19 +458,20 @@ class TestVerifyReplicatedGroup:
         other_a = torch.ones(4)
         other_b = torch.ones(4) * 2
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [baseline, other_a, other_b],
             axis=ParallelAxis.CP,
             group_index=1,
         )
-        assert len(warnings) == 2
+        assert len(checks) == 2
-        assert warnings[0].differing_index == 1
+        assert checks[0].compared_index == 1
-        assert warnings[1].differing_index == 2
+        assert not checks[0].passed
-        assert warnings[1].max_abs_diff == pytest.approx(2.0, abs=1e-5)
+        assert checks[1].compared_index == 2
+        assert not checks[1].passed
+        assert checks[1].diff.max_abs_diff == pytest.approx(2.0, abs=1e-5)
 
-    def test_execute_returns_warnings(self) -> None:
+    def test_execute_returns_replicated_checks(self) -> None:
-        """execute_unsharder_plan emits warnings for replicated mismatch."""
+        """execute_unsharder_plan returns replicated checks for mismatch."""
         dim_specs = parse_dims("h d")
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
@@ -474,56 +482,58 @@ class TestVerifyReplicatedGroup:
         tensor_a = torch.zeros(4)
         tensor_b = torch.ones(4)
 
-        with warning_sink.context() as warnings:
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], [tensor_a, tensor_b])
+            plans[0], [tensor_a, tensor_b]
-        assert len(result) == 1
+        )
-        assert len(warnings) == 1
+        assert len(unsharder_result.tensors) == 1
-        assert torch.allclose(result[0].rename(None), tensor_a)
+        assert len(unsharder_result.replicated_checks) == 1
+        assert not unsharder_result.replicated_checks[0].passed
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), tensor_a)
 
     def test_atol_boundary_within(self) -> None:
-        """Difference exactly at atol (1e-6) -> torch.allclose passes -> no warning."""
+        """Difference exactly at atol (1e-6) -> passed."""
         baseline = torch.zeros(4)
         other = torch.full((4,), 1e-6)
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [baseline, other],
             axis=ParallelAxis.TP,
             group_index=0,
         )
-        assert warnings == []
+        assert len(checks) == 1
+        assert checks[0].passed
 
     def test_atol_boundary_exceeded(self) -> None:
-        """Difference just above atol (1e-6 + 1e-9) -> torch.allclose fails -> warning."""
+        """Difference just above atol (1e-6 + 1e-9) -> failed."""
         baseline = torch.zeros(4)
         other = torch.full((4,), 1e-6 + 1e-9)
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [baseline, other],
             axis=ParallelAxis.TP,
             group_index=0,
         )
-        assert len(warnings) == 1
+        assert len(checks) == 1
-        assert warnings[0].differing_index == 1
+        assert not checks[0].passed
+        assert checks[0].compared_index == 1
 
-    def test_recompute_pseudo_mismatch_warns(self) -> None:
+    def test_recompute_pseudo_mismatch(self) -> None:
-        """_verify_replicated_group produces warning for RECOMPUTE_PSEUDO axis mismatch."""
+        """_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
         tensor_a = torch.ones(4)
         tensor_b = torch.ones(4) + 0.1
 
-        with warning_sink.context() as warnings:
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
-            _verify_replicated_group(
             [tensor_a, tensor_b],
             axis=ParallelAxis.RECOMPUTE_PSEUDO,
             group_index=0,
         )
-        assert len(warnings) == 1
+        assert len(checks) == 1
-        assert warnings[0].axis == "recompute_pseudo"
+        assert checks[0].axis == "recompute_pseudo"
-        assert warnings[0].group_index == 0
+        assert checks[0].group_index == 0
-        assert warnings[0].differing_index == 1
+        assert checks[0].compared_index == 1
-        assert warnings[0].baseline_index == 0
+        assert checks[0].baseline_index == 0
-        assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
+        assert not checks[0].passed
+        assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
 
 
 class TestThdCpConcat:
@@ -537,12 +547,11 @@ class TestThdCpConcat:
             params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
-            result = execute_unsharder_plan(plan, [rank0, rank1])
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
         expected = torch.tensor([1, 2, 3, 4, 5, 6])
-        assert torch.equal(result[0].rename(None), expected)
+        assert torch.equal(unsharder_result.tensors[0].rename(None), expected)
 
     def test_multi_seq(self) -> None:
         """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
@@ -563,11 +572,10 @@ class TestThdCpConcat:
             params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
-            result = execute_unsharder_plan(plan, [rank0, rank1])
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
+        unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None)
 
         # seqA: r0(50) + r1(50) = 100 tokens, values 0..99
         assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
@@ -595,11 +603,10 @@ class TestThdCpConcat:
             params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
-            result = execute_unsharder_plan(plan, [rank0, rank1])
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
+        unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None)
 
         assert unsharded.shape == (10, hidden)
         assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
@@ -625,11 +632,10 @@ class TestThdCpConcat:
             params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
-            result = execute_unsharder_plan(plan, [rank0, rank1])
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
+        unsharded: torch.Tensor = unsharder_result.tensors[0].rename(None)
 
         assert unsharded.shape == (batch, 10, hidden)
         # seqA: r0(3) + r1(3) = 6 tokens per batch
@@ -657,11 +663,12 @@ class TestReduceSum:
         assert isinstance(plans[0].params, ReduceSumParams)
 
         named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], named_parts)
+            plans[0], named_parts
+        )
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        assert torch.allclose(result[0].rename(None), full_tensor)
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
 
     def test_tp4_reduce(self) -> None:
         """4 partial tensors sum to full tensor."""
@@ -677,11 +684,12 @@ class TestReduceSum:
         assert len(plans) == 1
 
         named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], named_parts)
+            plans[0], named_parts
+        )
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        assert torch.allclose(result[0].rename(None), full_tensor)
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
 
     def test_multi_axis_concat_then_reduce(self) -> None:
         """CP concat + TP reduce end-to-end."""
@@ -707,9 +715,9 @@ class TestReduceSum:
         assert len(plans) == 2
 
         current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
-        with warning_sink.context():
         for plan in plans:
-                current = execute_unsharder_plan(plan, current)
+            unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
+            current = unsharder_result.tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)
@@ -735,11 +743,12 @@ class TestReduceSum:
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
 
         named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plans[0], named_parts)
+            plans[0], named_parts
+        )
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        assert torch.allclose(result[0].rename(None), full_tensor)
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
 
     def test_reduce_preserves_named_dims(self) -> None:
         """Named tensor dimensions are preserved through reduce_sum."""
@@ -752,13 +761,16 @@ class TestReduceSum:
             params=ReduceSumParams(),
             groups=[[0, 1]],
         )
-        with warning_sink.context():
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
-            result = execute_unsharder_plan(plan, [part_a, part_b])
+            plan, [part_a, part_b]
+        )
 
-        assert len(result) == 1
+        assert len(unsharder_result.tensors) == 1
-        assert result[0].names == ("h", "d")
+        assert unsharder_result.tensors[0].names == ("h", "d")
         expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d")
-        assert torch.allclose(result[0].rename(None), expected.rename(None))
+        assert torch.allclose(
+            unsharder_result.tensors[0].rename(None), expected.rename(None)
+        )
 
 
 if __name__ == "__main__":

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

@@ -6,9 +6,9 @@ import torch
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
     QUANTILE_NUMEL_THRESHOLD,
     SAMPLE_DIFF_THRESHOLD,
-    _compute_diff,
     _compute_tensor_stats,
     compare_tensor_pair,
+    compute_diff,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -55,7 +55,7 @@ class TestComputeTensorStats:
 class TestComputeDiff:
     def test_identical_tensors(self):
         x = torch.ones(10, 10)
-        diff = _compute_diff(x_baseline=x, x_target=x)
+        diff = compute_diff(x_baseline=x, x_target=x)
 
         assert diff.rel_diff == pytest.approx(0.0, abs=1e-5)
         assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
@@ -70,7 +70,7 @@ class TestComputeDiff:
         y = x.clone()
         y[3, 7] = 1.5
 
-        diff = _compute_diff(x_baseline=x, x_target=y)
+        diff = compute_diff(x_baseline=x, x_target=y)
 
         assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4)
         assert diff.max_diff_coord == [3, 7]
@@ -85,14 +85,14 @@ class TestComputeDiff:
     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)
+        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])
-        diff = _compute_diff(x_baseline=x, x_target=y)
+        diff = compute_diff(x_baseline=x, x_target=y)
 
         assert diff.rel_diff == pytest.approx(1.0, abs=1e-5)
         assert diff.passed is False
@@ -103,7 +103,7 @@ class TestComputeDiff:
         x: torch.Tensor = torch.randn(8, 16)
         y: torch.Tensor = x + torch.randn_like(x) * 0.01
 
-        diff: DiffInfo = _compute_diff(
+        diff: DiffInfo = compute_diff(
             x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0
         )
 
@@ -117,7 +117,7 @@ class TestComputeDiff:
         x: torch.Tensor = torch.randn(8, 16)
         y: torch.Tensor = x + torch.randn_like(x) * 0.01
 
-        diff: DiffInfo = _compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3)
+        diff: DiffInfo = compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3)
 
         assert diff.per_token_rel_diff is None
 
@@ -127,7 +127,7 @@ class TestComputeDiff:
         x: torch.Tensor = torch.randn(4, 8)
         y: torch.Tensor = x + torch.randn_like(x) * 0.01
 
-        diff: DiffInfo = _compute_diff(
+        diff: DiffInfo = compute_diff(
             x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0
         )
 

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

@@ -7,7 +7,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     ConfigRecord,
     GeneralWarning,
-    ReplicatedMismatchWarning,
+    ReplicatedCheckResult,
     SkipRecord,
     SummaryRecord,
     WarningRecord,
@@ -126,16 +126,24 @@ class TestRecordTypes:
             assert restored == record
 
 
-def _make_warning(**overrides) -> ReplicatedMismatchWarning:
+def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
     defaults: dict = dict(
         axis="tp",
         group_index=0,
-        differing_index=1,
+        compared_index=1,
         baseline_index=0,
+        passed=False,
+        atol=1e-6,
+        diff=_make_diff(
+            rel_diff=0.1,
             max_abs_diff=0.1,
+            mean_abs_diff=0.05,
+            diff_threshold=1e-6,
+            passed=False,
+        ),
     )
     defaults.update(overrides)
-    return ReplicatedMismatchWarning(**defaults)
+    return ReplicatedCheckResult(**defaults)
 
 
 class TestWarnings:
@@ -148,7 +156,7 @@ class TestWarnings:
             unified_shape=[4, 8],
             shape_mismatch=False,
             diff=_make_diff(passed=True),
-            warnings=[_make_warning()],
+            warnings=[GeneralWarning(category="test", message="some warning")],
         )
         assert record.category == "failed"
 
@@ -157,18 +165,44 @@ class TestWarnings:
         record = SkipRecord(
             name="x",
             reason="no_baseline",
-            warnings=[_make_warning()],
+            warnings=[GeneralWarning(category="test", message="some warning")],
         )
         assert record.category == "failed"
 
-    def test_warnings_json_round_trip(self):
+    def test_replicated_checks_all_passed(self):
-        """warnings survive model_dump_json → parse_record_json round-trip."""
+        """ComparisonRecord with all replicated_checks passed → category=='passed'."""
-        warning = _make_warning(
+        record = ComparisonRecord(
+            name="hidden",
+            baseline=_make_tensor_info(),
+            target=_make_tensor_info(),
+            unified_shape=[4, 8],
+            shape_mismatch=False,
+            diff=_make_diff(passed=True),
+            replicated_checks=[_make_replicated_check(passed=True)],
+        )
+        assert record.category == "passed"
+
+    def test_replicated_checks_failed_means_record_failed(self):
+        """ComparisonRecord with any replicated_check.passed=False → category=='failed'."""
+        record = ComparisonRecord(
+            name="hidden",
+            baseline=_make_tensor_info(),
+            target=_make_tensor_info(),
+            unified_shape=[4, 8],
+            shape_mismatch=False,
+            diff=_make_diff(passed=True),
+            replicated_checks=[_make_replicated_check(passed=False)],
+        )
+        assert record.category == "failed"
+
+    def test_replicated_check_json_round_trip(self):
+        """ReplicatedCheckResult survives JSON round-trip via ComparisonRecord."""
+        check = _make_replicated_check(
             axis="cp",
             group_index=2,
-            differing_index=3,
+            compared_index=3,
             baseline_index=0,
-            max_abs_diff=0.42,
+            passed=False,
         )
         record = ComparisonRecord(
             name="mlp",
@@ -177,30 +211,23 @@ class TestWarnings:
             unified_shape=[4, 8],
             shape_mismatch=False,
             diff=_make_diff(),
-            warnings=[warning],
+            replicated_checks=[check],
         )
 
         restored = parse_record_json(record.model_dump_json())
         assert isinstance(restored, ComparisonRecord)
-        assert len(restored.warnings) == 1
+        assert len(restored.replicated_checks) == 1
 
-        restored_warning = restored.warnings[0]
+        restored_check: ReplicatedCheckResult = restored.replicated_checks[0]
-        assert restored_warning.axis == "cp"
+        assert restored_check.axis == "cp"
-        assert restored_warning.group_index == 2
+        assert restored_check.group_index == 2
-        assert restored_warning.differing_index == 3
+        assert restored_check.compared_index == 3
-        assert restored_warning.baseline_index == 0
+        assert restored_check.baseline_index == 0
-        assert restored_warning.max_abs_diff == pytest.approx(0.42)
+        assert not restored_check.passed
 
     def test_any_warning_discriminated_union_round_trip(self):
         """All AnyWarning variants survive JSON round-trip via a WarningRecord."""
         all_warnings = [
-            ReplicatedMismatchWarning(
-                axis="tp",
-                group_index=0,
-                differing_index=1,
-                baseline_index=0,
-                max_abs_diff=0.1,
-            ),
             GeneralWarning(
                 category="aux_tensors_missing",
                 message="Aux tensors missing, skipping token alignment",

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -13,7 +13,7 @@ from sglang.srt.debug_utils.comparator.output_types import (
     ConfigRecord,
     GeneralWarning,
     NonTensorRecord,
-    ReplicatedMismatchWarning,
+    ReplicatedCheckResult,
     SkipRecord,
     SummaryRecord,
     WarningRecord,
@@ -910,7 +910,7 @@ class TestEntrypointGroupingLogical:
         assert comp.name == "hidden"
 
     def test_recompute_pseudo_mismatch_warning(self, tmp_path, capsys):
-        """Recompute pseudo-axis with differing original/recompute → ReplicatedMismatchWarning."""
+        """Recompute pseudo-axis with differing original/recompute → failed replicated_checks."""
         torch.manual_seed(42)
         tensor = torch.randn(4, 8)
         mismatched_tensor = tensor + torch.randn(4, 8) * 10.0
@@ -937,12 +937,112 @@ class TestEntrypointGroupingLogical:
         comparisons = _get_comparisons(records)
         assert len(comparisons) == 1
 
-        recompute_warnings = [
+        recompute_checks: list[ReplicatedCheckResult] = [
-            w
+            c for c in comparisons[0].replicated_checks if c.axis == "recompute_pseudo"
-            for w in comparisons[0].warnings
-            if isinstance(w, ReplicatedMismatchWarning) and w.axis == "recompute_pseudo"
         ]
-        assert len(recompute_warnings) > 0
+        assert len(recompute_checks) > 0
+        assert any(not c.passed for c in recompute_checks)
+
+    def test_tp_partial_reduction_unshard(self, tmp_path, capsys):
+        """TP=2 with partial reduction: element-wise sum reconstructs full tensor."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8)
+        full_target = full_baseline + torch.randn(4, 8) * 0.001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        baseline_path = _create_tp_partial_dumps(
+            baseline_dir,
+            full_tensor=full_baseline,
+            name="attn_out",
+            tp_size=2,
+            dims_str="b h(tp,partial)",
+        )
+        target_path = _create_tp_partial_dumps(
+            target_dir,
+            full_tensor=full_target,
+            name="attn_out",
+            tp_size=2,
+            dims_str="b h(tp,partial)",
+        )
+
+        args = _make_args(baseline_path, target_path, diff_threshold=0.01)
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "attn_out"
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.total == 1
+        assert summary.passed == 1
+
+    def test_tp_partial_vs_single_rank(self, tmp_path, capsys):
+        """Baseline single rank vs target TP=2 partial: unshard target then compare."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+        target_full = full_tensor + torch.randn(4, 8) * 0.001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        baseline_path = _create_rank_dump(
+            baseline_dir, rank=0, name="attn_out", tensor=full_tensor
+        )
+        target_path = _create_tp_partial_dumps(
+            target_dir,
+            full_tensor=target_full,
+            name="attn_out",
+            tp_size=2,
+            dims_str="b h(tp,partial)",
+        )
+
+        args = _make_args(baseline_path, target_path, diff_threshold=0.01)
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "attn_out"
+
+    def test_cp_concat_tp_partial_reduction(self, tmp_path, capsys):
+        """CP=2 concat + TP=2 partial reduction: multi-axis unshard."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 16)
+        full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
+
+        for side_dir, full_tensor in [
+            (tmp_path / "baseline", full_baseline),
+            (tmp_path / "target", full_target),
+        ]:
+            side_dir.mkdir()
+            cp_chunks = list(full_tensor.chunk(2, dim=1))
+            rank = 0
+            for cp_rank in range(2):
+                for tp_rank in range(2):
+                    _create_rank_dump(
+                        side_dir,
+                        rank=rank,
+                        name="hidden",
+                        tensor=cp_chunks[cp_rank] / 2,
+                        dims="b s(cp) h(tp,partial)",
+                        parallel_info={
+                            "cp_rank": cp_rank,
+                            "cp_size": 2,
+                            "tp_rank": tp_rank,
+                            "tp_size": 2,
+                        },
+                    )
+                    rank += 1
+
+        args = _make_args(
+            tmp_path / "baseline" / _FIXED_EXP_NAME,
+            tmp_path / "target" / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
 
     def test_tp_partial_reduction_unshard(self, tmp_path, capsys):
         """TP=2 with partial reduction: element-wise sum reconstructs full tensor."""
@@ -1233,7 +1333,7 @@ class TestEntrypointReplicatedAxis:
     """Test replicated-axis scenarios through the full entrypoint pipeline."""
 
     def test_replicated_axis_identical_replicas_passed(self, tmp_path, capsys):
-        """CP2 TP2, TP replicated and identical → passed, no warnings."""
+        """CP2 TP2, TP replicated and identical → passed, replicated_checks all passed."""
         torch.manual_seed(42)
         full_baseline = torch.randn(4, 8, 6)
         full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1264,13 +1364,14 @@ class TestEntrypointReplicatedAxis:
         records = _run_and_parse(args, capsys)
         comp = _assert_single_comparison_passed(records)
         assert comp.warnings == []
+        assert all(c.passed for c in comp.replicated_checks)
 
         summary = records[-1]
         assert isinstance(summary, SummaryRecord)
         assert summary.passed == 1
 
     def test_replicated_mismatch_fails(self, tmp_path, capsys):
-        """CP2 TP2, TP replicas differ (> atol) → failed with warnings."""
+        """CP2 TP2, TP replicas differ (> atol) → failed with replicated_checks."""
         torch.manual_seed(42)
         full_baseline = torch.randn(4, 8, 6)
         full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1303,14 +1404,14 @@ class TestEntrypointReplicatedAxis:
         comparisons = _get_comparisons(records)
         assert len(comparisons) == 1
         assert comparisons[0].category == "failed"
-        assert len(comparisons[0].warnings) > 0
+        assert any(not c.passed for c in comparisons[0].replicated_checks)
 
         summary = records[-1]
         assert isinstance(summary, SummaryRecord)
         assert summary.failed == 1
 
-    def test_summary_counts_failed_from_warnings_only(self, tmp_path, capsys):
+    def test_summary_counts_failed_from_replicated_checks_only(self, tmp_path, capsys):
-        """Diff itself passes but TP replicas differ → summary.failed=1 from warnings."""
+        """Diff itself passes but TP replicas differ → summary.failed=1 from replicated_checks."""
         torch.manual_seed(42)
         full_baseline = torch.randn(4, 8, 6)
         full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
@@ -1352,7 +1453,7 @@ class TestEntrypointReplicatedAxis:
         comp = comparisons[0]
         assert comp.diff is not None
         assert comp.diff.passed
-        assert len(comp.warnings) > 0
+        assert any(not c.passed for c in comp.replicated_checks)
         assert comp.category == "failed"
 
         summary = records[-1]

test/registered/debug_utils/comparator/test_warning_sink.py

@@ -3,23 +3,20 @@ import sys
 
 import pytest
 
-from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
+from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 
-def _make_warning(**overrides) -> ReplicatedMismatchWarning:
+def _make_warning(**overrides) -> GeneralWarning:
     defaults: dict = dict(
-        axis="tp",
+        category="test",
-        group_index=0,
+        message="test warning",
-        differing_index=1,
-        baseline_index=0,
-        max_abs_diff=0.1,
     )
     defaults.update(overrides)
-    return ReplicatedMismatchWarning(**defaults)
+    return GeneralWarning(**defaults)
 
 
 class TestWarningSink:
@@ -35,8 +32,8 @@ class TestWarningSink:
 
     def test_nested_contexts(self) -> None:
         sink = WarningSink()
-        outer_warning = _make_warning(group_index=0)
+        outer_warning = _make_warning(message="outer")
-        inner_warning = _make_warning(group_index=1)
+        inner_warning = _make_warning(message="inner")
 
         with sink.context() as outer:
             sink.add(outer_warning)
@@ -61,7 +58,7 @@ class TestWarningSink:
         sink.add(_make_warning())
 
         captured = capsys.readouterr()
-        assert "Replicated along tp" in captured.out
+        assert "test warning" in captured.out
 
     def test_context_captures_instead_of_printing(self, capsys) -> None:
         sink = WarningSink()
@@ -96,9 +93,9 @@ class TestWarningSink:
 
         assert len(collected) == 1
 
-        sink.add(_make_warning(group_index=99))
+        sink.add(_make_warning(message="after exception"))
         captured = capsys.readouterr()
-        assert "Replicated along tp" in captured.out
+        assert "after exception" in captured.out
 
 
 if __name__ == "__main__":