Support data parallel attention in dump comparator (#19602)

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

@@ -26,7 +26,7 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
 class TestComputeReordererPlans:
     def test_compute_reorderer_plans_zigzag(self) -> None:
         """s(cp:zigzag) produces a ReordererPlan."""
-        dim_specs = parse_dims("b s(cp:zigzag) h(tp)")
+        dim_specs = parse_dims("b s(cp:zigzag) h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -44,7 +44,7 @@ class TestComputeReordererPlans:
 
     def test_compute_reorderer_plans_thd_zigzag(self) -> None:
         """t(cp:zigzag) produces a ZigzagToNaturalThdParams plan."""
-        dim_specs = parse_dims("t(cp:zigzag) h(tp)")
+        dim_specs = parse_dims("t(cp:zigzag) h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -65,7 +65,7 @@ class TestComputeReordererPlans:
 
     def test_non_seq_dim_still_raises(self) -> None:
         """Zigzag on non-sequence/non-token dim (e.g. h(cp:zigzag)) raises ValueError."""
-        dim_specs = parse_dims("h(cp:zigzag) d")
+        dim_specs = parse_dims("h(cp:zigzag) d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)},
         ]
@@ -74,7 +74,7 @@ class TestComputeReordererPlans:
 
     def test_thd_zigzag_without_seq_lens_raises(self) -> None:
         """t(cp:zigzag) without thd_global_seq_lens raises ValueError."""
-        dim_specs = parse_dims("t(cp:zigzag) h(tp)")
+        dim_specs = parse_dims("t(cp:zigzag) h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -87,7 +87,7 @@ class TestComputeReordererPlans:
     def test_thd_natural_no_reorder(self) -> None:
         """t(cp:natural) and t(cp) produce no reorder plans."""
         for dims_str in ["t(cp:natural) h(tp)", "t(cp) h(tp)"]:
-            dim_specs = parse_dims(dims_str)
+            dim_specs = parse_dims(dims_str).dims
             parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
                 {
                     ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -102,7 +102,7 @@ class TestComputeReordererPlans:
     def test_compute_reorderer_plans_natural(self) -> None:
         """s(cp) and s(cp:natural) produce no reorder plans."""
         for dims_str in ["b s(cp) h(tp)", "b s(cp:natural) h(tp)"]:
-            dim_specs = parse_dims(dims_str)
+            dim_specs = parse_dims(dims_str).dims
             parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
                 {
                     ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -141,7 +141,7 @@ class TestCpZigzagTpE2E:
                     }
                 )
 
-        dim_specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp)")
+        dim_specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp)").dims
         dim_names: list[str] = [s.name for s in dim_specs]
 
         unsharder_plans = compute_unsharder_plan(
@@ -215,7 +215,7 @@ class TestCpZigzagSpSameDimE2E:
                     }
                 )
 
-        dim_specs: list[DimSpec] = parse_dims("t(cp:zigzag,sp) h")
+        dim_specs: list[DimSpec] = parse_dims("t(cp:zigzag,sp) h").dims
         dim_names: list[str] = [s.name for s in dim_specs]
 
         unsharder_plans = compute_unsharder_plan(

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

@@ -42,7 +42,7 @@ class TestExecuteUnsharderPlan:
         full_tensor = torch.randn(2, 8, 16)
         shards = list(full_tensor.chunk(4, dim=1))
 
-        dim_specs = parse_dims("b h(tp) d")
+        dim_specs = parse_dims("b h(tp) d").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
         ]
@@ -67,7 +67,7 @@ class TestExecuteUnsharderPlan:
             {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
         ]
-        dim_specs = parse_dims("h(tp) d")
+        dim_specs = parse_dims("h(tp) d").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 1
 
@@ -95,7 +95,7 @@ class TestExecuteUnsharderPlan:
         shards_a = list(full_a.chunk(4, dim=0))
         shards_b = list(full_b.chunk(4, dim=0))
 
-        dim_specs = parse_dims("s(cp) h(tp)")
+        dim_specs = parse_dims("s(cp) h(tp)").dims
         parallel_infos = []
         for cp_rank in range(2):
             for tp_rank in range(4):
@@ -145,7 +145,7 @@ class TestExecuteUnsharderPlan:
                     }
                 )
 
-        dim_specs = parse_dims("b s(cp) h(tp)")
+        dim_specs = parse_dims("b s(cp) h(tp)").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 2
 
@@ -187,7 +187,7 @@ class TestExecuteUnsharderPlan:
                 }
             )
 
-        dim_specs = parse_dims("b s(cp) h(tp)")
+        dim_specs = parse_dims("b s(cp) h(tp)").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 2
 
@@ -241,7 +241,7 @@ class TestExecuteUnsharderPlan:
                         }
                     )
 
-        dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
+        dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 3
 
@@ -290,7 +290,7 @@ class TestExecuteUnsharderPlan:
                 }
             )
 
-        dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
+        dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 3
 
@@ -307,7 +307,7 @@ 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")
+        dim_specs = parse_dims("h d").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -326,7 +326,7 @@ class TestPickOperation:
 
     def test_pick_multiple_groups(self) -> None:
         """PickParams with multiple groups picks one from each."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -378,7 +378,7 @@ class TestPickOperation:
                     }
                 )
 
-        dim_specs = parse_dims("b s(cp) d")
+        dim_specs = parse_dims("b s(cp) d").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 2
 
@@ -407,7 +407,7 @@ class TestPickOperation:
                     }
                 )
 
-        dim_specs = parse_dims("b h d")
+        dim_specs = parse_dims("b h d").dims
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert len(plans) == 2
         assert all(isinstance(p.params, PickParams) for p in plans)
@@ -472,7 +472,7 @@ class TestVerifyReplicatedGroup:
 
     def test_execute_returns_replicated_checks(self) -> None:
         """execute_unsharder_plan returns replicated checks for mismatch."""
-        dim_specs = parse_dims("h d")
+        dim_specs = parse_dims("h d").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -646,6 +646,261 @@ class TestThdCpConcat:
         )
 
 
+class TestReduceSum:
+    def test_basic_tp2_reduce(self) -> None:
+        """2 partial tensors sum to full tensor."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+        part_a = full_tensor * 0.6
+        part_b = full_tensor * 0.4
+
+        dim_specs = parse_dims("h(tp:partial) d").dims
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, ReduceSumParams)
+
+        named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
+            plans[0], named_parts
+        )
+
+        assert len(unsharder_result.tensors) == 1
+        assert torch.allclose(unsharder_result.tensors[0].rename(None), full_tensor)
+
+    def test_tp4_reduce(self) -> None:
+        """4 partial tensors sum to full tensor."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+        parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
+
+        dim_specs = parse_dims("h(tp:partial) d").dims
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+
+        named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
+            plans[0], named_parts
+        )
+
+        assert len(unsharder_result.tensors) == 1
+        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."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        cp_chunks = list(full_tensor.chunk(2, dim=1))
+        # Each CP chunk is held as partial sums across TP ranks
+        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] * 0.5)
+                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) h(tp:partial)").dims
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
+        for plan in plans:
+            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)
+
+    def test_reduce_scrambled_ranks(self) -> None:
+        """Scrambled rank order — sum is commutative so result is the same."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+        parts: list[torch.Tensor] = [
+            full_tensor * 0.1,
+            full_tensor * 0.2,
+            full_tensor * 0.3,
+            full_tensor * 0.4,
+        ]
+
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
+        ]
+        dim_specs = parse_dims("h(tp:partial) d").dims
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
+            plans[0], named_parts
+        )
+
+        assert len(unsharder_result.tensors) == 1
+        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."""
+        dim_specs = parse_dims("h(tp:partial) d").dims
+        part_a = torch.randn(4, 8).refine_names("h", "d")
+        part_b = torch.randn(4, 8).refine_names("h", "d")
+
+        plan = UnsharderPlan(
+            axis=ParallelAxis.TP,
+            params=ReduceSumParams(),
+            groups=[[0, 1]],
+        )
+        unsharder_result: UnsharderResult = execute_unsharder_plan(
+            plan, [part_a, part_b]
+        )
+
+        assert len(unsharder_result.tensors) == 1
+        assert unsharder_result.tensors[0].names == ("h", "d")
+        expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d")
+        assert torch.allclose(
+            unsharder_result.tensors[0].rename(None), expected.rename(None)
+        )
+
+    def test_recompute_pseudo_mismatch(self) -> None:
+        """_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
+        tensor_a = torch.ones(4)
+        tensor_b = torch.ones(4) + 0.1
+
+        checks: list[ReplicatedCheckResult] = _verify_replicated_group(
+            [tensor_a, tensor_b],
+            axis=ParallelAxis.RECOMPUTE_PSEUDO,
+            group_index=0,
+        )
+        assert len(checks) == 1
+        assert checks[0].axis == "recompute_pseudo"
+        assert checks[0].group_index == 0
+        assert checks[0].compared_index == 1
+        assert checks[0].baseline_index == 0
+        assert not checks[0].passed
+        assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
+
+
+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]],
+        )
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(unsharder_result.tensors) == 1
+        expected = torch.tensor([1, 2, 3, 4, 5, 6])
+        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]."""
+        # 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]],
+        )
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(unsharder_result.tensors) == 1
+        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]))
+        # 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]],
+        )
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(unsharder_result.tensors) == 1
+        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]))
+        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]],
+        )
+        unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
+
+        assert len(unsharder_result.tensors) == 1
+        unsharded: torch.Tensor = unsharder_result.tensors[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)
+        )
+
+
 class TestReduceSum:
     def test_basic_tp2_reduce(self) -> None:
         """2 partial tensors sum to full tensor."""

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

@@ -19,7 +19,7 @@ register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 class TestComputeUnsharderPlan:
     def test_tp4_plan(self) -> None:
-        dim_specs = parse_dims("b s h(tp) d")
+        dim_specs = parse_dims("b s h(tp) d").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
         ]
@@ -31,7 +31,7 @@ class TestComputeUnsharderPlan:
         assert plans[0].groups == [[0, 1, 2, 3]]
 
     def test_inconsistent_axis_size_raises(self) -> None:
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -41,7 +41,7 @@ class TestComputeUnsharderPlan:
 
     def test_missing_axis_in_all_parallel_infos_skipped(self) -> None:
         """Axis in dims but absent from all parallel_infos -> axis_size=1, auto-skip."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
         # TP not in any parallel_info → skipped; CP is replicated but only 1 rank
         # with size=2 → incomplete coverage
@@ -49,13 +49,13 @@ class TestComputeUnsharderPlan:
             compute_unsharder_plan(dim_specs, parallel_infos)
 
     def test_empty_parallel_infos_raises(self) -> None:
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         with pytest.raises(ValueError, match="must not be empty"):
             compute_unsharder_plan(dim_specs, [])
 
     def test_scrambled_world_ranks(self) -> None:
         """world_rank order != axis_rank order."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -67,14 +67,14 @@ class TestComputeUnsharderPlan:
         assert plans[0].groups == [[1, 3, 0, 2]]
 
     def test_no_sharded_axes_returns_empty(self) -> None:
-        dim_specs = parse_dims("b s d")
+        dim_specs = parse_dims("b s d").dims
         parallel_infos = [{}]
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert plans == []
 
     def test_multi_axis_plan(self) -> None:
         """Multi-axis (TP + CP) produces a 2-step plan."""
-        dim_specs = parse_dims("s(cp) h(tp)")
+        dim_specs = parse_dims("s(cp) h(tp)").dims
         parallel_infos = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -101,7 +101,7 @@ class TestComputeUnsharderPlan:
 
     def test_cp_tp_plan(self) -> None:
         """CP=2 + TP=4 produces correct 2-step plan with correct groups."""
-        dim_specs = parse_dims("s(cp) h(tp)")
+        dim_specs = parse_dims("s(cp) h(tp)").dims
         parallel_infos = []
         for cp_rank in range(2):
             for tp_rank in range(4):
@@ -129,7 +129,7 @@ class TestComputeUnsharderPlan:
 
     def test_cp_tp_scrambled_ranks(self) -> None:
         """Scrambled rank assignment still produces correct plan."""
-        dim_specs = parse_dims("s(cp) h(tp)")
+        dim_specs = parse_dims("s(cp) h(tp)").dims
         parallel_infos = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
@@ -165,7 +165,7 @@ class TestComputeUnsharderPlan:
 
     def test_axis_rank_coverage_incomplete_raises(self) -> None:
         """TP size=4 but only ranks 0,1,3 provided (missing rank 2)."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
@@ -175,7 +175,7 @@ class TestComputeUnsharderPlan:
             compute_unsharder_plan(dim_specs, parallel_infos)
 
     def test_reduction_partial_returns_reduce_sum(self) -> None:
-        dim_specs = parse_dims("h(tp:partial)")
+        dim_specs = parse_dims("h(tp:partial)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
         ]
@@ -188,7 +188,7 @@ class TestComputeUnsharderPlan:
 
     def test_reduction_partial_tp4(self) -> None:
         """TP=4 with partial reduction produces a single ReduceSumParams step."""
-        dim_specs = parse_dims("h(tp:partial)")
+        dim_specs = parse_dims("h(tp:partial)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
         ]
@@ -200,7 +200,7 @@ class TestComputeUnsharderPlan:
 
     def test_multi_axis_with_reduction_on_one(self) -> None:
         """CP concat + TP reduce produces a 2-step plan."""
-        dim_specs = parse_dims("s(cp) h(tp:partial)")
+        dim_specs = parse_dims("s(cp) h(tp:partial)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
         for cp_rank in range(2):
             for tp_rank in range(2):
@@ -221,7 +221,7 @@ class TestComputeUnsharderPlan:
 
     def test_reduction_scrambled_ranks(self) -> None:
         """Scrambled world_rank order with partial reduction."""
-        dim_specs = parse_dims("h(tp:partial)")
+        dim_specs = parse_dims("h(tp:partial)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -235,7 +235,7 @@ class TestComputeUnsharderPlan:
         assert plans[0].groups == [[1, 3, 0, 2]]
 
     def test_ordering_zigzag_accepted(self) -> None:
-        dim_specs = parse_dims("s(cp:zigzag)")
+        dim_specs = parse_dims("s(cp:zigzag)").dims
         parallel_infos = [
             {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
         ]
@@ -244,7 +244,7 @@ class TestComputeUnsharderPlan:
         assert plans[0].axis == ParallelAxis.CP
 
     def test_ordering_natural_accepted(self) -> None:
-        dim_specs = parse_dims("s(cp:natural)")
+        dim_specs = parse_dims("s(cp:natural)").dims
         parallel_infos = [
             {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
         ]
@@ -254,7 +254,7 @@ class TestComputeUnsharderPlan:
 
     def test_three_axis_plan(self) -> None:
         """EP=2 + CP=2 + TP=2 produces a 3-step plan."""
-        dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
+        dim_specs = parse_dims("b e(ep) s(cp) h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
         for ep_rank in range(2):
             for cp_rank in range(2):
@@ -290,7 +290,7 @@ class TestComputeUnsharderPlan:
 
     def test_same_dim_cp_sp_plan(self) -> None:
         """t(cp:zigzag,sp) with CP=2 SP=2: SP unshards first (inner), then CP."""
-        dim_specs = parse_dims("t(cp:zigzag,sp) 1 h")
+        dim_specs = parse_dims("t(cp:zigzag,sp) 1 h").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
         for cp_rank in range(2):
             for sp_rank in range(2):
@@ -328,7 +328,7 @@ class TestComputeUnsharderPlan:
             CpThdConcatParams,
         )
 
-        dim_specs = parse_dims("t(cp:zigzag,sp) h")
+        dim_specs = parse_dims("t(cp:zigzag,sp) h").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
         for cp_rank in range(2):
             for sp_rank in range(2):
@@ -361,7 +361,7 @@ class TestComputeUnsharderPlan:
 
     def test_sp_in_dims_but_not_in_parallel_info(self) -> None:
         """s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip."""
-        dim_specs = parse_dims("s(sp) b h(tp)")
+        dim_specs = parse_dims("s(sp) b h(tp)").dims
         parallel_infos = [
             {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2)},
@@ -372,14 +372,14 @@ class TestComputeUnsharderPlan:
 
     def test_all_dims_sharded_but_single_gpu(self) -> None:
         """Single GPU (TP=1, CP=1), dims has s(cp) h(tp) but parallel_info is empty."""
-        dim_specs = parse_dims("b s(cp) h(tp) d")
+        dim_specs = parse_dims("b s(cp) h(tp) d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [{}]
         plans = compute_unsharder_plan(dim_specs, parallel_infos)
         assert plans == []
 
     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)")
+        dim_specs = parse_dims("s(cp) h(tp)").dims
         parallel_infos = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -397,7 +397,7 @@ class TestComputeUnsharderPlan:
 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")
+        dim_specs = parse_dims("b s(cp) d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -431,7 +431,7 @@ class TestReplicatedAxes:
 
     def test_fully_replicated(self) -> None:
         """CP2 TP2, dims='b h d' → PickPlan(CP) + PickPlan(TP)."""
-        dim_specs = parse_dims("b h d")
+        dim_specs = parse_dims("b h d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -459,7 +459,7 @@ class TestReplicatedAxes:
 
     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)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
         for cp_rank in range(2):
             for ep_rank in range(2):
@@ -486,7 +486,7 @@ class TestReplicatedAxes:
 
     def test_replicated_scrambled_ranks(self) -> None:
         """Scrambled world_rank order with replicated axis."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
@@ -515,7 +515,7 @@ class TestReplicatedAxes:
 
     def test_replicated_axis_inconsistent_size_raises(self) -> None:
         """Replicated axis with inconsistent sizes raises ValueError."""
-        dim_specs = parse_dims("h(tp)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -531,7 +531,7 @@ class TestReplicatedAxes:
 
     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)")
+        dim_specs = parse_dims("h(tp)").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {
                 ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
@@ -547,7 +547,7 @@ class TestReplicatedAxes:
 
     def test_recompute_pseudo_replicated(self) -> None:
         """RECOMPUTE_PSEUDO with no dim annotation → replicated → PickParams."""
-        dim_specs = parse_dims("h d")
+        dim_specs = parse_dims("h d").dims
         parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
             {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=0, axis_size=2)},
             {ParallelAxis.RECOMPUTE_PSEUDO: AxisInfo(axis_rank=1, axis_size=2)},

test/registered/debug_utils/comparator/test_dims.py

@@ -9,6 +9,7 @@ from sglang.srt.debug_utils.comparator.dims import (
     SQUEEZE_DIM_NAME,
     TOKEN_DIM_NAME,
     DimSpec,
+    DimsSpec,
     Ordering,
     ParallelAxis,
     ParallelModifier,
@@ -17,7 +18,6 @@ from sglang.srt.debug_utils.comparator.dims import (
     apply_dim_names,
     find_dim_index,
     parse_dim,
-    parse_dim_names,
     parse_dims,
     resolve_dim_by_name,
     resolve_dim_names,
@@ -113,10 +113,17 @@ class TestParseDim:
         with pytest.raises(ValueError, match="Invalid dim token"):
             parse_dim("1(tp)")
 
+    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:
-        assert parse_dims("b s h d") == [
+        assert parse_dims("b s h d").dims == [
             DimSpec(name="b"),
             DimSpec(name="s"),
             DimSpec(name="h"),
@@ -124,10 +131,10 @@ class TestParseDims:
         ]
 
     def test_single_dim(self) -> None:
-        assert parse_dims("t") == [DimSpec(name="t")]
+        assert parse_dims("t").dims == [DimSpec(name="t")]
 
     def test_mixed_annotated(self) -> None:
-        assert parse_dims("b s(cp:zigzag) h(tp) d") == [
+        assert parse_dims("b s(cp:zigzag) h(tp) d").dims == [
             DimSpec(name="b"),
             DimSpec(
                 name="s",
@@ -155,17 +162,17 @@ class TestParseDims:
             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")
+        dims: list[DimSpec] = parse_dims("t 1 h").dims
+        assert len(dims) == 3
+        assert dims[0] == DimSpec(name="t")
+        assert dims[1] == DimSpec(name="1")
+        assert dims[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")
+        dims: list[DimSpec] = parse_dims("t 1 h 1 d").dims
+        assert len(dims) == 5
+        assert dims[1] == DimSpec(name="1")
+        assert dims[3] == DimSpec(name="1")
 
 
 class TestDimConstants:
@@ -181,23 +188,23 @@ class TestDimConstants:
 
 class TestFindDimIndex:
     def test_found(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d")
+        specs: list[DimSpec] = parse_dims("b s h d").dims
         assert find_dim_index(specs, "s") == 1
 
     def test_not_found(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d")
+        specs: list[DimSpec] = parse_dims("b s h d").dims
         assert find_dim_index(specs, "t") is None
 
     def test_first_dim(self) -> None:
-        specs: list[DimSpec] = parse_dims("t h d")
+        specs: list[DimSpec] = parse_dims("t h d").dims
         assert find_dim_index(specs, "t") == 0
 
     def test_last_dim(self) -> None:
-        specs: list[DimSpec] = parse_dims("b s h d")
+        specs: list[DimSpec] = parse_dims("b s h d").dims
         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")
+        specs: list[DimSpec] = parse_dims("b s(cp:zigzag) h(tp) d").dims
         assert find_dim_index(specs, "h") == 2
 
     def test_empty_list(self) -> None:
@@ -265,18 +272,18 @@ class TestResolveDimNames:
 
 class TestSingletonDimUtilFilterOut:
     def test_no_squeeze(self) -> None:
-        specs: list[DimSpec] = parse_dims("t h d")
+        specs: list[DimSpec] = parse_dims("t h d").dims
         assert _SingletonDimUtil.filter_out(specs) == specs
 
     def test_with_squeeze(self) -> None:
-        specs: list[DimSpec] = parse_dims("t 1 h")
+        specs: list[DimSpec] = parse_dims("t 1 h").dims
         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")
+        specs: list[DimSpec] = parse_dims("1 1").dims
         assert _SingletonDimUtil.filter_out(specs) == []
 
 
@@ -318,91 +325,51 @@ class TestSingletonDimUtilSanitizeNames:
         assert _SingletonDimUtil.sanitize_names([]) == []
 
 
-class TestParseDimNames:
-    def test_plain(self) -> None:
-        assert parse_dim_names("b s h d") == ["b", "s", "h", "d"]
+class TestParseDimsWithHash:
+    """parse_dims strips the ``#`` declaration section from dims."""
 
-    def test_strips_modifiers(self) -> None:
-        assert parse_dim_names("b s(cp,zigzag) h(tp) d") == ["b", "s", "h", "d"]
+    def test_shape_dims_unchanged(self) -> None:
+        assert parse_dims("b s h(tp) # dp:=moe_dp").dims == parse_dims("b s h(tp)").dims
+
+    def test_dp_group_alias_extracted(self) -> None:
+        assert parse_dims("b s h(tp) # dp:=moe_dp").dp_group_alias == "moe_dp"
+
+    def test_no_hash_no_alias(self) -> None:
+        assert parse_dims("b s h(tp)").dp_group_alias is None
+
+    def test_whitespace_around_hash(self) -> None:
+        assert parse_dims("t h #   dp:=foo  ").dims == parse_dims("t h").dims
+        assert parse_dims("t h #   dp:=foo  ").dp_group_alias == "foo"
+
+    def test_multiple_declarations_picks_dp(self) -> None:
+        result: DimsSpec = parse_dims("t h(tp) # dp:=moe_dp ep:replicated")
+        assert result.dims == parse_dims("t h(tp)").dims
+        assert result.dp_group_alias == "moe_dp"
+
+    def test_no_dp_alias_token(self) -> None:
+        assert parse_dims("t h(tp) # ep:replicated").dp_group_alias is None
 
 
-class TestDimConstants:
-    def test_token_dim_name(self) -> None:
-        assert TOKEN_DIM_NAME == "t"
+class TestDpGroupAlias:
+    def test_basic(self) -> None:
+        assert parse_dims("b s h(tp) # dp:=moe_dp").dp_group_alias == "moe_dp"
 
-    def test_batch_dim_name(self) -> None:
-        assert BATCH_DIM_NAME == "b"
+    def test_no_hash_returns_none(self) -> None:
+        assert parse_dims("t h").dp_group_alias is None
 
-    def test_seq_dim_name(self) -> None:
-        assert SEQ_DIM_NAME == "s"
+    def test_no_dp_alias_token(self) -> None:
+        assert parse_dims("t h(tp) # ep:replicated").dp_group_alias is None
+
+    def test_multiple_tokens_picks_dp(self) -> None:
+        assert (
+            parse_dims("b s # ep:replicated dp:=custom_dp").dp_group_alias
+            == "custom_dp"
+        )
 
 
-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 TestResolveDimNamesWithHash:
+    def test_hash_stripped(self) -> None:
+        assert resolve_dim_names("t h # dp:=moe_dp") == ["t", "h"]
 
 
 if __name__ == "__main__":

test/registered/debug_utils/comparator/test_dp_utils.py

@@ -214,5 +214,122 @@ class TestFilterToNonEmptyDpRank:
         assert torch.equal(result[1].value, torch.tensor([2.0]))
 
 
+# ---------------------------------------------------------------------------
+# dp_group_alias tests
+# ---------------------------------------------------------------------------
+
+
+class TestExtractDpInfoWithAlias:
+    def test_alias_found(self) -> None:
+        meta: dict = {
+            "sglang_parallel_info": {
+                "dp_rank": 0,
+                "dp_size": 2,
+                "moe_dp_rank": 1,
+                "moe_dp_size": 4,
+            }
+        }
+        assert _extract_dp_info(meta, dp_group_alias="moe_dp") == (1, 4)
+
+    def test_alias_not_found_returns_none(self) -> None:
+        meta: dict = _make_sglang_meta(dp_rank=0, dp_size=2)
+        assert _extract_dp_info(meta, dp_group_alias="moe_dp") is None
+
+    def test_alias_none_uses_default(self) -> None:
+        meta: dict = _make_sglang_meta(dp_rank=1, dp_size=4)
+        assert _extract_dp_info(meta, dp_group_alias=None) == (1, 4)
+
+
+class TestFilterToNonEmptyDpRankWithAlias:
+    def test_alias_none_unchanged_behavior(self) -> None:
+        """dp_group_alias=None → same behavior as before (regression)."""
+        items: list[ValueWithMeta] = [
+            _make_item(
+                value=torch.tensor([1.0, 2.0]),
+                meta=_make_sglang_meta(dp_rank=0, dp_size=2),
+            ),
+            _make_item(
+                value=torch.tensor([]),
+                meta=_make_sglang_meta(dp_rank=1, dp_size=2),
+            ),
+        ]
+
+        result: list[ValueWithMeta] = filter_to_non_empty_dp_rank(
+            items, dp_group_alias=None
+        )
+
+        assert len(result) == 1
+        assert torch.equal(result[0].value, torch.tensor([1.0, 2.0]))
+
+    def test_alias_group_absent_noop(self) -> None:
+        """Alias group not in metadata → noop, return items unchanged."""
+        items: list[ValueWithMeta] = [
+            _make_item(
+                value=torch.tensor([1.0]),
+                meta=_make_sglang_meta(dp_rank=0, dp_size=2),
+            ),
+            _make_item(
+                value=torch.tensor([2.0]),
+                meta=_make_sglang_meta(dp_rank=1, dp_size=2),
+            ),
+        ]
+
+        result: list[ValueWithMeta] = filter_to_non_empty_dp_rank(
+            items, dp_group_alias="moe_dp"
+        )
+
+        assert result is items
+
+    def test_alias_size_1_noop(self) -> None:
+        """Alias group present but size=1 → noop."""
+        meta: dict = {
+            "sglang_parallel_info": {
+                "dp_rank": 0,
+                "dp_size": 2,
+                "moe_dp_rank": 0,
+                "moe_dp_size": 1,
+            }
+        }
+        items: list[ValueWithMeta] = [
+            _make_item(value=torch.tensor([1.0]), meta=meta),
+        ]
+
+        result: list[ValueWithMeta] = filter_to_non_empty_dp_rank(
+            items, dp_group_alias="moe_dp"
+        )
+
+        assert result is items
+
+    def test_alias_filters_correctly(self) -> None:
+        """Alias group size=2, one empty rank → correctly filters."""
+        meta_rank0: dict = {
+            "sglang_parallel_info": {
+                "dp_rank": 0,
+                "dp_size": 2,
+                "moe_dp_rank": 0,
+                "moe_dp_size": 2,
+            }
+        }
+        meta_rank1: dict = {
+            "sglang_parallel_info": {
+                "dp_rank": 0,
+                "dp_size": 2,
+                "moe_dp_rank": 1,
+                "moe_dp_size": 2,
+            }
+        }
+        items: list[ValueWithMeta] = [
+            _make_item(value=torch.tensor([1.0, 2.0]), meta=meta_rank0),
+            _make_item(value=torch.tensor([]), meta=meta_rank1),
+        ]
+
+        result: list[ValueWithMeta] = filter_to_non_empty_dp_rank(
+            items, dp_group_alias="moe_dp"
+        )
+
+        assert len(result) == 1
+        assert torch.equal(result[0].value, torch.tensor([1.0, 2.0]))
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -1630,80 +1630,6 @@ class TestEntrypointAxisAligner:
         assert comp.target.shape == [4, 8]
 
 
-class TestEntrypointAxisSwapper:
-    """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"
-
-
 class TestEntrypointReplicatedAxis:
     """Test replicated-axis scenarios through the full entrypoint pipeline."""
 
@@ -3360,6 +3286,151 @@ class TestEntrypointDpFilter:
             _run_and_parse(args, capsys)
 
 
+class TestEntrypointDpGroupAlias:
+    """E2E tests for the ``# dp:=<group>`` dp group alias feature.
+
+    In dp_attn mode, dp_size > 1 but MLP tensors after dp_gather have data
+    on all ranks.  With ``# dp:=moe_dp`` in dims, the dp filter uses
+    ``moe_dp_rank/moe_dp_size`` instead of ``dp_rank/dp_size``.
+    """
+
+    def test_dp_alias_absent_group_noop(self, tmp_path: Path, capsys) -> None:
+        """Single rank with ``# dp:=moe_dp`` in dims → parse_dims strips ``#``, comparison OK."""
+        torch.manual_seed(42)
+        tensor_data: torch.Tensor = torch.randn(10, 8)
+        target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001
+
+        for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]:
+            side_dir: Path = tmp_path / side_dir_name
+            side_dir.mkdir()
+
+            _create_rank_dump(
+                side_dir,
+                rank=0,
+                name="hidden",
+                tensor=data,
+                dims="t h # dp:=moe_dp",
+                parallel_info={
+                    "tp_rank": 0,
+                    "tp_size": 1,
+                    "dp_rank": 0,
+                    "dp_size": 1,
+                },
+                framework="sglang",
+            )
+
+        args: Namespace = _make_args(
+            tmp_path / "baseline" / _FIXED_EXP_NAME,
+            tmp_path / "target" / _FIXED_EXP_NAME,
+            grouping="logical",
+            diff_threshold=1e-3,
+        )
+        records, _ = _run_and_parse(args, capsys)
+
+        comparison: ComparisonRecord = _assert_single_comparison_passed(records)
+        assert comparison.name == "hidden"
+
+    def test_dp_alias_via_override_dims(self, tmp_path: Path, capsys) -> None:
+        """--override-dims adds ``# dp:=moe_dp`` → dp filter uses alias, filters correctly."""
+        torch.manual_seed(42)
+        tensor_data: torch.Tensor = torch.randn(10, 8)
+        target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001
+
+        for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]:
+            side_dir: Path = tmp_path / side_dir_name
+            side_dir.mkdir()
+
+            # moe_dp_rank=0: non-empty
+            _create_rank_dump(
+                side_dir,
+                rank=0,
+                name="hidden",
+                tensor=data,
+                dims="t h",
+                parallel_info={
+                    "tp_rank": 0,
+                    "tp_size": 1,
+                    "dp_rank": 0,
+                    "dp_size": 1,
+                    "moe_dp_rank": 0,
+                    "moe_dp_size": 2,
+                },
+                framework="sglang",
+            )
+
+            # moe_dp_rank=1: empty
+            _create_rank_dump(
+                side_dir,
+                rank=1,
+                name="hidden",
+                tensor=torch.empty(0, 8),
+                dims="t h",
+                parallel_info={
+                    "tp_rank": 0,
+                    "tp_size": 1,
+                    "dp_rank": 0,
+                    "dp_size": 1,
+                    "moe_dp_rank": 1,
+                    "moe_dp_size": 2,
+                },
+                framework="sglang",
+            )
+
+        args: Namespace = _make_args(
+            tmp_path / "baseline" / _FIXED_EXP_NAME,
+            tmp_path / "target" / _FIXED_EXP_NAME,
+            grouping="logical",
+            diff_threshold=1e-3,
+            override_dims=["hidden:t h # dp:=moe_dp"],
+        )
+        records, _ = _run_and_parse(args, capsys)
+
+        comparison: ComparisonRecord = _assert_single_comparison_passed(records)
+        assert comparison.name == "hidden"
+
+    def test_dp_alias_with_real_alias_group_filters(
+        self, tmp_path: Path, capsys
+    ) -> None:
+        """Alias group present with moe_dp_size=2, one empty rank → filters correctly."""
+        torch.manual_seed(42)
+        tensor_data: torch.Tensor = torch.randn(10, 8)
+        target_data: torch.Tensor = tensor_data + torch.randn(10, 8) * 0.001
+
+        for side_dir_name, data in [("baseline", tensor_data), ("target", target_data)]:
+            side_dir: Path = tmp_path / side_dir_name
+            side_dir.mkdir()
+
+            for moe_dp_rank in range(2):
+                tensor: torch.Tensor = data if moe_dp_rank == 0 else torch.empty(0, 8)
+                _create_rank_dump(
+                    side_dir,
+                    rank=moe_dp_rank,
+                    name="hidden",
+                    tensor=tensor,
+                    dims="t h # dp:=moe_dp",
+                    parallel_info={
+                        "tp_rank": 0,
+                        "tp_size": 1,
+                        "dp_rank": 0,
+                        "dp_size": 1,
+                        "moe_dp_rank": moe_dp_rank,
+                        "moe_dp_size": 2,
+                    },
+                    framework="sglang",
+                )
+
+        args: Namespace = _make_args(
+            tmp_path / "baseline" / _FIXED_EXP_NAME,
+            tmp_path / "target" / _FIXED_EXP_NAME,
+            grouping="logical",
+            diff_threshold=1e-3,
+        )
+        records, _ = _run_and_parse(args, capsys)
+
+        comparison: ComparisonRecord = _assert_single_comparison_passed(records)
+        assert comparison.name == "hidden"
+
+
 class TestEntrypointMetaOverride:
     """E2E: dump with wrong dims → --override-dims / --override-config corrects at comparison time."""
 

test/registered/debug_utils/comparator/test_model_validation.py

@@ -403,79 +403,5 @@ 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

@@ -10,7 +10,7 @@ from sglang.srt.debug_utils.source_patcher.code_patcher import (
 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")
+register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module"
 

test/registered/debug_utils/source_patcher/test_dumper_integration.py

@@ -8,7 +8,7 @@ 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")
+register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module"
 

test/registered/debug_utils/source_patcher/test_source_editor.py

@@ -5,7 +5,7 @@ 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")
+register_cpu_ci(est_time=10, suite="default", nightly=True)
 
 
 class TestApplyEdits:

test/registered/debug_utils/test_dumper.py

@@ -2074,6 +2074,33 @@ class TestDumperE2E:
             assert "rank" in loaded["meta"]
             assert "step" in loaded["meta"]
 
+            par = loaded["meta"].get("sglang_parallel_info", {})
+            expected_keys = [
+                "tp_rank",
+                "tp_size",
+                "pp_rank",
+                "pp_size",
+                "moe_ep_rank",
+                "moe_ep_size",
+                "moe_tp_rank",
+                "moe_tp_size",
+                "moe_dp_rank",
+                "moe_dp_size",
+                "enable_dp_attention",
+                "attn_tp_rank",
+                "attn_tp_size",
+                "attn_dp_rank",
+                "attn_dp_size",
+                "local_attn_dp_rank",
+                "local_attn_dp_size",
+                "attn_cp_rank",
+                "attn_cp_size",
+            ]
+            for key in expected_keys:
+                assert (
+                    key in par
+                ), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
+
             rids_files = [f for f in dump_files if "name=rids" in f.name]
             rids_loaded = torch.load(
                 rids_files[0], map_location="cpu", weights_only=False