Support non-packed format when aligning tokens in dump comparator (#19459)

test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py

@@ -23,7 +23,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     UnsharderPlan,
 )
-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -229,7 +229,10 @@ class TestExecuteAlignerPlanWithTokenDim:
             steps=[0, 0, 0],
             token_index_in_step=[0, 1, 2],
         )
-        token_plan = TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y))
+        token_plan = TokenAlignerPlan(
+            locators=Pair(x=locator_x, y=locator_y),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
+        )
 
         plan = AlignerPlan(
             per_step_plans=Pair(
@@ -262,6 +265,54 @@ class TestExecuteAlignerPlanWithTokenDim:
                 plain_y.select(dim=1, index=i),
             )
 
+    def test_bshd_cross_layout_e2e(self) -> None:
+        """x=SGLang THD, y=Megatron BSHD: planner->executor full flow."""
+        torch.manual_seed(42)
+
+        # x side: THD layout, shape [6, 8] (6 tokens, hidden=8), pre-named
+        tensor_x: torch.Tensor = torch.randn(6, 8).refine_names("t", "h")
+
+        # y side: BSHD layout, shape [2, 3, 8] (B=2, S=3, H=8), pre-named
+        tensor_y: torch.Tensor = torch.randn(2, 3, 8).refine_names("b", "s", "h")
+        flat_y: torch.Tensor = tensor_y.rename(None).reshape(6, 8)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 2, 5],
+        )
+        token_plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS),
+        )
+
+        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,
+        )
+
+        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
+
+        assert result.tensors.x.shape == (3, 8)
+        assert result.tensors.y.shape == (3, 8)
+
+        plain_x: torch.Tensor = tensor_x.rename(None)
+        assert torch.equal(result.tensors.x[0], plain_x[0])
+        assert torch.equal(result.tensors.x[1], plain_x[2])
+        assert torch.equal(result.tensors.x[2], plain_x[5])
+
+        assert torch.equal(result.tensors.y[0], flat_y[0])
+        assert torch.equal(result.tensors.y[1], flat_y[2])
+        assert torch.equal(result.tensors.y[2], flat_y[5])
+
 
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py

@@ -15,6 +15,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
 )
 from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -154,6 +155,7 @@ class TestComputeAlignerPlan:
                 x=TokenLocator(steps=[0], token_index_in_step=[0]),
                 y=TokenLocator(steps=[0], token_index_in_step=[0]),
             ),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
         )
 
         plan: AlignerPlan = compute_aligner_plan(