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

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