Enhance replicated tensor checker in dump comparator (#19597)

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():
-            unsharded: list[torch.Tensor] = execute_unsharder_plan(
-                unshard_plan, rank_tensors
-            )
+        unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
+        unsharded: list[torch.Tensor] = unsharder_result.tensors
         assert len(unsharded) == 1
 
         # Step 2: THD reorder
@@ -266,10 +263,8 @@ class TestThdCpZigzagE2E:
             ),
             groups=[list(range(cp_size))],
         )
-        with warning_sink.context():
-            unsharded: list[torch.Tensor] = execute_unsharder_plan(
-                unshard_plan, rank_tensors
-            )
+        unsharder_result = execute_unsharder_plan(unshard_plan, rank_tensors)
+        unsharded: list[torch.Tensor] = unsharder_result.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)
+        for plan in all_plans:
+            if isinstance(plan, ReordererPlan):
+                current = execute_reorderer_plan(plan, current)
+            else:
+                current = execute_unsharder_plan(plan, current).tensors
 
         assert len(current) == 1
         assert torch.allclose(current[0].rename(None), full_tensor)