Support unifying axis ordering in dump comparator (#19456)

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."""