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: