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 5354b58f2..7cfb00a86 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py @@ -54,6 +54,7 @@ def execute_aligner_plan( combined: Pair[torch.Tensor] = execute_token_aligner( plan=plan.token_aligner_plan, tensor_of_step_pair=Pair(x=step_tensors_x, y=step_tensors_y), + token_dims=plan.token_dims, ) else: assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1 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 5fec60cfa..2e4fcc0dc 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py @@ -19,7 +19,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 parse_dims +from sglang.srt.debug_utils.comparator.dims import ( + TOKEN_DIM_NAME, + find_dim_index, + parse_dims, +) from sglang.srt.debug_utils.comparator.utils import Pair @@ -28,14 +32,33 @@ def compute_aligner_plan( metas_pair: Pair[list[dict[str, Any]]], token_aligner_plan: Optional[TokenAlignerPlan], ) -> AlignerPlan: + token_dims: Pair[int] = metas_pair.map(_compute_token_dim) 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, ) +def _compute_token_dim(metas: list[dict[str, Any]]) -> int: + fallback_dim = 0 + + if not metas: + return fallback_dim + + dims_str: Optional[str] = metas[0].get("dims") + if dims_str is None: + return fallback_dim + + idx: Optional[int] = find_dim_index(parse_dims(dims_str), TOKEN_DIM_NAME) + if idx is None: + return fallback_dim + + return idx + + def _compute_per_step_plans(metas: list[dict[str, Any]]) -> list[AlignerPerStepPlan]: step_to_input_indices: dict[int, list[int]] = {} for i, meta in enumerate(metas): 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 0a0018bed..7a9284569 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py @@ -24,3 +24,4 @@ class AlignerPerStepPlan: class AlignerPlan: per_step_plans: Pair[list[AlignerPerStepPlan]] token_aligner_plan: Optional[TokenAlignerPlan] + token_dims: Pair[int] diff --git a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py index d792d3a4b..d4c8de358 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/reorderer/planner.py @@ -3,9 +3,14 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( ZigzagToNaturalParams, ) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo -from sglang.srt.debug_utils.comparator.dims import DimSpec, Ordering, ParallelAxis +from sglang.srt.debug_utils.comparator.dims import ( + SEQ_DIM_NAME, + DimSpec, + Ordering, + ParallelAxis, +) -_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {"s"} +_ALLOWED_ZIGZAG_DIM_NAMES: set[str] = {SEQ_DIM_NAME} def compute_reorderer_plans( diff --git a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py index f3e4418c0..11d1eded5 100644 --- a/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py +++ b/python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py @@ -12,31 +12,46 @@ from sglang.srt.debug_utils.comparator.utils import Pair def execute_token_aligner( plan: TokenAlignerPlan, tensor_of_step_pair: Pair[dict[int, torch.Tensor]], + *, + token_dims: Pair[int] = Pair(x=0, y=0), ) -> Pair[torch.Tensor]: if not plan.locators.x.steps: - dummy: torch.Tensor = next(iter(tensor_of_step_pair.x.values())) - empty_shape: list[int] = [0] + list(dummy.shape[1:]) - empty: torch.Tensor = torch.empty(empty_shape, dtype=dummy.dtype) - return Pair(x=empty, y=empty.clone()) + return Pair( + x=_make_empty(tensor_of_step=tensor_of_step_pair.x, token_dim=token_dims.x), + y=_make_empty(tensor_of_step=tensor_of_step_pair.y, token_dim=token_dims.y), + ) return Pair( x=_extract_and_stack_tokens( tensor_of_step=tensor_of_step_pair.x, locator=plan.locators.x, + token_dim=token_dims.x, ), y=_extract_and_stack_tokens( tensor_of_step=tensor_of_step_pair.y, locator=plan.locators.y, + token_dim=token_dims.y, ), ) +def _make_empty( + *, tensor_of_step: dict[int, torch.Tensor], token_dim: int +) -> torch.Tensor: + dummy: torch.Tensor = next(iter(tensor_of_step.values())) + shape: list[int] = list(dummy.shape) + shape[token_dim] = 0 + return torch.empty(shape, dtype=dummy.dtype) + + def _extract_and_stack_tokens( *, tensor_of_step: dict[int, torch.Tensor], locator: TokenLocator, + token_dim: int, ) -> torch.Tensor: tokens: list[torch.Tensor] = [ - tensor_of_step[s][i] for s, i in zip(locator.steps, locator.token_index_in_step) + tensor_of_step[s].select(dim=token_dim, index=i) + for s, i in zip(locator.steps, locator.token_index_in_step) ] - return torch.stack(tokens) + return torch.stack(tokens, dim=token_dim) diff --git a/python/sglang/srt/debug_utils/comparator/dims.py b/python/sglang/srt/debug_utils/comparator/dims.py index 55f6c3b78..3f2519e2f 100644 --- a/python/sglang/srt/debug_utils/comparator/dims.py +++ b/python/sglang/srt/debug_utils/comparator/dims.py @@ -3,6 +3,10 @@ from dataclasses import dataclass from enum import Enum from typing import Optional +TOKEN_DIM_NAME: str = "t" +BATCH_DIM_NAME: str = "b" +SEQ_DIM_NAME: str = "s" + class ParallelAxis(Enum): TP = "tp" @@ -78,3 +82,8 @@ def parse_dims(dims_str: str) -> list[DimSpec]: raise ValueError(f"Duplicate dim names: {duplicates}") return result + + +def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]: + names: list[str] = [spec.name for spec in dim_specs] + return names.index(name) if name in names else None diff --git a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py index 0653cda58..7d3d7d836 100644 --- a/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py @@ -15,6 +15,10 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( AlignerPerStepPlan, AlignerPlan, ) +from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( + TokenAlignerPlan, + TokenLocator, +) from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( ConcatParams, UnsharderPlan, @@ -120,6 +124,7 @@ class TestExecuteAlignerPlan: y=[self._make_step_plan(step=0, indices=[0])], ), token_aligner_plan=None, + token_dims=Pair(x=0, y=0), ) tensors_pair: Pair[list[torch.Tensor]] = Pair( @@ -141,6 +146,7 @@ class TestExecuteAlignerPlan: y=[self._make_step_plan(step=0, indices=[0, 1])], ), token_aligner_plan=None, + token_dims=Pair(x=0, y=0), ) tensors_pair: Pair[list[torch.Tensor]] = Pair( @@ -162,6 +168,7 @@ class TestExecuteAlignerPlan: y=[self._make_step_plan(step=0, indices=[0])], ), token_aligner_plan=None, + token_dims=Pair(x=0, y=0), ) t_x: torch.Tensor = torch.tensor([1.0, 2.0]) @@ -184,6 +191,7 @@ class TestExecuteAlignerPlan: y=[self._make_step_plan(step=0, indices=[0])], ), token_aligner_plan=None, + token_dims=Pair(x=0, y=0), ) tensors_pair: Pair[list[torch.Tensor]] = Pair( @@ -199,5 +207,60 @@ class TestExecuteAlignerPlan: assert result.tensors is not None +class TestExecuteAlignerPlanWithTokenDim: + """End-to-end tests for AlignerPlan with non-zero token_dim.""" + + def _make_step_plan(self, *, step: int, indices: list[int]) -> AlignerPerStepPlan: + return AlignerPerStepPlan(step=step, input_object_indices=indices, sub_plans=[]) + + def test_token_dim_nonzero_e2e(self) -> None: + """AlignerPlan with token_dim=1 passes through to token aligner correctly.""" + torch.manual_seed(42) + + # shape [3, 4, 8]: dim0=batch, dim1=token(4 tokens), dim2=hidden + tensor_x: torch.Tensor = torch.randn(3, 4, 8) + tensor_y: torch.Tensor = torch.randn(3, 4, 8) + + locator_x = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 1, 2], + ) + locator_y = TokenLocator( + steps=[0, 0, 0], + token_index_in_step=[0, 1, 2], + ) + token_plan = TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y)) + + plan = AlignerPlan( + per_step_plans=Pair( + x=[self._make_step_plan(step=0, indices=[0])], + y=[self._make_step_plan(step=0, indices=[0])], + ), + token_aligner_plan=token_plan, + token_dims=Pair(x=1, y=1), + ) + + tensors_pair: Pair[list[torch.Tensor]] = Pair(x=[tensor_x], y=[tensor_y]) + result: AlignerResult = execute_aligner_plan( + tensors_pair=tensors_pair, plan=plan + ) + + assert result.tensors is not None + assert result.failed_side_xy is None + # token dim stays at dim 1 -> shape [3, 3, 8] (3 tokens selected from 4) + assert result.tensors.x.shape == (3, 3, 8) + assert result.tensors.y.shape == (3, 3, 8) + + for i in range(3): + assert torch.equal( + result.tensors.x.select(dim=1, index=i), + tensor_x.select(dim=1, index=i), + ) + assert torch.equal( + result.tensors.y.select(dim=1, index=i), + tensor_y.select(dim=1, index=i), + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py index a1db75e5d..072a8f391 100644 --- a/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py +++ b/test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py @@ -88,5 +88,114 @@ class TestExecuteAlignment: assert aligned.y.shape[1:] == (8,) +class TestTokenDim: + """Tests for non-zero token_dim support.""" + + def _make_simple_plan(self, *, num_tokens: int) -> TokenAlignerPlan: + locator = TokenLocator( + steps=[0] * num_tokens, + token_index_in_step=list(range(num_tokens)), + ) + return TokenAlignerPlan(locators=Pair(x=locator, y=locator)) + + def test_token_dim_nonzero(self) -> None: + """tensor shape [3, 5, 8], token_dim=1 -> token dim stays at dim 1.""" + torch.manual_seed(42) + tensor: torch.Tensor = torch.randn(3, 5, 8) + plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + token_dims=Pair(x=1, y=1), + ) + + assert aligned.x.shape == (3, 5, 8) + assert torch.equal(aligned.x, aligned.y) + for i in range(5): + assert torch.equal( + aligned.x.select(dim=1, index=i), tensor.select(dim=1, index=i) + ) + + def test_token_dim_last(self) -> None: + """tensor shape [3, 8, 5], token_dim=2 -> token dim stays at dim 2.""" + torch.manual_seed(42) + tensor: torch.Tensor = torch.randn(3, 8, 5) + plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + token_dims=Pair(x=2, y=2), + ) + + assert aligned.x.shape == (3, 8, 5) + for i in range(5): + assert torch.equal( + aligned.x.select(dim=2, index=i), tensor.select(dim=2, index=i) + ) + + def test_token_dim_zero(self) -> None: + """token_dim=0 selects along first dimension (standard t-h-d layout).""" + torch.manual_seed(42) + tensor: torch.Tensor = torch.randn(5, 8) + plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + token_dims=Pair(x=0, y=0), + ) + + assert aligned.x.shape == (5, 8) + for i in range(5): + assert torch.equal(aligned.x[i], tensor.select(dim=0, index=i)) + + def test_zero_matched_tokens_nonzero_token_dim(self) -> None: + """Empty plan with token_dim=1 produces correct empty shape.""" + torch.manual_seed(42) + + plan = TokenAlignerPlan( + locators=Pair( + x=TokenLocator(steps=[], token_index_in_step=[]), + y=TokenLocator(steps=[], token_index_in_step=[]), + ), + ) + + # tensor shape [3, 5, 8], token_dim=1 + tensors: dict[int, torch.Tensor] = {0: torch.randn(3, 5, 8)} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + token_dims=Pair(x=1, y=1), + ) + + # token dim (dim 1) set to 0, other dims preserved -> [3, 0, 8] + assert aligned.x.shape == (3, 0, 8) + assert aligned.y.shape == (3, 0, 8) + + def test_high_rank_tensor(self) -> None: + """tensor shape [2, 3, 5, 4, 8] (a b t c d), token_dim=2 -> stays at dim 2.""" + torch.manual_seed(42) + tensor: torch.Tensor = torch.randn(2, 3, 5, 4, 8) + plan: TokenAlignerPlan = self._make_simple_plan(num_tokens=5) + + tensors: dict[int, torch.Tensor] = {0: tensor} + aligned: Pair[torch.Tensor] = execute_token_aligner( + plan=plan, + tensor_of_step_pair=Pair(x=tensors, y=tensors), + token_dims=Pair(x=2, y=2), + ) + + assert aligned.x.shape == (2, 3, 5, 4, 8) + for i in range(5): + assert torch.equal( + aligned.x.select(dim=2, index=i), tensor.select(dim=2, index=i) + ) + + 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 8e5e1a322..34dc5115c 100644 --- a/test/registered/debug_utils/comparator/test_dims.py +++ b/test/registered/debug_utils/comparator/test_dims.py @@ -3,10 +3,14 @@ import sys import pytest from sglang.srt.debug_utils.comparator.dims import ( + BATCH_DIM_NAME, + SEQ_DIM_NAME, + TOKEN_DIM_NAME, DimSpec, Ordering, ParallelAxis, Reduction, + find_dim_index, parse_dim, parse_dims, ) @@ -97,5 +101,41 @@ class TestParseDims: parse_dims("h h") +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 + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))