Support replication axis in dump comparator (#19282)

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

@@ -2,30 +2,80 @@ import torch
 
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     ConcatParams,
+    PickParams,
     UnshardParams,
     UnshardPlan,
 )
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.output_types import (
+    AlignWarning,
+    ReplicatedMismatchWarning,
+)
 
 
 def execute_unshard_plan(
     plan: UnshardPlan,
     tensors: list[torch.Tensor],
-) -> list[torch.Tensor]:
+) -> tuple[list[torch.Tensor], list[AlignWarning]]:
+    all_warnings: list[AlignWarning] = []
     result: list[torch.Tensor] = []
-    for group in plan.groups:
+
+    for group_idx, group in enumerate(plan.groups):
         group_tensors = [tensors[i] for i in group]
-        result.append(_apply_unshard(plan.params, group_tensors))
+        tensor, warnings = _apply_unshard(
-    return result
+            plan.params,
+            group_tensors,
+            axis=plan.axis,
+            group_index=group_idx,
+        )
+        result.append(tensor)
+        all_warnings.extend(warnings)
+
+    return result, all_warnings
 
 
 def _apply_unshard(
-    params: UnshardParams, ordered_tensors: list[torch.Tensor]
+    params: UnshardParams,
-) -> torch.Tensor:
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> tuple[torch.Tensor, list[AlignWarning]]:
+    if isinstance(params, PickParams):
+        warnings = _verify_replicated_group(
+            ordered_tensors,
+            axis=axis,
+            group_index=group_index,
+        )
+        return ordered_tensors[0], warnings
+
     if isinstance(params, ConcatParams):
-        return _unshard_concat(ordered_tensors, dim=params.dim)
+        return torch.cat(ordered_tensors, dim=params.dim), []
+
     # Phase 2: ReduceSumParams, CpZigzagParams
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
 
 
-def _unshard_concat(tensors: list[torch.Tensor], dim: int) -> torch.Tensor:
+def _verify_replicated_group(
-    return torch.cat(tensors, dim=dim)
+    ordered_tensors: list[torch.Tensor],
+    *,
+    axis: ParallelAxis,
+    group_index: int,
+) -> list[ReplicatedMismatchWarning]:
+    warnings: list[ReplicatedMismatchWarning] = []
+    baseline = ordered_tensors[0]
+
+    for i in range(1, len(ordered_tensors)):
+        other = ordered_tensors[i]
+        if not torch.allclose(baseline, other, atol=1e-6):
+            warnings.append(
+                ReplicatedMismatchWarning(
+                    axis=axis.value,
+                    group_index=group_index,
+                    differing_index=i,
+                    baseline_index=0,
+                    max_abs_diff=(baseline - other).abs().max().item(),
+                )
+            )
+
+    return warnings

python/sglang/srt/debug_utils/comparator/aligner/unshard/planner.py

@@ -4,6 +4,7 @@ from typing import NamedTuple
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
     AxisInfo,
     ConcatParams,
+    PickParams,
     UnshardParams,
     UnshardPlan,
 )
@@ -32,31 +33,38 @@ def compute_unshard_plan(
         for dim_idx, spec in enumerate(dim_specs)
         if spec.parallel is not None
     }
-    if not sharded_axis_infos:
+    sharded_axes: set[ParallelAxis] = set(sharded_axis_infos)
+
+    all_axes: set[ParallelAxis] = {axis for info in parallel_infos for axis in info}
+    replicated_axes: set[ParallelAxis] = all_axes - sharded_axes
+
+    if not sharded_axes and not replicated_axes:
         return []
 
-    _validate(sharded_axes=set(sharded_axis_infos), parallel_infos=parallel_infos)
+    _validate(
+        axes_to_validate=sharded_axes | replicated_axes,
+        parallel_infos=parallel_infos,
+    )
 
     current_coords: _CoordsList = [
-        {axis: info[axis].axis_rank for axis in sharded_axis_infos}
+        {axis: info[axis].axis_rank for axis in sharded_axes | replicated_axes}
         for info in parallel_infos
     ]
 
+    axis_and_params: list[tuple[ParallelAxis, UnshardParams]] = [
+        (axis, PickParams()) for axis in sorted(replicated_axes, key=lambda a: a.value)
+    ] + [
+        (axis, _resolve_unshard_params(spec=spec, dim_index=dim_index))
+        for axis, (dim_index, spec) in sharded_axis_infos.items()
+    ]
+
     plans: list[UnshardPlan] = []
-    for axis, (dim_index, spec) in sharded_axis_infos.items():
+    for axis, params in axis_and_params:
         result = _group_and_project(
             current_coords=current_coords,
             target_axis=axis,
         )
-
+        plans.append(UnshardPlan(axis=axis, params=params, groups=result.groups))
-        plans.append(
-            UnshardPlan(
-                axis=axis,
-                params=_resolve_unshard_params(spec=spec, dim_index=dim_index),
-                groups=result.groups,
-            )
-        )
-
         current_coords = result.projected_coords
 
     return plans
@@ -64,18 +72,18 @@ def compute_unshard_plan(
 
 def _validate(
     *,
-    sharded_axes: set[ParallelAxis],
+    axes_to_validate: set[ParallelAxis],
     parallel_infos: list[dict[ParallelAxis, AxisInfo]],
 ) -> None:
-    """Check that every rank has all sharded axes, sizes are consistent, and ranks are complete."""
+    """Check that every rank has all axes, sizes are consistent, and ranks are complete."""
     axis_sizes: dict[ParallelAxis, int] = {}
 
     for world_rank, parallel_info in enumerate(parallel_infos):
-        for axis in sharded_axes:
+        for axis in axes_to_validate:
             if axis not in parallel_info:
                 raise ValueError(
                     f"world_rank={world_rank} missing parallel_info for "
-                    f"sharded axis {axis.value!r}"
+                    f"axis {axis.value!r}"
                 )
 
             axis_info = parallel_info[axis]

python/sglang/srt/debug_utils/comparator/aligner/unshard/types.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import Literal
+from typing import Annotated, Literal, Union
+
+from pydantic import Field
 
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 from sglang.srt.debug_utils.comparator.utils import _FrozenBase
@@ -16,7 +18,14 @@ class ConcatParams(_FrozenBase):
     dim: int
 
 
-UnshardParams = ConcatParams
+class PickParams(_FrozenBase):
+    op: Literal["pick"] = "pick"
+
+
+UnshardParams = Annotated[
+    Union[ConcatParams, PickParams],
+    Field(discriminator="op"),
+]
 
 
 class UnshardPlan(_FrozenBase):

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

@@ -1,7 +1,7 @@
 from abc import abstractmethod
 from typing import Annotated, Literal, Union
 
-from pydantic import Discriminator, TypeAdapter
+from pydantic import Discriminator, Field, TypeAdapter
 
 from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
     format_comparison,
@@ -12,9 +12,38 @@ from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
 from sglang.srt.debug_utils.comparator.utils import _StrictBase
 
 
+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})"
+        )
+
+
+AlignWarning = (
+    ReplicatedMismatchWarning  # future: Annotated[Union[...], Discriminator("kind")]
+)
+
+
 class _OutputRecord(_StrictBase):
+    align_warnings: list[AlignWarning] = Field(default_factory=list)
+
     @abstractmethod
-    def to_text(self) -> str: ...
+    def _format_body(self) -> str: ...
+
+    def to_text(self) -> str:
+        body = self._format_body()
+        if self.align_warnings:
+            body += "\n" + "\n".join(f"  ⚠ {w.to_text()}" for w in self.align_warnings)
+        return body
 
 
 class ConfigRecord(_OutputRecord):
@@ -25,7 +54,7 @@ class ConfigRecord(_OutputRecord):
     start_step: int
     end_step: int
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return (
             f"Config: baseline={self.baseline_path} target={self.target_path}\n"
             f"diff_threshold={self.diff_threshold} "
@@ -39,10 +68,12 @@ class SkipRecord(_OutputRecord):
     reason: str
 
     @property
-    def category(self):
+    def category(self) -> str:
+        if self.align_warnings:
+            return "failed"
         return "skipped"
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return f"Skip: {self.name} ({self.reason})"
 
 
@@ -50,10 +81,12 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
     type: Literal["comparison"] = "comparison"
 
     @property
-    def category(self):
+    def category(self) -> str:
+        if self.align_warnings:
+            return "failed"
         return "passed" if self.diff is not None and self.diff.passed else "failed"
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return format_comparison(self)
 
 
@@ -64,7 +97,7 @@ class SummaryRecord(_OutputRecord):
     failed: int
     skipped: int
 
-    def to_text(self) -> str:
+    def _format_body(self) -> str:
         return (
             f"Summary: {self.passed} passed, {self.failed} failed, "
             f"{self.skipped} skipped (total {self.total})"

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

@@ -20,6 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
 from sglang.srt.debug_utils.comparator.dims import parse_dims
 from sglang.srt.debug_utils.comparator.output_types import (
+    AlignWarning,
     ComparisonRecord,
     SkipRecord,
 )
@@ -50,12 +51,13 @@ def process_tensor_group(
     t_extracted = _extract_tensors(t_tensors)
     del b_tensors, t_tensors
 
-    b_tensor = _execute_plans(b_extracted, b_plans)
+    b_tensor, b_warns = _execute_plans(b_extracted, b_plans)
-    t_tensor = _execute_plans(t_extracted, t_plans)
+    t_tensor, t_warns = _execute_plans(t_extracted, t_plans)
+    all_warnings: list[AlignWarning] = b_warns + t_warns
 
     if b_tensor is None or t_tensor is None:
         reason = "baseline_load_failed" if b_tensor is None else "target_load_failed"
-        return SkipRecord(name=name, reason=reason)
+        return SkipRecord(name=name, reason=reason, align_warnings=all_warnings)
 
     info = compare_tensors(
         x_baseline=b_tensor,
@@ -64,7 +66,7 @@ def process_tensor_group(
         diff_threshold=diff_threshold,
     )
 
-    return ComparisonRecord(**info.model_dump())
+    return ComparisonRecord(**info.model_dump(), align_warnings=all_warnings)
 
 
 def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
@@ -112,27 +114,32 @@ def _extract_tensors(
 def _execute_plans(
     tensors: list[torch.Tensor],
     plans: list[Plan],
-) -> Optional[torch.Tensor]:
+) -> tuple[Optional[torch.Tensor], list[AlignWarning]]:
     if not tensors:
-        return None
+        return None, []
 
     if not plans:
         if len(tensors) != 1:
-            return None
+            return None, []
-        return tensors[0]
+        return tensors[0], []
 
+    warnings: list[AlignWarning] = []
     current = tensors
     for plan in plans:
-        current = _execute_plan(current, plan)
+        current, new_warnings = _execute_plan(current, plan)
+        warnings.extend(new_warnings)
 
     assert len(current) == 1
-    return current[0]
+    return current[0], warnings
 
 
-def _execute_plan(tensors, plan):
+def _execute_plan(
+    tensors: list[torch.Tensor],
+    plan: Plan,
+) -> tuple[list[torch.Tensor], list[AlignWarning]]:
     if isinstance(plan, UnshardPlan):
         return execute_unshard_plan(plan, tensors)
     elif isinstance(plan, ReorderPlan):
-        return execute_reorder_plan(plan, tensors)
+        return execute_reorder_plan(plan, tensors), []
     else:
         raise NotImplementedError(f"Unknown {plan=}")

test/registered/debug_utils/comparator/aligner/test_reorder.py

@@ -148,7 +148,7 @@ class TestCpZigzagTpE2E:
             if isinstance(plan, ReorderPlan):
                 current = execute_reorder_plan(plan, current)
             else:
-                current = execute_unshard_plan(plan, current)
+                current, _ = execute_unshard_plan(plan, current)
 
         assert len(current) == 1
         assert torch.allclose(current[0], full_tensor)

test/registered/debug_utils/comparator/aligner/unshard/test_execute.py

@@ -5,12 +5,16 @@ import torch
 
 from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
     _apply_unshard,
+    _verify_replicated_group,
     execute_unshard_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
     compute_unshard_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
+    AxisInfo,
+    PickParams,
+)
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -29,9 +33,10 @@ class TestExecuteUnshardPlan:
         plans = compute_unshard_plan(dim_specs, parallel_infos)
         assert len(plans) == 1
 
-        result = execute_unshard_plan(plans[0], shards)
+        result, warnings = execute_unshard_plan(plans[0], shards)
         assert len(result) == 1
         assert torch.allclose(result[0], full_tensor)
+        assert warnings == []
 
     def test_scrambled_world_ranks_correct_result(self) -> None:
         full_tensor = torch.randn(4, 8)
@@ -54,9 +59,10 @@ class TestExecuteUnshardPlan:
             shards[1],  # world_rank=3, axis_rank=1
         ]
 
-        result = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
+        result, warnings = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
         assert len(result) == 1
         assert torch.allclose(result[0], full_tensor)
+        assert warnings == []
 
     def test_single_step_reduces_tensor_count(self) -> None:
         """8 tensors with 2 groups of 4 produce 2 output tensors."""
@@ -85,10 +91,10 @@ class TestExecuteUnshardPlan:
             for tp_rank in range(4):
                 tensors.append(source[tp_rank])
 
-        intermediate = execute_unshard_plan(plans[0], tensors)
+        intermediate, _ = execute_unshard_plan(plans[0], tensors)
         assert len(intermediate) == 4
 
-        final = execute_unshard_plan(plans[1], intermediate)
+        final, _ = execute_unshard_plan(plans[1], intermediate)
         assert len(final) == 1
 
     def test_cp_tp_concat(self) -> None:
@@ -116,7 +122,7 @@ class TestExecuteUnshardPlan:
 
         current = tensors
         for plan in plans:
-            current = execute_unshard_plan(plan, current)
+            current, _ = execute_unshard_plan(plan, current)
 
         assert len(current) == 1
         assert torch.allclose(current[0], full_tensor)
@@ -157,7 +163,7 @@ class TestExecuteUnshardPlan:
 
         current = tensors
         for plan in plans:
-            current = execute_unshard_plan(plan, current)
+            current, _ = execute_unshard_plan(plan, current)
 
         assert len(current) == 1
         assert torch.allclose(current[0], full_tensor)
@@ -169,7 +175,12 @@ class TestExecuteUnshardPlan:
             pass
 
         with pytest.raises(ValueError, match="Unsupported unshard"):
-            _apply_unshard(_FakeParams(), [torch.randn(2, 2)])
+            _apply_unshard(
+                _FakeParams(),
+                [torch.randn(2, 2)],
+                axis=ParallelAxis.TP,
+                group_index=0,
+            )
 
     def test_cp_tp_ep_three_axis_concat(self) -> None:
         """CP=2 + TP=2 + EP=2: three-step unshard reconstructs original tensor."""
@@ -205,7 +216,7 @@ class TestExecuteUnshardPlan:
 
         current = tensors
         for plan in plans:
-            current = execute_unshard_plan(plan, current)
+            current, _ = execute_unshard_plan(plan, current)
 
         assert len(current) == 1
         assert torch.allclose(current[0], full_tensor)
@@ -253,11 +264,211 @@ class TestExecuteUnshardPlan:
 
         current = tensors
         for plan in plans:
-            current = execute_unshard_plan(plan, current)
+            current, _ = execute_unshard_plan(plan, current)
 
         assert len(current) == 1
         assert torch.allclose(current[0], full_tensor)
 
 
+class TestPickOperation:
+    def test_pick_single_group(self) -> None:
+        """PickParams picks the first tensor from a single group."""
+        tensor = torch.randn(4, 8)
+        dim_specs = parse_dims("h d")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
+        ]
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, PickParams)
+
+        result, warnings = execute_unshard_plan(plans[0], [tensor, tensor.clone()])
+        assert len(result) == 1
+        assert torch.allclose(result[0], tensor)
+        assert warnings == []
+
+    def test_pick_multiple_groups(self) -> None:
+        """PickParams with multiple groups picks one from each."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
+        assert len(pick_plans) == 1
+        assert pick_plans[0].axis == ParallelAxis.CP
+
+        tensor = torch.randn(4)
+        tensors = [tensor.clone() for _ in range(4)]
+
+        result, warnings = execute_unshard_plan(pick_plans[0], tensors)
+        assert len(result) == 2
+        assert warnings == []
+
+    def test_replicated_tp_sharded_cp_e2e(self) -> None:
+        """CP2 TP2, dims='b s(cp) d': replicated TP pick + sharded CP concat round-trip."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+        cp_chunks = list(full_tensor.chunk(2, dim=1))
+
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for tp_rank in range(2):
+                tensors.append(cp_chunks[cp_rank].clone())
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                    }
+                )
+
+        dim_specs = parse_dims("b s(cp) d")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        current = tensors
+        for plan in plans:
+            current, _ = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+    def test_fully_replicated_e2e(self) -> None:
+        """CP2 TP2, dims='b h d': fully replicated → 2 pick steps → 1 tensor."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for tp_rank in range(2):
+                tensors.append(full_tensor.clone())
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                    }
+                )
+
+        dim_specs = parse_dims("b h d")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+        assert all(isinstance(p.params, PickParams) for p in plans)
+
+        current = tensors
+        for plan in plans:
+            current, _ = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+
+class TestVerifyReplicatedGroup:
+    def test_warns_on_mismatch(self) -> None:
+        """_verify_replicated_group produces warning when replicas differ."""
+        tensor_a = torch.ones(4)
+        tensor_b = torch.ones(4) + 0.1
+
+        warnings = _verify_replicated_group(
+            [tensor_a, tensor_b],
+            axis=ParallelAxis.TP,
+            group_index=0,
+        )
+        assert len(warnings) == 1
+        assert warnings[0].axis == "tp"
+        assert warnings[0].group_index == 0
+        assert warnings[0].differing_index == 1
+        assert warnings[0].baseline_index == 0
+        assert warnings[0].max_abs_diff == pytest.approx(0.1, abs=1e-5)
+
+    def test_no_warn_when_identical(self) -> None:
+        """_verify_replicated_group produces no warning for identical replicas."""
+        tensor = torch.randn(4, 8)
+
+        warnings = _verify_replicated_group(
+            [tensor, tensor.clone()],
+            axis=ParallelAxis.TP,
+            group_index=0,
+        )
+        assert warnings == []
+
+    def test_multiple_mismatches(self) -> None:
+        """_verify_replicated_group reports each differing replica."""
+        baseline = torch.zeros(4)
+        other_a = torch.ones(4)
+        other_b = torch.ones(4) * 2
+
+        warnings = _verify_replicated_group(
+            [baseline, other_a, other_b],
+            axis=ParallelAxis.CP,
+            group_index=1,
+        )
+        assert len(warnings) == 2
+        assert warnings[0].differing_index == 1
+        assert warnings[1].differing_index == 2
+        assert warnings[1].max_abs_diff == pytest.approx(2.0, abs=1e-5)
+
+    def test_execute_returns_warnings(self) -> None:
+        """execute_unshard_plan returns warnings for replicated mismatch."""
+        dim_specs = parse_dims("h d")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        tensor_a = torch.zeros(4)
+        tensor_b = torch.ones(4)
+
+        result, warnings = execute_unshard_plan(plans[0], [tensor_a, tensor_b])
+        assert len(result) == 1
+        assert len(warnings) == 1
+        assert torch.allclose(result[0], tensor_a)
+
+    def test_atol_boundary_within(self) -> None:
+        """Difference exactly at atol (1e-6) → torch.allclose passes → no warning."""
+        baseline = torch.zeros(4)
+        other = torch.full((4,), 1e-6)
+
+        warnings = _verify_replicated_group(
+            [baseline, other],
+            axis=ParallelAxis.TP,
+            group_index=0,
+        )
+        assert warnings == []
+
+    def test_atol_boundary_exceeded(self) -> None:
+        """Difference just above atol (1e-6 + 1e-9) → torch.allclose fails → warning."""
+        baseline = torch.zeros(4)
+        other = torch.full((4,), 1e-6 + 1e-9)
+
+        warnings = _verify_replicated_group(
+            [baseline, other],
+            axis=ParallelAxis.TP,
+            group_index=0,
+        )
+        assert len(warnings) == 1
+        assert warnings[0].differing_index == 1
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/aligner/unshard/test_plan.py

@@ -5,7 +5,11 @@ import pytest
 from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
     compute_unshard_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
+    AxisInfo,
+    ConcatParams,
+    PickParams,
+)
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -228,7 +232,7 @@ class TestComputeUnshardPlan:
         assert len(plans[2].groups) == 1
         assert len(plans[2].groups[0]) == 2
 
-    def test_replicated_axis_raises(self) -> None:
+    def test_sharded_axis_missing_from_rank_raises(self) -> None:
         """A world_rank missing a sharded axis raises ValueError."""
         dim_specs = parse_dims("s(cp) h(tp)")
         parallel_infos = [
@@ -238,7 +242,159 @@ class TestComputeUnshardPlan:
             },
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
-                # missing TP — replicated
+                # missing TP — sharded axis absent from rank
+            },
+        ]
+        with pytest.raises(ValueError, match="missing parallel_info"):
+            compute_unshard_plan(dim_specs, parallel_infos)
+
+
+class TestReplicatedAxes:
+    def test_replicated_tp_with_sharded_cp(self) -> None:
+        """CP2 TP2, dims='b s(cp) d' → PickPlan(TP) + ConcatPlan(CP)."""
+        dim_specs = parse_dims("b s(cp) d")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.TP
+        assert isinstance(plans[0].params, PickParams)
+        assert len(plans[0].groups) == 2
+        for group in plans[0].groups:
+            assert len(group) == 2
+
+        assert plans[1].axis == ParallelAxis.CP
+        assert isinstance(plans[1].params, ConcatParams)
+        assert plans[1].params.dim == 1
+
+    def test_fully_replicated(self) -> None:
+        """CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP)."""
+        dim_specs = parse_dims("b h d")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+        assert all(isinstance(p.params, PickParams) for p in plans)
+        axes = {p.axis for p in plans}
+        assert axes == {ParallelAxis.CP, ParallelAxis.TP}
+
+    def test_multiple_replicated_one_sharded(self) -> None:
+        """CP2 TP2 EP2, dims='h(tp)' → PickPlan(CP) + PickPlan(EP) + ConcatPlan(TP)."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for ep_rank in range(2):
+                for tp_rank in range(2):
+                    parallel_infos.append(
+                        {
+                            ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                            ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
+                            ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                        }
+                    )
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 3
+        pick_plans = [p for p in plans if isinstance(p.params, PickParams)]
+        concat_plans = [p for p in plans if isinstance(p.params, ConcatParams)]
+        assert len(pick_plans) == 2
+        assert len(concat_plans) == 1
+        assert concat_plans[0].axis == ParallelAxis.TP
+
+        replicated_axes = {p.axis for p in pick_plans}
+        assert replicated_axes == {ParallelAxis.CP, ParallelAxis.EP}
+
+    def test_replicated_scrambled_ranks(self) -> None:
+        """Scrambled world_rank order with replicated axis."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.CP
+        assert isinstance(plans[0].params, PickParams)
+        assert plans[1].axis == ParallelAxis.TP
+        assert isinstance(plans[1].params, ConcatParams)
+
+    def test_replicated_axis_inconsistent_size_raises(self) -> None:
+        """Replicated axis with inconsistent sizes raises ValueError."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=4),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+        ]
+        with pytest.raises(ValueError, match="Inconsistent axis_size"):
+            compute_unshard_plan(dim_specs, parallel_infos)
+
+    def test_replicated_axis_missing_from_rank_raises(self) -> None:
+        """A rank missing a replicated axis that other ranks have raises ValueError."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                # missing CP — replicated axis absent from this rank
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
             },
         ]
         with pytest.raises(ValueError, match="missing parallel_info"):

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

@@ -6,6 +6,7 @@ import pytest
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     ConfigRecord,
+    ReplicatedMismatchWarning,
     SkipRecord,
     SummaryRecord,
     parse_record_json,
@@ -118,5 +119,71 @@ class TestRecordTypes:
             assert restored == record
 
 
+def _make_warning(**overrides) -> ReplicatedMismatchWarning:
+    defaults: dict = dict(
+        axis="tp",
+        group_index=0,
+        differing_index=1,
+        baseline_index=0,
+        max_abs_diff=0.1,
+    )
+    defaults.update(overrides)
+    return ReplicatedMismatchWarning(**defaults)
+
+
+class TestAlignWarnings:
+    def test_comparison_record_failed_when_diff_passed_but_warnings(self):
+        """ComparisonRecord with diff.passed=True but align_warnings → 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),
+            align_warnings=[_make_warning()],
+        )
+        assert record.category == "failed"
+
+    def test_skip_record_failed_when_warnings(self):
+        """SkipRecord with align_warnings → category=='failed' instead of 'skipped'."""
+        record = SkipRecord(
+            name="x",
+            reason="no_baseline",
+            align_warnings=[_make_warning()],
+        )
+        assert record.category == "failed"
+
+    def test_align_warnings_json_round_trip(self):
+        """align_warnings survive model_dump_json → parse_record_json round-trip."""
+        warning = _make_warning(
+            axis="cp",
+            group_index=2,
+            differing_index=3,
+            baseline_index=0,
+            max_abs_diff=0.42,
+        )
+        record = ComparisonRecord(
+            name="mlp",
+            baseline=_make_tensor_info(),
+            target=_make_tensor_info(),
+            unified_shape=[4, 8],
+            shape_mismatch=False,
+            diff=_make_diff(),
+            align_warnings=[warning],
+        )
+
+        restored = parse_record_json(record.model_dump_json())
+        assert isinstance(restored, ComparisonRecord)
+        assert len(restored.align_warnings) == 1
+
+        restored_warning = restored.align_warnings[0]
+        assert restored_warning.axis == "cp"
+        assert restored_warning.group_index == 2
+        assert restored_warning.differing_index == 3
+        assert restored_warning.baseline_index == 0
+        assert restored_warning.max_abs_diff == pytest.approx(0.42)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -878,6 +878,138 @@ class TestEntrypointGroupingLogical:
         assert comp.name == "hidden"
 
 
+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 align_warnings."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 6)
+        full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        for side_dir, full_tensor in [
+            (baseline_dir, full_baseline),
+            (target_dir, full_target),
+        ]:
+            _create_replicated_tp_sharded_cp_dumps(
+                side_dir,
+                full_tensor=full_tensor,
+                name="attn_out",
+                cp_size=2,
+                tp_size=2,
+                seq_dim=1,
+                dims_str="b s(cp) d",
+            )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.align_warnings == []
+
+        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 align_warnings."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 6)
+        full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        for side_dir, full_tensor in [
+            (baseline_dir, full_baseline),
+            (target_dir, full_target),
+        ]:
+            _create_replicated_tp_sharded_cp_dumps(
+                side_dir,
+                full_tensor=full_tensor,
+                name="attn_out",
+                cp_size=2,
+                tp_size=2,
+                seq_dim=1,
+                dims_str="b s(cp) d",
+                tp_noise=0.5,
+            )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        assert comparisons[0].category == "failed"
+        assert len(comparisons[0].align_warnings) > 0
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.failed == 1
+
+    def test_summary_counts_failed_from_align_warnings_only(self, tmp_path, capsys):
+        """Diff itself passes but TP replicas differ → summary.failed=1 from align_warnings."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 6)
+        full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_replicated_tp_sharded_cp_dumps(
+            baseline_dir,
+            full_tensor=full_baseline,
+            name="attn_out",
+            cp_size=2,
+            tp_size=2,
+            seq_dim=1,
+            dims_str="b s(cp) d",
+            tp_noise=0.5,
+        )
+        _create_replicated_tp_sharded_cp_dumps(
+            target_dir,
+            full_tensor=full_target,
+            name="attn_out",
+            cp_size=2,
+            tp_size=2,
+            seq_dim=1,
+            dims_str="b s(cp) d",
+            tp_noise=0.5,
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.5,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+
+        comp = comparisons[0]
+        assert comp.diff is not None
+        assert comp.diff.passed
+        assert len(comp.align_warnings) > 0
+        assert comp.category == "failed"
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.failed == 1
+        assert summary.passed == 0
+
+
 # --------------------------- Assertion helpers -------------------
 
 
@@ -1138,6 +1270,49 @@ def _create_cp_zigzag_tp_sharded_dumps(
     return directory / _FIXED_EXP_NAME
 
 
+def _create_replicated_tp_sharded_cp_dumps(
+    directory: Path,
+    *,
+    full_tensor: torch.Tensor,
+    name: str,
+    cp_size: int,
+    tp_size: int,
+    seq_dim: int,
+    dims_str: str,
+    tp_noise: float = 0.0,
+) -> Path:
+    """Create CP-sharded + TP-replicated dump files from a full tensor.
+
+    CP direction: chunks along seq_dim (sharded).
+    TP direction: clones (replicated), with optional noise to simulate mismatch.
+    """
+    cp_chunks: list[torch.Tensor] = list(full_tensor.chunk(cp_size, dim=seq_dim))
+
+    rank: int = 0
+    for cp_rank in range(cp_size):
+        for tp_rank in range(tp_size):
+            shard = cp_chunks[cp_rank].clone()
+            if tp_noise > 0 and tp_rank > 0:
+                shard = shard + torch.randn_like(shard) * tp_noise
+
+            _create_rank_dump(
+                directory,
+                rank=rank,
+                name=name,
+                tensor=shard,
+                dims=dims_str,
+                parallel_info={
+                    "cp_rank": cp_rank,
+                    "cp_size": cp_size,
+                    "tp_rank": tp_rank,
+                    "tp_size": tp_size,
+                },
+            )
+            rank += 1
+
+    return directory / _FIXED_EXP_NAME
+
+
 def _create_tp_sharded_dumps(
     directory: Path,
     *,