Support unifying axis ordering in dump comparator (#19456)

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

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from einops import rearrange
+
+from sglang.srt.debug_utils.comparator.dims import parse_dims
+from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
+from sglang.srt.debug_utils.comparator.utils import Pair, _FrozenBase
+from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
+
+# --- types ---
+
+
+class AxisSwapperPlan(_FrozenBase):
+    pattern: str  # einops pattern, e.g. "t h d -> t d h"
+
+
+# --- planner ---
+
+
+def compute_axis_swapper_plan(
+    dims_str_pair: Pair[Optional[str]],
+) -> Optional[AxisSwapperPlan]:
+    if dims_str_pair.x is None or dims_str_pair.y is None:
+        return None
+
+    x_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.x)]
+    y_names: list[str] = [spec.name for spec in parse_dims(dims_str_pair.y)]
+
+    if x_names == y_names:
+        return None
+
+    if set(x_names) != set(y_names):
+        warning_sink.add(
+            GeneralWarning(
+                category="axis_swapper_dim_mismatch",
+                message=(
+                    f"AxisSwapper: dim name sets differ (x={x_names}, y={y_names}), "
+                    f"skipping axis swap"
+                ),
+            )
+        )
+        return None
+
+    pattern: str = f"{' '.join(x_names)} -> {' '.join(y_names)}"
+    return AxisSwapperPlan(pattern=pattern)
+
+
+# --- executor ---
+
+
+def execute_axis_swapper_plan(
+    tensor: torch.Tensor, plan: AxisSwapperPlan
+) -> torch.Tensor:
+    return rearrange(tensor, plan.pattern)

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

@@ -5,6 +5,9 @@ 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.entrypoint.types import (
     AlignerPerStepPlan,
     AlignerPerStepSubPlan,
@@ -63,6 +66,13 @@ 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:
+        combined = Pair(
+            x=execute_axis_swapper_plan(tensor=combined.x, plan=swap_plan),
+            y=combined.y,
+        )
+
     return AlignerResult(tensors=combined, failed_side_xy=None)
 
 

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

@@ -2,6 +2,10 @@ 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.entrypoint.types import (
     AlignerPerStepPlan,
     AlignerPerStepSubPlan,
@@ -33,12 +37,21 @@ def compute_aligner_plan(
     token_aligner_plan: Optional[TokenAlignerPlan],
 ) -> AlignerPlan:
     token_dims: Pair[int] = metas_pair.map(_compute_token_dim)
+
+    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(
+        dims_str_pair=dims_str_pair
+    )
+
     return AlignerPlan(
         per_step_plans=metas_pair.map(
             lambda metas: _compute_per_step_plans(metas=metas)
         ),
         token_aligner_plan=token_aligner_plan,
         token_dims=token_dims,
+        axis_swapper_plan=axis_swapper_plan,
     )
 
 

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

@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from sglang.srt.debug_utils.comparator.aligner.axis_swapper import AxisSwapperPlan
 from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerPlan,
@@ -25,3 +26,4 @@ class AlignerPlan:
     per_step_plans: Pair[list[AlignerPerStepPlan]]
     token_aligner_plan: Optional[TokenAlignerPlan]
     token_dims: Pair[int]
+    axis_swapper_plan: Optional[AxisSwapperPlan] = None

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

@@ -0,0 +1,74 @@
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.axis_swapper import (
+    AxisSwapperPlan,
+    compute_axis_swapper_plan,
+    execute_axis_swapper_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 TestComputeAxisSwapperPlan:
+    def test_no_dims_returns_none(self) -> None:
+        assert compute_axis_swapper_plan(Pair(x=None, y=None)) is None
+        assert compute_axis_swapper_plan(Pair(x="t h d", y=None)) is None
+        assert compute_axis_swapper_plan(Pair(x=None, y="t h d")) is None
+
+    def test_same_order_returns_none(self) -> None:
+        result: Optional[AxisSwapperPlan] = compute_axis_swapper_plan(
+            Pair(x="t h d", y="t h d")
+        )
+        assert result is None
+
+    def test_different_order(self) -> None:
+        result: Optional[AxisSwapperPlan] = compute_axis_swapper_plan(
+            Pair(x="t h d", y="t d h")
+        )
+        assert result is not None
+        assert result.pattern == "t h d -> t d h"
+
+    def test_name_mismatch_returns_none_with_warning(self) -> None:
+        with warning_sink.context() as warnings:
+            result: Optional[AxisSwapperPlan] = compute_axis_swapper_plan(
+                Pair(x="t h d", y="t h e")
+            )
+
+        assert result is None
+        assert len(warnings) == 1
+        assert warnings[0].category == "axis_swapper_dim_mismatch"
+        assert "dim name sets differ" in warnings[0].message
+
+    def test_modifiers_ignored_for_name_extraction(self) -> None:
+        result: Optional[AxisSwapperPlan] = compute_axis_swapper_plan(
+            Pair(x="t h(tp) d", y="t d h(tp)")
+        )
+        assert result is not None
+        assert result.pattern == "t h d -> t d h"
+
+
+class TestExecuteAxisSwapperPlan:
+    def test_rearrange(self) -> None:
+        torch.manual_seed(42)
+        tensor: torch.Tensor = torch.randn(4, 8, 16)
+        plan = AxisSwapperPlan(pattern="t h d -> t d h")
+
+        result: torch.Tensor = execute_axis_swapper_plan(tensor=tensor, plan=plan)
+
+        assert result.shape == (4, 16, 8)
+        for i in range(4):
+            assert torch.equal(
+                result[i],
+                tensor[i].T,
+            )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -882,6 +882,80 @@ class TestEntrypointGroupingLogical:
         assert comp.name == "hidden"
 
 
+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."""