diff --git a/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py b/python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py new file mode 100644 index 000000000..67dea458b --- /dev/null +++ b/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 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 ce4960b01..e9781711e 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/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) 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 9859dec2b..e7aaf20fa 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py +++ b/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( 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 dc9954652..84be6b27f 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py +++ b/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 diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py index 6285c396f..20b1051de 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py +++ b/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) diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 58a093211..a14936c1f 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/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] diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py index 084340524..0e74e30d3 100644 --- a/python/sglang/srt/debug_utils/comparator/dims.py +++ b/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[a-zA-Z_]\w*)(?:\((?P[^)]+)\))?$") _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]: diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index 61a0f8ff8..1aa784d91 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/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) diff --git a/test/registered/debug_utils/comparator/aligner/test_axis_aligner.py b/test/registered/debug_utils/comparator/aligner/test_axis_aligner.py new file mode 100644 index 000000000..44981c6cd --- /dev/null +++ b/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__])) diff --git a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py b/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py index c1cbf5a29..eca1fe27a 100644 --- a/test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py +++ b/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__])) diff --git a/test/registered/debug_utils/comparator/test_dims.py b/test/registered/debug_utils/comparator/test_dims.py index df3dabd89..0717e6800 100644 --- a/test/registered/debug_utils/comparator/test_dims.py +++ b/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: diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 97751663b..0be688b21 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/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.""" diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index d003a38e3..1bbc1973a 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/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__])) diff --git a/test/registered/debug_utils/source_patcher/test_code_patcher.py b/test/registered/debug_utils/source_patcher/test_code_patcher.py index 56e0ab721..a6fa9270f 100644 --- a/test/registered/debug_utils/source_patcher/test_code_patcher.py +++ b/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" diff --git a/test/registered/debug_utils/source_patcher/test_dumper_integration.py b/test/registered/debug_utils/source_patcher/test_dumper_integration.py index 70a06eb45..278003525 100644 --- a/test/registered/debug_utils/source_patcher/test_dumper_integration.py +++ b/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" diff --git a/test/registered/debug_utils/source_patcher/test_source_editor.py b/test/registered/debug_utils/source_patcher/test_source_editor.py index 10510732f..ed75e6fd2 100644 --- a/test/registered/debug_utils/source_patcher/test_source_editor.py +++ b/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: