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

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

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

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]

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(

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)

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

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

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

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