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

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py

@@ -14,12 +14,6 @@ from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
 from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
 
-_BSHD_NOT_SUPPORTED_MSG: str = (
-    "BSHD layout is not currently supported. "
-    "Use aux_loader BSHD→THD conversion (planned)."
-)
-
-
 # ── plugin ABC ─────────────────────────────────────────────────────
 
 
@@ -153,26 +147,26 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
         return frozenset({"cu_seqlens_q", "cu_seqlens_kv", "qkv_format"})
 
     def has_required_names(self, names: set[str]) -> bool:
-        return "input_ids" in names and "cu_seqlens_q" in names
+        return "input_ids" in names
 
     def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
         for step_data in raw.values():
             if (qkv_format := step_data.get("qkv_format")) is not None:
                 fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
                 if "bshd" in fmt.lower():
-                    raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
+                    return TokenLayout.BS
                 return TokenLayout.T
 
             input_ids = step_data.get("input_ids")
             if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
-                raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
+                return TokenLayout.BS
 
         warning_sink.add(
             GeneralWarning(
                 category="layout_detection_fallback",
                 message=(
                     "Megatron layout detection: no qkv_format or 2D input_ids found, "
-                    "falling back to thd"
+                    "falling back to T"
                 ),
             )
         )
@@ -182,18 +176,27 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
         self, step_data: dict[str, object], *, layout: TokenLayout, step: int
     ) -> TokenAlignerStepAux:
         input_ids: torch.Tensor = step_data["input_ids"]
+        is_bshd: bool = layout == TokenLayout.BS
+
+        # BSHD [B, S] → flat [B*S]; THD [T] stays as-is
+        flat_ids: list[int] = input_ids.reshape(-1).tolist()
 
         if (cu_seqlens_q := step_data.get("cu_seqlens_q")) is not None:
-            seq_lens: torch.Tensor = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+            seq_lens_list: list[int] = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
+        elif is_bshd:
+            seq_lens_list = [input_ids.shape[1]] * input_ids.shape[0]
         else:
-            seq_lens = torch.tensor([input_ids.shape[0]], dtype=torch.long)
+            seq_lens_list = [input_ids.shape[0]]
 
         if (position_ids := step_data.get("position_ids")) is not None:
-            positions: torch.Tensor = position_ids
+            flat_positions: list[int] = position_ids.reshape(-1).tolist()
+        elif is_bshd:
+            flat_positions = list(range(input_ids.shape[1])) * input_ids.shape[0]
         else:
-            positions = _infer_positions(seq_lens=seq_lens)
+            flat_positions = _infer_positions(
+                seq_lens=torch.tensor(seq_lens_list)
+            ).tolist()
 
-        seq_lens_list: list[int] = seq_lens.tolist()
         num_seqs: int = len(seq_lens_list)
         seq_ids: list[SeqId] = [
             PositionalSeqId(step=step, seq_index=seq_index)
@@ -201,8 +204,8 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
         ]
 
         return TokenAlignerStepAux(
-            input_ids=input_ids.tolist(),
-            positions=positions.tolist(),
+            input_ids=flat_ids,
+            positions=flat_positions,
             seq_lens=seq_lens_list,
             seq_ids=seq_ids,
         )

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
 import torch
+from einops import rearrange
 
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerPlan,
     TokenLocator,
 )
 from sglang.srt.debug_utils.comparator.dims import (
+    BATCH_DIM_NAME,
+    SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
+    TokenLayout,
     resolve_dim_by_name,
     strip_dim_names,
 )
@@ -16,42 +20,115 @@ from sglang.srt.debug_utils.comparator.utils import Pair
 _UNNAMED_TOKEN_DIM_FALLBACK: int = 0
 
 
-def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int:
-    if tensor.names[0] is None:
-        return _UNNAMED_TOKEN_DIM_FALLBACK
-    return resolve_dim_by_name(tensor, name)
-
-
 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]:
+    flat_pair: Pair[dict[int, torch.Tensor]] = Pair(
+        x=_collapse_bs_to_t(
+            tensor_of_step=tensor_of_step_pair.x, layout=plan.layouts.x
+        ),
+        y=_collapse_bs_to_t(
+            tensor_of_step=tensor_of_step_pair.y, layout=plan.layouts.y
+        ),
+    )
+
     if not plan.locators.x.steps:
         return Pair(
-            x=_make_empty(tensor_of_step=tensor_of_step_pair.x),
-            y=_make_empty(tensor_of_step=tensor_of_step_pair.y),
+            x=_make_empty(tensor_of_step=flat_pair.x),
+            y=_make_empty(tensor_of_step=flat_pair.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,
+            tensor_of_step=flat_pair.x, locator=plan.locators.x
         ),
         y=_extract_and_stack_tokens(
-            tensor_of_step=tensor_of_step_pair.y,
-            locator=plan.locators.y,
-            token_dim=token_dims.y,
+            tensor_of_step=flat_pair.y, locator=plan.locators.y
         ),
     )
 
 
-def _make_empty(
+# ── BS → T preprocessing ─────────────────────────────────────────
+
+
+def _collapse_bs_to_t(
     *,
     tensor_of_step: dict[int, torch.Tensor],
-) -> torch.Tensor:
+    layout: TokenLayout,
+) -> dict[int, torch.Tensor]:
+    """Collapse B and S dims into a single flat token dim (always batch-major).
+
+    Handles both ``b s`` and ``s b`` orderings correctly via einops rearrange.
+    Returns the original tensors unchanged if layout is T.
+    """
+    if layout != TokenLayout.BS:
+        return tensor_of_step
+
+    some_tensor: torch.Tensor = next(iter(tensor_of_step.values()))
+    batch_dim: int = _resolve_dim_or_fallback(some_tensor, BATCH_DIM_NAME)
+    seq_dim: int = _resolve_dim_or_fallback(some_tensor, SEQ_DIM_NAME)
+
+    if abs(batch_dim - seq_dim) != 1:
+        raise ValueError(
+            f"BS dims must be adjacent: "
+            f"{BATCH_DIM_NAME}={batch_dim}, "
+            f"{SEQ_DIM_NAME}={seq_dim}"
+        )
+
+    lhs_pattern, rhs_pattern, new_names = _build_bs_collapse_pattern(
+        names=list(some_tensor.names),
+        batch_dim=batch_dim,
+        seq_dim=seq_dim,
+    )
+
+    result: dict[int, torch.Tensor] = {}
+    for step, tensor in tensor_of_step.items():
+        plain: torch.Tensor = strip_dim_names(tensor)
+        collapsed: torch.Tensor = rearrange(plain, f"{lhs_pattern} -> {rhs_pattern}")
+        result[step] = collapsed.refine_names(*new_names)
+
+    return result
+
+
+def _build_bs_collapse_pattern(
+    *,
+    names: list[str | None],
+    batch_dim: int,
+    seq_dim: int,
+) -> tuple[str, str, list[str | None]]:
+    """Build einops lhs/rhs patterns and output dim names for BS→T collapse.
+
+    Always produces batch-major order ``(b s)`` regardless of input ordering.
+    Uses the tensor's own dim names as einops axis names.
+    """
+    lo: int = min(batch_dim, seq_dim)
+    hi: int = max(batch_dim, seq_dim)
+
+    lhs: str = " ".join(names)  # type: ignore[arg-type]
+
+    rhs_names: list[str] = list(names[:lo]) + [f"({BATCH_DIM_NAME} {SEQ_DIM_NAME})"] + list(names[hi + 1 :])  # type: ignore[misc]
+    rhs: str = " ".join(rhs_names)
+
+    new_names: list[str | None] = (
+        list(names[:lo]) + [TOKEN_DIM_NAME] + list(names[hi + 1 :])
+    )
+
+    return lhs, rhs, new_names
+
+
+# ── core logic (T layout only) ───────────────────────────────────
+
+
+def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int:
+    if tensor.names[0] is None:
+        return _UNNAMED_TOKEN_DIM_FALLBACK
+    return resolve_dim_by_name(tensor, name)
+
+
+def _make_empty(*, tensor_of_step: dict[int, torch.Tensor]) -> torch.Tensor:
     dummy: torch.Tensor = next(iter(tensor_of_step.values()))
     token_dim: int = _resolve_dim_or_fallback(dummy, TOKEN_DIM_NAME)
     shape: list[int] = list(dummy.shape)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py

@@ -48,7 +48,10 @@ def compute_token_aligner_plan(
             token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
         )
 
-    return TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y))
+    return TokenAlignerPlan(
+        locators=Pair(x=locator_x, y=locator_y),
+        layouts=seqs_info_pair.map(lambda s: s.layout),
+    )
 
 
 # -------------------- Sequence matcher --------------------

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py

@@ -114,6 +114,7 @@ class TokenAlignerPlan(_FrozenBase):
     """Token alignment plan. locators.x[i] and locators.y[i] correspond to the same logical token."""
 
     locators: Pair[TokenLocator]
+    layouts: Pair[TokenLayout]
 
     @model_validator(mode="after")
     def _validate_fields(self) -> TokenAlignerPlan:

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(

test/registered/debug_utils/comparator/aligner/token_aligner/test_aux_plugins.py

@@ -130,6 +130,79 @@ class TestNormalizeMegatron:
         ]
 
 
+class TestDetectLayoutMegatron:
+    """Tests for Megatron layout detection."""
+
+    def test_detect_layout_bshd_via_qkv_format(self):
+        """qkv_format containing 'bshd' → layout 'bshd'."""
+        raw: dict[int, dict[str, object]] = {
+            0: {"qkv_format": "bshd", "input_ids": torch.tensor([1, 2, 3])}
+        }
+        assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS
+
+    def test_detect_layout_bshd_via_ndim(self):
+        """2D input_ids → layout BS."""
+        raw: dict[int, dict[str, object]] = {
+            0: {"input_ids": torch.tensor([[1, 2], [3, 4]])}
+        }
+        assert _megatron_plugin.detect_layout(raw) == TokenLayout.BS
+
+    def test_detect_layout_thd_via_qkv_format(self):
+        """qkv_format 'thd' → layout T."""
+        raw: dict[int, dict[str, object]] = {
+            0: {"qkv_format": "thd", "input_ids": torch.tensor([1, 2, 3])}
+        }
+        assert _megatron_plugin.detect_layout(raw) == TokenLayout.T
+
+
+class TestNormalizeMegatronBSHD:
+    """Tests for Megatron BSHD normalization."""
+
+    def test_basic_bshd(self):
+        """2D input_ids [2,4] → flat [8], seq_lens=[4,4], positions=[0,1,2,3,0,1,2,3]."""
+        step_data: dict = {
+            "input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout=TokenLayout.BS, step=0
+        )
+
+        assert result.input_ids == [10, 20, 30, 40, 50, 60, 70, 80]
+        assert result.seq_lens == [4, 4]
+        assert result.positions == [0, 1, 2, 3, 0, 1, 2, 3]
+        assert result.seq_ids == [
+            PositionalSeqId(step=0, seq_index=0),
+            PositionalSeqId(step=0, seq_index=1),
+        ]
+
+    def test_bshd_with_cu_seqlens(self):
+        """BSHD with cu_seqlens_q → uses cu_seqlens for seq_lens."""
+        step_data: dict = {
+            "input_ids": torch.tensor([[10, 20, 30, 40], [50, 60, 70, 80]]),
+            "cu_seqlens_q": torch.tensor([0, 3, 8]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout=TokenLayout.BS, step=0
+        )
+
+        assert result.seq_lens == [3, 5]
+
+    def test_bshd_with_position_ids(self):
+        """BSHD with 2D position_ids → flattened positions."""
+        step_data: dict = {
+            "input_ids": torch.tensor([[10, 20], [30, 40]]),
+            "position_ids": torch.tensor([[5, 6], [10, 11]]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout=TokenLayout.BS, step=0
+        )
+
+        assert result.positions == [5, 6, 10, 11]
+
+
 class TestInferPositions:
     """Tests for position inference helper."""
 

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

@@ -80,6 +80,7 @@ class TestExecuteAlignment:
                 x=TokenLocator(steps=[], token_index_in_step=[]),
                 y=TokenLocator(steps=[], token_index_in_step=[]),
             ),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
         )
 
         tensors = {0: torch.randn(5, 8).refine_names("t", "h")}
@@ -101,7 +102,10 @@ class TestTokenDim:
             steps=[0] * num_tokens,
             token_index_in_step=list(range(num_tokens)),
         )
-        return TokenAlignerPlan(locators=Pair(x=locator, y=locator))
+        return TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
+        )
 
     def test_token_dim_nonzero(self) -> None:
         """tensor shape [3, 5, 8], token_dim=1 -> token dim stays at dim 1."""
@@ -168,6 +172,7 @@ class TestTokenDim:
                 x=TokenLocator(steps=[], token_index_in_step=[]),
                 y=TokenLocator(steps=[], token_index_in_step=[]),
             ),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
         )
 
         tensors: dict[int, torch.Tensor] = {
@@ -204,5 +209,246 @@ class TestTokenDim:
             )
 
 
+class TestBSHDExecutor:
+    """BSHD tensor collapse: B+S dims -> flat token dim for alignment."""
+
+    def test_bshd_standard_bs_at_front(self):
+        """Standard "b s h d": B=dim0, S=dim1. [2, 3, 4, 5] -> collapse -> [6, 4, 5]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(torch.randn(2, 3, 4, 5), ["b", "s", "h", "d"])
+        flat: torch.Tensor = tensor.rename(None).reshape(6, 4, 5)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 3, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (3, 4, 5)
+        assert torch.equal(aligned.x[0], flat[0])
+        assert torch.equal(aligned.x[1], flat[3])
+        assert torch.equal(aligned.x[2], flat[5])
+
+    def test_bshd_3d_bs_at_front(self):
+        """Minimal 3D "b s h": B=dim0, S=dim1. [2, 3, 4] -> collapse -> [6, 4]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(torch.randn(2, 3, 4), ["b", "s", "h"])
+        flat: torch.Tensor = tensor.rename(None).reshape(6, 4)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0, 0],
+            token_index_in_step=[0, 2, 3, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (4, 4)
+        assert torch.equal(aligned.x[0], flat[0])
+        assert torch.equal(aligned.x[1], flat[2])
+        assert torch.equal(aligned.x[2], flat[3])
+        assert torch.equal(aligned.x[3], flat[5])
+
+    def test_bshd_bs_not_at_front(self):
+        """Non-leading "h b s d": B=dim1, S=dim2. [4, 2, 3, 5] -> collapse -> [4, 6, 5]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"])
+        flat: torch.Tensor = tensor.rename(None).reshape(4, 6, 5)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 3, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (4, 3, 5)
+        for idx, flat_idx in enumerate([0, 3, 5]):
+            assert torch.equal(
+                aligned.x.select(dim=1, index=idx),
+                flat.select(dim=1, index=flat_idx),
+            )
+
+    def test_bshd_expert_before_bs(self):
+        """Expert dim before B: "e b s h d". [2, 3, 4, 5, 6] -> collapse -> [2, 12, 5, 6]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(
+            torch.randn(2, 3, 4, 5, 6), ["e", "b", "s", "h", "d"]
+        )
+        flat: torch.Tensor = tensor.rename(None).reshape(2, 12, 5, 6)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 5, 11],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (2, 3, 5, 6)
+        for idx, flat_idx in enumerate([0, 5, 11]):
+            assert torch.equal(
+                aligned.x.select(dim=1, index=idx),
+                flat.select(dim=1, index=flat_idx),
+            )
+
+    def test_bshd_bs_at_end(self):
+        """B and S at end: "h d b s". [4, 5, 2, 3] -> collapse -> [4, 5, 6]."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(torch.randn(4, 5, 2, 3), ["h", "d", "b", "s"])
+        flat: torch.Tensor = tensor.rename(None).reshape(4, 5, 6)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[1, 3, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (4, 5, 3)
+        for idx, flat_idx in enumerate([1, 3, 5]):
+            assert torch.equal(
+                aligned.x.select(dim=2, index=idx),
+                flat.select(dim=2, index=flat_idx),
+            )
+
+    def test_cross_layout_thd_vs_bshd(self):
+        """Cross-layout: x=THD [6, 8], y=BSHD [2, 3, 8] -> y collapse -> [6, 8]."""
+        torch.manual_seed(42)
+        tensor_thd: torch.Tensor = _named(torch.randn(6, 8), ["t", "h"])
+        tensor_bshd: torch.Tensor = _named(torch.randn(2, 3, 8), ["b", "s", "h"])
+        flat_bshd: torch.Tensor = tensor_bshd.rename(None).reshape(6, 8)
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 2, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.BS),
+        )
+
+        aligned: Pair[torch.Tensor] = execute_token_aligner(
+            plan=plan,
+            tensor_of_step_pair=Pair(x={0: tensor_thd}, y={0: tensor_bshd}),
+        )
+
+        assert aligned.x.shape == (3, 8)
+        assert aligned.y.shape == (3, 8)
+        assert torch.equal(aligned.x[0], tensor_thd.rename(None)[0])
+        assert torch.equal(aligned.y[0], flat_bshd[0])
+        assert torch.equal(aligned.y[2], flat_bshd[5])
+
+    def test_bshd_reversed_sb_order(self):
+        """Reversed "s b h": S=dim0, B=dim1. Collapse is batch-major: (b s)."""
+        torch.manual_seed(42)
+        tensor: torch.Tensor = _named(torch.randn(3, 2, 4), ["s", "b", "h"])
+        # batch-major flatten: rearrange("s b h -> (b s) h")
+        from einops import rearrange
+
+        flat: torch.Tensor = rearrange(tensor.rename(None), "s b h -> (b s) h")
+
+        locator = TokenLocator(
+            steps=[0, 0, 0],
+            token_index_in_step=[0, 2, 5],
+        )
+        plan = TokenAlignerPlan(
+            locators=Pair(x=locator, y=locator),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        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),
+        )
+
+        assert aligned.x.shape == (3, 4)
+        assert torch.equal(aligned.x[0], flat[0])
+        assert torch.equal(aligned.x[1], flat[2])
+        assert torch.equal(aligned.x[2], flat[5])
+
+    def test_bshd_empty_plan_bs_not_at_front(self):
+        """Empty plan with non-leading B,S: "h b s d". [4, 2, 3, 5] -> collapse -> [4, 0, 5]."""
+        plan = TokenAlignerPlan(
+            locators=Pair(
+                x=TokenLocator(steps=[], token_index_in_step=[]),
+                y=TokenLocator(steps=[], token_index_in_step=[]),
+            ),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        tensors: dict[int, torch.Tensor] = {
+            0: _named(torch.randn(4, 2, 3, 5), ["h", "b", "s", "d"])
+        }
+        aligned: Pair[torch.Tensor] = execute_token_aligner(
+            plan=plan,
+            tensor_of_step_pair=Pair(x=tensors, y=tensors),
+        )
+
+        assert aligned.x.shape == (4, 0, 5)
+        assert aligned.y.shape == (4, 0, 5)
+
+    def test_bshd_empty_plan_bs_at_front(self):
+        """Empty plan with standard BSHD: "b s h". [2, 3, 4] -> collapse -> [0, 4]."""
+        plan = TokenAlignerPlan(
+            locators=Pair(
+                x=TokenLocator(steps=[], token_index_in_step=[]),
+                y=TokenLocator(steps=[], token_index_in_step=[]),
+            ),
+            layouts=Pair(x=TokenLayout.BS, y=TokenLayout.BS),
+        )
+
+        tensors: dict[int, torch.Tensor] = {
+            0: _named(torch.randn(2, 3, 4), ["b", "s", "h"])
+        }
+        aligned: Pair[torch.Tensor] = execute_token_aligner(
+            plan=plan,
+            tensor_of_step_pair=Pair(x=tensors, y=tensors),
+        )
+
+        assert aligned.x.shape == (0, 4)
+        assert aligned.y.shape == (0, 4)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

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

@@ -369,8 +369,8 @@ class TestMatchSequences:
         assert matched == []
 
 
-class TestComputeAlignmentPlanCrossLayout:
-    """Tests for alignment plan across different step distributions."""
+class TestComputeAlignmentPlanCrossFramework:
+    """Tests for alignment plan across different frameworks and layouts."""
 
     def test_thd_vs_thd_different_step_splits(self):
         """Two thd sides with same tokens but different step distributions."""
@@ -454,6 +454,57 @@ class TestComputeAlignmentPlanCrossLayout:
 
         assert len(plan.locators.x.steps) == 7
 
+    def test_cross_layout_sglang_thd_vs_megatron_bshd(self):
+        """SGLang THD vs Megatron BSHD end-to-end alignment via planner.
+
+        SGLang side: two sequences [10,20,30] and [40,50] across 2 steps.
+        Megatron BSHD side: same tokens as 2 batch slots [10,20,30,PAD] and [40,50,PAD,PAD],
+        where PAD tokens (99) are included because BSHD treats whole padded row as one seq.
+        Planner should match by prefix and align the common 5 tokens.
+        """
+        side_sglang = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30, 40, 50],
+                    positions=[0, 1, 2, 0, 1],
+                    seq_lens=[3, 2],
+                    seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
+                ),
+            },
+            framework="sglang",
+            layout=TokenLayout.T,
+        )
+
+        side_megatron_bshd = TokenAlignerGlobalAux(
+            step_auxs={
+                # BSHD normalized: flat [B*S] with each batch slot as one seq
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30, 99, 40, 50, 99, 99],
+                    positions=[0, 1, 2, 3, 0, 1, 2, 3],
+                    seq_lens=[4, 4],
+                    seq_ids=[
+                        PositionalSeqId(step=0, seq_index=0),
+                        PositionalSeqId(step=0, seq_index=1),
+                    ],
+                ),
+            },
+            framework="megatron",
+            layout=TokenLayout.BS,
+        )
+
+        index_sglang = build_seqs_info(side_sglang)
+        index_megatron = build_seqs_info(side_megatron_bshd)
+
+        plan = compute_token_aligner_plan(
+            seqs_info_pair=Pair(x=index_sglang, y=index_megatron)
+        )
+
+        # Seq A: [10,20,30] matches prefix of [10,20,30,99] → 3 tokens
+        # Seq B: [40,50] matches prefix of [40,50,99,99] → 2 tokens
+        assert len(plan.locators.x.steps) == 5
+        assert plan.layouts.x == TokenLayout.T
+        assert plan.layouts.y == TokenLayout.BS
+
 
 # ---------------------------------------------------------------------------
 # Helpers

test/registered/debug_utils/comparator/test_model_validation.py

@@ -11,6 +11,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenLocator,
 )
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     GeneralWarning,
@@ -120,6 +121,7 @@ class TestTokenAlignerPlan:
                 x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
                 y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
             ),
+            layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
         )
         assert len(plan.locators.x.steps) == 3
 
@@ -130,6 +132,7 @@ class TestTokenAlignerPlan:
                     x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]),
                     y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
                 ),
+                layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
             )