chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support unifying axis ordering in dump comparator (#19456)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:08:32 +08:00
committed by GitHub
parent 425d333ee3
commit 8ac64e1487
6 changed files with 230 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
python/sglang/srt/debug_utils/comparator/aligner/axis_swapper.py Normal file
View File

@@ -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)

10
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py
View File

@@ -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)

13
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
View File

@@ -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,
)

2
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py
View File

@@ -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

74
test/registered/debug_utils/comparator/aligner/test_axis_swapper.py Normal file
View File

@@ -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__]))

74
test/registered/debug_utils/comparator/test_entrypoint.py
View File

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