diff --git a/python/sglang/srt/debug_utils/comparator/aligner/axis_swapper.py b/python/sglang/srt/debug_utils/comparator/aligner/axis_swapper.py new file mode 100644 index 000000000..6e3f31a44 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py index 7cfb00a86..8bf037c9f 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/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) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py index 2e4fcc0dc..e555bee50 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py +++ b/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, ) diff --git a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py index 7a9284569..11be92b0d 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py +++ b/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 diff --git a/test/registered/debug_utils/comparator/aligner/test_axis_swapper.py b/test/registered/debug_utils/comparator/aligner/test_axis_swapper.py new file mode 100644 index 000000000..78e337802 --- /dev/null +++ b/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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 842a87aa2..2a32b1b49 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/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."""