Support singleton dimension squeezing in dump comparator (#19566)

python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from einops import rearrange
+
+from sglang.srt.debug_utils.comparator.dims import (
+    _SingletonDimUtil,
+    parse_dims,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair, _FrozenBase
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+
+# --- types ---
+
+
+class AxisAlignerPlan(_FrozenBase):
+    pattern: Pair[Optional[str]]  # einops pattern per side, None = no-op
+
+
+# --- planner ---
+
+
+def compute_axis_aligner_plan(
+    dims_str_pair: Pair[Optional[str]],
+) -> Optional[AxisAlignerPlan]:
+    if dims_str_pair.x is None or dims_str_pair.y is None:
+        return None
+
+    # Pair[str] after None check
+    dims_pair: Pair[str] = Pair(x=dims_str_pair.x, y=dims_str_pair.y)
+
+    raw_names: Pair[list[str]] = dims_pair.map(
+        lambda s: [spec.name for spec in parse_dims(s)]
+    )
+    filtered_names: Pair[list[str]] = dims_pair.map(
+        lambda s: [spec.name for spec in _SingletonDimUtil.filter_out(parse_dims(s))]
+    )
+
+    target_order: Optional[list[str]] = _resolve_target_order(
+        x_names=filtered_names.x, y_names=filtered_names.y
+    )
+    if target_order is None:
+        return None
+
+    pattern: Pair[Optional[str]] = raw_names.map(
+        lambda names: _build_pattern(source=names, target=target_order)
+    )
+
+    if pattern.x is None and pattern.y is None:
+        return None
+
+    return AxisAlignerPlan(pattern=pattern)
+
+
+def _resolve_target_order(
+    x_names: list[str], y_names: list[str]
+) -> Optional[list[str]]:
+    """Determine the canonical dim order both sides should align to.
+
+    Returns y_names (the target ordering) if name sets match, or None on mismatch.
+    If both sides are identical, returns the shared order.
+    """
+    if x_names == y_names:
+        return y_names
+
+    if set(x_names) != set(y_names):
+        # Local import to avoid circular dependency:
+        # output_types -> aligner/entrypoint/types -> axis_aligner -> output_types
+        from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
+
+        warning_sink.add(
+            GeneralWarning(
+                category="axis_aligner_dim_mismatch",
+                message=(
+                    f"AxisAligner: dim name sets differ (x={x_names}, y={y_names}), "
+                    f"skipping axis swap"
+                ),
+            )
+        )
+        return None
+
+    return y_names
+
+
+def _build_pattern(*, source: list[str], target: list[str]) -> Optional[str]:
+    """Build an einops rearrange pattern from source dim names to target dim names.
+
+    Returns None if source already matches target (no rearrange needed).
+    """
+    if source == target:
+        return None
+
+    return f"{' '.join(source)} -> {' '.join(target)}"
+
+
+# --- executor ---
+
+
+def execute_axis_aligner_plan(
+    tensor: torch.Tensor, plan: AxisAlignerPlan, *, side: str
+) -> torch.Tensor:
+    pattern: Optional[str] = plan.pattern.x if side == "x" else plan.pattern.y
+
+    if pattern is not None:
+        tensor = rearrange(tensor.rename(None), pattern)
+
+    return tensor

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

@@ -5,8 +5,8 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.debug_utils.comparator.aligner.axis_swapper import (
-    execute_axis_swapper_plan,
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
+    execute_axis_aligner_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
     AlignerPerStepPlan,
@@ -65,11 +65,11 @@ def execute_aligner_plan(
             y=list(step_tensors_y.values())[0],
         )
 
-    # Cross-side: axis swap (rearrange x to match y's dim order)
-    if (swap_plan := plan.axis_swapper_plan) is not None:
+    # Cross-side: axis alignment (squeeze singletons + rearrange dim order)
+    if (aligner_plan := plan.axis_aligner_plan) is not None:
         combined = Pair(
-            x=execute_axis_swapper_plan(tensor=combined.x, plan=swap_plan),
-            y=combined.y,
+            x=execute_axis_aligner_plan(tensor=combined.x, plan=aligner_plan, side="x"),
+            y=execute_axis_aligner_plan(tensor=combined.y, plan=aligner_plan, side="y"),
         )
 
     return AlignerResult(tensors=combined, failed_side_xy=None)

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

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 from typing import Any, Optional
 
-from sglang.srt.debug_utils.comparator.aligner.axis_swapper import (
-    AxisSwapperPlan,
-    compute_axis_swapper_plan,
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
+    AxisAlignerPlan,
+    compute_axis_aligner_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
     AlignerPerStepPlan,
@@ -23,7 +23,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
 from sglang.srt.debug_utils.comparator.aligner.unsharder.planner import (
     compute_unsharder_plan,
 )
-from sglang.srt.debug_utils.comparator.dims import DimSpec, parse_dims
+from sglang.srt.debug_utils.comparator.dims import (
+    DimSpec,
+    _SingletonDimUtil,
+    parse_dims,
+)
 from sglang.srt.debug_utils.comparator.utils import Pair
 
 
@@ -38,7 +42,7 @@ def compute_aligner_plan(
     dims_str_pair: Pair[Optional[str]] = metas_pair.map(
         lambda metas: metas[0].get("dims") if metas else None
     )
-    axis_swapper_plan: Optional[AxisSwapperPlan] = compute_axis_swapper_plan(
+    axis_aligner_plan: Optional[AxisAlignerPlan] = compute_axis_aligner_plan(
         dims_str_pair=dims_str_pair
     )
 
@@ -54,7 +58,7 @@ def compute_aligner_plan(
             ),
         ),
         token_aligner_plan=token_aligner_plan,
-        axis_swapper_plan=axis_swapper_plan,
+        axis_aligner_plan=axis_aligner_plan,
     )
 
 
@@ -100,7 +104,7 @@ def compute_per_step_sub_plans(
     if dims_str is None:
         return []
 
-    dim_specs: list[DimSpec] = parse_dims(dims_str)
+    dim_specs: list[DimSpec] = _SingletonDimUtil.filter_out(parse_dims(dims_str))
     parallel_infos = [normalize_parallel_info(meta) for meta in metas]
 
     unsharder_plans = compute_unsharder_plan(

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

@@ -4,7 +4,7 @@ from typing import Annotated, Optional, Union
 
 from pydantic import Discriminator
 
-from sglang.srt.debug_utils.comparator.aligner.axis_swapper import AxisSwapperPlan
+from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan
 from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerPlan,
@@ -27,4 +27,4 @@ class AlignerPerStepPlan(_FrozenBase):
 class AlignerPlan(_FrozenBase):
     per_step_plans: Pair[list[AlignerPerStepPlan]]
     token_aligner_plan: Optional[TokenAlignerPlan] = None
-    axis_swapper_plan: Optional[AxisSwapperPlan] = None
+    axis_aligner_plan: Optional[AxisAlignerPlan] = None

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

@@ -28,7 +28,7 @@ from sglang.srt.debug_utils.comparator.dims import (
     ParallelAxis,
     TokenLayout,
     apply_dim_names,
-    parse_dim_names,
+    resolve_dim_names,
 )
 from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
@@ -230,7 +230,7 @@ def _load_and_align_aux_tensor(
     if sub_plans:
         dims_str: Optional[str] = metas[0].get("dims")
         if dims_str is not None:
-            dim_names: list[str] = parse_dim_names(dims_str)
+            dim_names: list[str] = resolve_dim_names(dims_str)
             tensors = [apply_dim_names(t, dim_names) for t in tensors]
 
         result = execute_sub_plans(tensors=tensors, plans=sub_plans)

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

@@ -18,7 +18,10 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPl
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import apply_dim_names, parse_dim_names
+from sglang.srt.debug_utils.comparator.dims import (
+    apply_dim_names,
+    resolve_dim_names,
+)
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     GeneralWarning,
@@ -241,7 +244,7 @@ def _apply_dim_names_from_meta(
     if dims_str is None:
         return tensors
 
-    dim_names: list[str] = parse_dim_names(dims_str)
+    dim_names: list[str] = resolve_dim_names(dims_str)
     return [apply_dim_names(t, dim_names) for t in tensors]
 
 

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

@@ -8,6 +8,7 @@ import torch
 TOKEN_DIM_NAME: str = "t"
 BATCH_DIM_NAME: str = "b"
 SEQ_DIM_NAME: str = "s"
+SQUEEZE_DIM_NAME: str = "1"
 
 
 class TokenLayout(Enum):
@@ -40,6 +41,46 @@ class DimSpec:
     reduction: Optional[Reduction] = None
 
 
+class _SingletonDimUtil:
+    """Utilities for squeeze dims (name="1") and their singleton tensor-name mapping."""
+
+    PREFIX: str = "singleton"
+
+    @staticmethod
+    def is_squeeze(spec: DimSpec) -> bool:
+        return spec.name == SQUEEZE_DIM_NAME
+
+    @staticmethod
+    def filter_out(dim_specs: list[DimSpec]) -> list[DimSpec]:
+        return [s for s in dim_specs if not _SingletonDimUtil.is_squeeze(s)]
+
+    @staticmethod
+    def make_name(index: int) -> str:
+        return f"{_SingletonDimUtil.PREFIX}{index}"
+
+    @staticmethod
+    def is_singleton_name(name: str) -> bool:
+        return (
+            name.startswith(_SingletonDimUtil.PREFIX)
+            and name[len(_SingletonDimUtil.PREFIX) :].isdigit()
+        )
+
+    @staticmethod
+    def sanitize_names(names: list[str]) -> list[str]:
+        """Replace '1' with 'singleton0', 'singleton1', ... for named tensor compatibility."""
+        result: list[str] = []
+        sq_idx: int = 0
+
+        for name in names:
+            if name == SQUEEZE_DIM_NAME:
+                result.append(_SingletonDimUtil.make_name(sq_idx))
+                sq_idx += 1
+            else:
+                result.append(name)
+
+        return result
+
+
 _DIM_PATTERN = re.compile(r"^(?P<name>[a-zA-Z_]\w*)(?:\((?P<modifiers>[^)]+)\))?$")
 
 _MODIFIER_FIELDS: list[tuple[type[Enum], str]] = [
@@ -55,6 +96,9 @@ for _enum_cls, _field in _MODIFIER_FIELDS:
 
 
 def parse_dim(token: str) -> DimSpec:
+    if token == SQUEEZE_DIM_NAME:
+        return DimSpec(name=SQUEEZE_DIM_NAME)
+
     match = _DIM_PATTERN.match(token)
     if match is None:
         raise ValueError(f"Invalid dim token: {token!r}")
@@ -84,16 +128,22 @@ def parse_dims(dims_str: str) -> list[DimSpec]:
 
     result = [parse_dim(token) for token in dims_str.strip().split()]
 
-    names = [spec.name for spec in result]
-    if len(names) != len(set(names)):
-        duplicates = sorted({n for n in names if names.count(n) > 1})
+    non_squeeze_names: list[str] = [
+        spec.name for spec in result if not _SingletonDimUtil.is_squeeze(spec)
+    ]
+    if len(non_squeeze_names) != len(set(non_squeeze_names)):
+        duplicates = sorted(
+            {n for n in non_squeeze_names if non_squeeze_names.count(n) > 1}
+        )
         raise ValueError(f"Duplicate dim names: {duplicates}")
 
     return result
 
 
-def parse_dim_names(dims_str: str) -> list[str]:
-    return [spec.name for spec in parse_dims(dims_str)]
+def resolve_dim_names(dims_str: str) -> list[str]:
+    """Parse dims string and return tensor-compatible names ('1' → 'singleton0', ...)."""
+    names: list[str] = [spec.name for spec in parse_dims(dims_str)]
+    return _SingletonDimUtil.sanitize_names(names)
 
 
 def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:

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

@@ -219,8 +219,13 @@ def _format_aligner_plan(plan: AlignerPlan) -> str:
         num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
         lines.append(f"  token_aligner: {num_tokens} tokens aligned")
 
-    if plan.axis_swapper_plan is not None:
-        lines.append(f"  axis_swapper: {plan.axis_swapper_plan.pattern}")
+    if plan.axis_aligner_plan is not None:
+        parts: list[str] = []
+        if plan.axis_aligner_plan.pattern.x:
+            parts.append(f"x: {plan.axis_aligner_plan.pattern.x}")
+        if plan.axis_aligner_plan.pattern.y:
+            parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
+        lines.append(f"  axis_aligner: {', '.join(parts)}")
 
     return "\n".join(lines)
 

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: