Support token dim in arbitrary location in dump comparator (#19455)

2026-02-27 08:07:38 +08:00
parent 5172c37845
commit 425d333ee3
9 changed files with 275 additions and 9 deletions
@@ -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__]))
@@ -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__]))
@@ -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__]))