Support singleton dimension squeezing in dump comparator (#19566)

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

@@ -0,0 +1,165 @@
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
+    AxisAlignerPlan,
+    compute_axis_aligner_plan,
+    execute_axis_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+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=15, suite="default", nightly=True)
+
+
+class TestComputeAxisAlignerPlan:
+    def test_no_dims_returns_none(self) -> None:
+        assert compute_axis_aligner_plan(Pair(x=None, y=None)) is None
+        assert compute_axis_aligner_plan(Pair(x="t h d", y=None)) is None
+        assert compute_axis_aligner_plan(Pair(x=None, y="t h d")) is None
+
+    def test_same_order_returns_none(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h d", y="t h d")
+        )
+        assert result is None
+
+    def test_different_order(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h d", y="t d h")
+        )
+        assert result is not None
+        assert result.pattern.x == "t h d -> t d h"
+        assert result.pattern.y is None
+
+    def test_name_mismatch_returns_none_with_warning(self) -> None:
+        with warning_sink.context() as warnings:
+            result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+                Pair(x="t h d", y="t h e")
+            )
+
+        assert result is None
+        assert len(warnings) == 1
+        assert warnings[0].category == "axis_aligner_dim_mismatch"
+        assert "dim name sets differ" in warnings[0].message
+
+    def test_modifiers_ignored_for_name_extraction(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h(tp) d", y="t d h(tp)")
+        )
+        assert result is not None
+        assert result.pattern.x == "t h d -> t d h"
+
+    def test_squeeze_only_no_swap(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t 1 h", y="t h")
+        )
+        assert result is not None
+        assert result.pattern.x == "t 1 h -> t h"
+        assert result.pattern.y is None
+
+    def test_squeeze_both_sides(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t 1 h", y="1 t h")
+        )
+        assert result is not None
+        assert result.pattern.x == "t 1 h -> t h"
+        assert result.pattern.y == "1 t h -> t h"
+
+    def test_squeeze_plus_swap(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t 1 h d", y="t d h")
+        )
+        assert result is not None
+        assert result.pattern.x == "t 1 h d -> t d h"
+        assert result.pattern.y is None
+
+    def test_squeeze_y_only(self) -> None:
+        result: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
+            Pair(x="t h", y="t 1 h")
+        )
+        assert result is not None
+        assert result.pattern.x is None
+        assert result.pattern.y == "t 1 h -> t h"
+
+
+class TestExecuteAxisAlignerPlan:
+    def test_rearrange(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 8, 16).refine_names("t", "h", "d")
+        plan = AxisAlignerPlan(
+            pattern=Pair(x="t h d -> t d h", y=None),
+        )
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="x"
+        )
+
+        assert result.shape == (4, 16, 8)
+        for i in range(4):
+            assert torch.equal(
+                result[i],
+                tensor.rename(None)[i].T,
+            )
+
+    def test_execute_squeeze(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 1, 8).refine_names("t", "singleton0", "h")
+        plan = AxisAlignerPlan(
+            pattern=Pair(x="t 1 h -> t h", y=None),
+        )
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="x"
+        )
+
+        assert result.shape == (4, 8)
+
+    def test_execute_squeeze_then_swap(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 1, 8, 16).refine_names(
+            "t", "singleton0", "h", "d"
+        )
+        plan = AxisAlignerPlan(
+            pattern=Pair(x="t 1 h d -> t d h", y=None),
+        )
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="x"
+        )
+
+        assert result.shape == (4, 16, 8)
+
+    def test_execute_y_side(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 1, 8).refine_names("t", "singleton0", "h")
+        plan = AxisAlignerPlan(
+            pattern=Pair(x=None, y="t 1 h -> t h"),
+        )
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="y"
+        )
+
+        assert result.shape == (4, 8)
+
+    def test_noop_side(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 8, 16).refine_names("t", "h", "d")
+        plan = AxisAlignerPlan(
+            pattern=Pair(x="t h d -> t d h", y=None),
+        )
+
+        result: torch.Tensor = execute_axis_aligner_plan(
+            tensor=tensor, plan=plan, side="y"
+        )
+
+        assert result.shape == (4, 8, 16)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

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

@@ -639,5 +639,119 @@ class TestThdCpConcat:
         )
 
 
+class TestThdCpConcat:
+    def test_single_seq(self) -> None:
+        """Single seq THD unshard: 2 ranks → per-seq concat."""
+        rank0 = torch.tensor([1, 2, 3]).refine_names("t")
+        rank1 = torch.tensor([4, 5, 6]).refine_names("t")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        expected = torch.tensor([1, 2, 3, 4, 5, 6])
+        assert torch.equal(result[0].rename(None), expected)
+
+    def test_multi_seq(self) -> None:
+        """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
+        # rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
+        # rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
+        seq_a_r0 = torch.arange(0, 50)
+        seq_b_r0 = torch.arange(100, 132)
+        pad_r0 = torch.full((46,), -1)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
+
+        seq_a_r1 = torch.arange(50, 100)
+        seq_b_r1 = torch.arange(132, 164)
+        pad_r1 = torch.full((46,), -2)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[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]))
+        # seqB: r0(32) + r1(32) = 64 tokens
+        assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
+        # pad: r0(46) + r1(46) = 92 tokens
+        assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
+
+    def test_with_hidden_dim(self) -> None:
+        """THD unshard with trailing hidden dim: shape [T, H]."""
+        torch.manual_seed(42)
+        hidden: int = 4
+        # rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
+        # rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
+        seq_a_r0 = torch.randn(3, hidden)
+        seq_b_r0 = torch.randn(2, hidden)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
+
+        seq_a_r1 = torch.randn(3, hidden)
+        seq_b_r1 = torch.randn(2, hidden)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[0].rename(None)
+
+        assert unsharded.shape == (10, hidden)
+        assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
+        assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
+
+    def test_with_leading_batch_dim(self) -> None:
+        """THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
+        torch.manual_seed(42)
+        batch: int = 2
+        hidden: int = 4
+        # rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
+        # rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
+        seq_a_r0 = torch.randn(batch, 3, hidden)
+        seq_b_r0 = torch.randn(batch, 2, hidden)
+        rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
+
+        seq_a_r1 = torch.randn(batch, 3, hidden)
+        seq_b_r1 = torch.randn(batch, 2, hidden)
+        rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.CP,
+            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
+            groups=[[0, 1]],
+        )
+        with warning_sink.context():
+            result = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(result) == 1
+        unsharded: torch.Tensor = result[0].rename(None)
+
+        assert unsharded.shape == (batch, 10, hidden)
+        # seqA: r0(3) + r1(3) = 6 tokens per batch
+        assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
+        # seqB: r0(2) + r1(2) = 4 tokens per batch
+        assert torch.equal(
+            unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
+        )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_dims.py

@@ -6,17 +6,20 @@ import torch
 from sglang.srt.debug_utils.comparator.dims import (
     BATCH_DIM_NAME,
     SEQ_DIM_NAME,
+    SQUEEZE_DIM_NAME,
     TOKEN_DIM_NAME,
     DimSpec,
     Ordering,
     ParallelAxis,
     Reduction,
+    _SingletonDimUtil,
     apply_dim_names,
     find_dim_index,
     parse_dim,
     parse_dim_names,
     parse_dims,
     resolve_dim_by_name,
+    resolve_dim_names,
     strip_dim_names,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -72,6 +75,13 @@ class TestParseDim:
         with pytest.raises(ValueError, match="Multiple reduction"):
             parse_dim("h(tp,partial,partial)")
 
+    def test_squeeze_dim(self) -> None:
+        assert parse_dim("1") == DimSpec(name="1")
+
+    def test_squeeze_dim_rejects_modifiers(self) -> None:
+        with pytest.raises(ValueError, match="Invalid dim token"):
+            parse_dim("1(tp)")
+
 
 class TestParseDims:
     def test_multi_dims(self) -> None:
@@ -105,6 +115,169 @@ class TestParseDims:
         with pytest.raises(ValueError, match="Duplicate"):
             parse_dims("h h")
 
+    def test_with_squeeze_dims(self) -> None:
+        result: list[DimSpec] = parse_dims("t 1 h")
+        assert len(result) == 3
+        assert result[0] == DimSpec(name="t")
+        assert result[1] == DimSpec(name="1")
+        assert result[2] == DimSpec(name="h")
+
+    def test_multiple_squeeze_dims_no_duplicate_error(self) -> None:
+        result: list[DimSpec] = parse_dims("t 1 h 1 d")
+        assert len(result) == 5
+        assert result[1] == DimSpec(name="1")
+        assert result[3] == DimSpec(name="1")
+
+
+class TestDimConstants:
+    def test_token_dim_name(self) -> None:
+        assert TOKEN_DIM_NAME == "t"
+
+    def test_batch_dim_name(self) -> None:
+        assert BATCH_DIM_NAME == "b"
+
+    def test_seq_dim_name(self) -> None:
+        assert SEQ_DIM_NAME == "s"
+
+
+class TestFindDimIndex:
+    def test_found(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d")
+        assert find_dim_index(specs, "s") == 1
+
+    def test_not_found(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d")
+        assert find_dim_index(specs, "t") is None
+
+    def test_first_dim(self) -> None:
+        specs: list[DimSpec] = parse_dims("t h d")
+        assert find_dim_index(specs, "t") == 0
+
+    def test_last_dim(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s h d")
+        assert find_dim_index(specs, "d") == 3
+
+    def test_with_modifiers(self) -> None:
+        specs: list[DimSpec] = parse_dims("b s(cp,zigzag) h(tp) d")
+        assert find_dim_index(specs, "h") == 2
+
+    def test_empty_list(self) -> None:
+        assert find_dim_index([], "t") is None
+
+
+class TestResolveDimByName:
+    def test_resolve_found(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3, 4).refine_names("b", "s", "h")
+        assert resolve_dim_by_name(tensor, "b") == 0
+        assert resolve_dim_by_name(tensor, "s") == 1
+        assert resolve_dim_by_name(tensor, "h") == 2
+
+    def test_resolve_not_found_raises(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3).refine_names("b", "s")
+        with pytest.raises(ValueError, match="not in tensor names"):
+            resolve_dim_by_name(tensor, "h")
+
+    def test_resolve_unnamed_raises(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        with pytest.raises(ValueError, match="no names"):
+            resolve_dim_by_name(tensor, "b")
+
+
+class TestApplyDimNames:
+    def test_apply(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3, 4)
+        named: torch.Tensor = apply_dim_names(tensor, ["b", "s", "h"])
+        assert named.names == ("b", "s", "h")
+        assert named.shape == (2, 3, 4)
+
+    def test_apply_preserves_data(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        named: torch.Tensor = apply_dim_names(tensor, ["x", "y"])
+        assert torch.equal(strip_dim_names(named), tensor)
+
+
+class TestStripDimNames:
+    def test_strip(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3).refine_names("a", "b")
+        stripped: torch.Tensor = strip_dim_names(tensor)
+        assert stripped.names == (None, None)
+
+    def test_strip_already_unnamed(self) -> None:
+        tensor: torch.Tensor = torch.randn(2, 3)
+        stripped: torch.Tensor = strip_dim_names(tensor)
+        assert stripped.names == (None, None)
+
+
+class TestResolveDimNames:
+    def test_no_squeeze(self) -> None:
+        assert resolve_dim_names("t h d") == ["t", "h", "d"]
+
+    def test_single_squeeze(self) -> None:
+        assert resolve_dim_names("t 1 h") == ["t", "singleton0", "h"]
+
+    def test_multiple_squeeze(self) -> None:
+        assert resolve_dim_names("1 t 1 h") == [
+            "singleton0",
+            "t",
+            "singleton1",
+            "h",
+        ]
+
+
+class TestSingletonDimUtilFilterOut:
+    def test_no_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("t h d")
+        assert _SingletonDimUtil.filter_out(specs) == specs
+
+    def test_with_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("t 1 h")
+        filtered: list[DimSpec] = _SingletonDimUtil.filter_out(specs)
+        assert len(filtered) == 2
+        assert filtered[0].name == "t"
+        assert filtered[1].name == "h"
+
+    def test_all_squeeze(self) -> None:
+        specs: list[DimSpec] = parse_dims("1 1")
+        assert _SingletonDimUtil.filter_out(specs) == []
+
+
+class TestSingletonDimUtilIsSqueeze:
+    def test_squeeze(self) -> None:
+        assert _SingletonDimUtil.is_squeeze(DimSpec(name=SQUEEZE_DIM_NAME)) is True
+
+    def test_non_squeeze(self) -> None:
+        assert _SingletonDimUtil.is_squeeze(DimSpec(name="t")) is False
+
+
+class TestSingletonDimUtilMakeName:
+    def test_indices(self) -> None:
+        assert _SingletonDimUtil.make_name(0) == "singleton0"
+        assert _SingletonDimUtil.make_name(1) == "singleton1"
+        assert _SingletonDimUtil.make_name(99) == "singleton99"
+
+
+class TestSingletonDimUtilSanitizeNames:
+    def test_no_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["t", "h", "d"]) == ["t", "h", "d"]
+
+    def test_single_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["t", "1", "h"]) == [
+            "t",
+            "singleton0",
+            "h",
+        ]
+
+    def test_multiple_squeeze(self) -> None:
+        assert _SingletonDimUtil.sanitize_names(["1", "t", "1", "h"]) == [
+            "singleton0",
+            "t",
+            "singleton1",
+            "h",
+        ]
+
+    def test_empty(self) -> None:
+        assert _SingletonDimUtil.sanitize_names([]) == []
+
 
 class TestParseDimNames:
     def test_plain(self) -> None:

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -945,6 +945,115 @@ class TestEntrypointGroupingLogical:
         assert len(recompute_warnings) > 0
 
 
+class TestEntrypointAxisAligner:
+    """Test cross-framework dim reordering through the full entrypoint pipeline."""
+
+    def test_axis_swap_different_dim_order(self, tmp_path, capsys):
+        """Baseline dims 'b h d' vs target dims 'b d h': axis swapper rearranges baseline to match."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="b h d",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.permute(0, 2, 1).contiguous(),
+            dims="b d h",
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+        assert comp.baseline.shape == [4, 16, 8]
+        assert comp.target.shape == [4, 16, 8]
+
+    def test_axis_swap_with_tp_unshard(self, tmp_path, capsys):
+        """Baseline TP=2 with dims 'b h(tp) d' vs target TP=2 with dims 'b d h(tp)': unshard + axis swap."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_tensor,
+            name="hidden",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="b h(tp) d",
+        )
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_tensor.permute(0, 2, 1).contiguous(),
+            name="hidden",
+            tp_size=2,
+            shard_dim=2,
+            dims_str="b d h(tp)",
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+
+    def test_squeeze_dim_one_side(self, tmp_path, capsys):
+        """SGLang dims 't h' vs Megatron dims 't 1 h': axis aligner squeezes the singleton dim."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor,
+            dims="t h",
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden",
+            tensor=full_tensor.unsqueeze(1),
+            dims="t 1 h",
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=1e-3,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+        assert comp.baseline.shape == [4, 8]
+        assert comp.target.shape == [4, 8]
+
+
 class TestEntrypointAxisSwapper:
     """Test cross-framework dim reordering through the full entrypoint pipeline."""
 

test/registered/debug_utils/comparator/test_model_validation.py

@@ -403,5 +403,79 @@ class TestAlignerPlanInComparisonRecord:
         assert "unsharder" in text
 
 
+def _make_aligner_plan() -> AlignerPlan:
+    unsharder = UnsharderPlan(
+        axis=ParallelAxis.TP,
+        params=ConcatParams(dim_name="h"),
+        groups=[[0, 1]],
+    )
+    return AlignerPlan(
+        per_step_plans=Pair(
+            x=[
+                AlignerPerStepPlan(
+                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
+                )
+            ],
+            y=[
+                AlignerPerStepPlan(
+                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
+                )
+            ],
+        ),
+    )
+
+
+class TestAlignerPlanInComparisonRecord:
+    def test_comparison_record_with_aligner_plan(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+        assert record_with_plan.aligner_plan is not None
+        assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
+
+    def test_aligner_plan_json_roundtrip(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+
+        json_str: str = record_with_plan.model_dump_json()
+        parsed = json.loads(json_str)
+        assert "aligner_plan" in parsed
+        assert (
+            parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
+            == "unsharder"
+        )
+
+        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        assert roundtripped.aligner_plan is not None
+        assert (
+            roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
+            == "unsharder"
+        )
+
+    def test_comparison_record_without_aligner_plan(self) -> None:
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        json_str: str = record.model_dump_json()
+        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        assert roundtripped.aligner_plan is None
+
+    def test_aligner_plan_text_format(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+
+        text: str = record_with_plan.to_text()
+        assert "Aligner Plan:" in text
+        assert "unsharder" in text
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/source_patcher/test_code_patcher.py

@@ -8,6 +8,9 @@ from sglang.srt.debug_utils.source_patcher.code_patcher import (
     patch_function,
 )
 from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchSpec
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default")
 
 SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module"
 

test/registered/debug_utils/source_patcher/test_dumper_integration.py

@@ -6,6 +6,9 @@ from types import ModuleType
 import yaml
 
 from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default")
 
 SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module"
 

test/registered/debug_utils/source_patcher/test_source_editor.py

@@ -3,6 +3,9 @@ from pydantic import ValidationError
 
 from sglang.srt.debug_utils.source_patcher.source_editor import apply_edits
 from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchApplicationError
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default")
 
 
 class TestApplyEdits: