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