Support multi-step alignment and pipeline integration in dump comparator (#19378)

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

@@ -51,13 +51,9 @@ def execute_aligner_plan(
 
     # Cross-side: token alignment (or direct extraction for single-step)
     if plan.token_aligner_plan is not None:
-        assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
         combined: Pair[torch.Tensor] = execute_token_aligner(
             plan=plan.token_aligner_plan,
-            tensor_pair=Pair(
+            tensor_of_step_pair=Pair(x=step_tensors_x, y=step_tensors_y),
-                x=list(step_tensors_x.values())[0],
-                y=list(step_tensors_y.values())[0],
-            ),
         )
     else:
         assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1

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

@@ -43,7 +43,6 @@ def load_and_normalize_aux(
 
     available_names: set[str] = set(df["name"].unique().to_list()) & plugin.all_names
     steps: list[int] = sorted(df["step"].unique().to_list())
-    assert len(steps) == 1, f"Multi-step not yet supported, got {len(steps)} steps"
     tensor_names: set[str] = available_names & plugin.tensor_names
     non_tensor_names: set[str] = available_names & plugin.non_tensor_names
 

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

@@ -4,20 +4,39 @@ import torch
 
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
     TokenAlignerPlan,
+    TokenLocator,
 )
 from sglang.srt.debug_utils.comparator.utils import Pair
 
 
 def execute_token_aligner(
     plan: TokenAlignerPlan,
-    tensor_pair: Pair[torch.Tensor],
+    tensor_of_step_pair: Pair[dict[int, torch.Tensor]],
 ) -> Pair[torch.Tensor]:
-    if not plan.locators.x.token_index_in_step:
+    if not plan.locators.x.steps:
-        empty_shape: list[int] = [0] + list(tensor_pair.x.shape[1:])
+        dummy: torch.Tensor = next(iter(tensor_of_step_pair.x.values()))
-        empty: torch.Tensor = torch.empty(empty_shape, dtype=tensor_pair.x.dtype)
+        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=tensor_pair.x[plan.locators.x.token_index_in_step],
+        x=_extract_and_stack_tokens(
-        y=tensor_pair.y[plan.locators.y.token_index_in_step],
+            tensor_of_step=tensor_of_step_pair.x,
+            locator=plan.locators.x,
+        ),
+        y=_extract_and_stack_tokens(
+            tensor_of_step=tensor_of_step_pair.y,
+            locator=plan.locators.y,
+        ),
     )
+
+
+def _extract_and_stack_tokens(
+    *,
+    tensor_of_step: dict[int, torch.Tensor],
+    locator: TokenLocator,
+) -> torch.Tensor:
+    tokens: list[torch.Tensor] = [
+        tensor_of_step[s][i] for s, i in zip(locator.steps, locator.token_index_in_step)
+    ]
+    return torch.stack(tokens)

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

@@ -21,7 +21,7 @@ def compute_token_aligner_plan(
         seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences)
     )
 
-    _empty = TokenLocator(token_index_in_step=[])
+    _empty = TokenLocator(steps=[], token_index_in_step=[])
     locator_x: TokenLocator = _empty
     locator_y: TokenLocator = _empty
 
@@ -40,9 +40,11 @@ def compute_token_aligner_plan(
         assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}"
 
         locator_x = locator_x + TokenLocator(
+            steps=rec.x.locator.steps[:common_len],
             token_index_in_step=rec.x.locator.token_index_in_step[:common_len],
         )
         locator_y = locator_y + TokenLocator(
+            steps=rec.y.locator.steps[:common_len],
             token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
         )
 

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

@@ -18,6 +18,7 @@ class _SeqInfoAccumulator:
 
     input_ids: list[int] = field(default_factory=list)
     positions: list[int] = field(default_factory=list)
+    steps: list[int] = field(default_factory=list)
     token_index_in_step: list[int] = field(default_factory=list)
 
     def extend(
@@ -25,10 +26,12 @@ class _SeqInfoAccumulator:
         *,
         input_ids: list[int],
         positions: list[int],
+        steps: list[int],
         token_index_in_step: list[int],
     ) -> None:
         self.input_ids.extend(input_ids)
         self.positions.extend(positions)
+        self.steps.extend(steps)
         self.token_index_in_step.extend(token_index_in_step)
 
     def build(self) -> TokenAlignerSeqInfo:
@@ -36,6 +39,7 @@ class _SeqInfoAccumulator:
             input_ids=self.input_ids,
             positions=self.positions,
             locator=TokenLocator(
+                steps=self.steps,
                 token_index_in_step=self.token_index_in_step,
             ),
         )
@@ -68,6 +72,7 @@ def _build_token_aligner_seq_infos(
             accum[seq_id].extend(
                 input_ids=aux.input_ids[offset : offset + seq_len],
                 positions=aux.positions[offset : offset + seq_len],
+                steps=[step] * seq_len,
                 token_index_in_step=list(range(offset, offset + seq_len)),
             )
 

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

@@ -54,15 +54,17 @@ class TokenAlignerGlobalAux:
 
 
 class TokenLocator(_FrozenBase):
-    """Locates tokens within a single-step tensor.
+    """Locates tokens within a multi-step tensor store.
 
-    token i is at tensor[token_index_in_step[i]].
+    token i is at tensor_of_step[steps[i]][token_index_in_step[i]].
     """
 
+    steps: list[int]
     token_index_in_step: list[int]
 
     def __add__(self, other: TokenLocator) -> TokenLocator:
         return TokenLocator(
+            steps=self.steps + other.steps,
             token_index_in_step=self.token_index_in_step + other.token_index_in_step,
         )
 
@@ -81,6 +83,7 @@ class TokenAlignerSeqInfo(_FrozenBase):
         _check_equal_lengths(
             input_ids=self.input_ids,
             positions=self.positions,
+            locator_steps=self.locator.steps,
             locator_token_index_in_step=self.locator.token_index_in_step,
         )
 
@@ -114,7 +117,9 @@ class TokenAlignerPlan(_FrozenBase):
     @model_validator(mode="after")
     def _validate_fields(self) -> TokenAlignerPlan:
         _check_equal_lengths(
+            locators_x_steps=self.locators.x.steps,
             locators_x_token_index_in_step=self.locators.x.token_index_in_step,
+            locators_y_steps=self.locators.y.steps,
             locators_y_token_index_in_step=self.locators.y.token_index_in_step,
         )
         return self

python/sglang/srt/debug_utils/comparator/bundle_comparator.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import torch
 
@@ -15,6 +15,9 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
     compute_aligner_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+    TokenAlignerPlan,
+)
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
     SkipRecord,
@@ -35,6 +38,7 @@ def compare_bundle_pair(
     filenames_pair: Pair[list[str]],
     baseline_path: Path,
     target_path: Path,
+    token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold: float,
 ) -> Union[ComparisonRecord, SkipRecord]:
     with warning_sink.context() as collected_warnings:
@@ -43,6 +47,7 @@ def compare_bundle_pair(
             filenames_pair=filenames_pair,
             baseline_path=baseline_path,
             target_path=target_path,
+            token_aligner_plan=token_aligner_plan,
             diff_threshold=diff_threshold,
         )
 
@@ -55,6 +60,7 @@ def _compare_bundle_pair_raw(
     filenames_pair: Pair[list[str]],
     baseline_path: Path,
     target_path: Path,
+    token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold: float,
 ) -> Union[ComparisonRecord, SkipRecord]:
     # 1. Load (tensor + meta, ungrouped)
@@ -72,7 +78,7 @@ def _compare_bundle_pair_raw(
         lambda items: [it.meta for it in items]
     )
     plan: AlignerPlan = compute_aligner_plan(
-        metas_pair=metas_pair, token_aligner_plan=None
+        metas_pair=metas_pair, token_aligner_plan=token_aligner_plan
     )
 
     # 3. Execute (tensor + plan only, no meta)

python/sglang/srt/debug_utils/comparator/entrypoint.py

@@ -2,10 +2,19 @@ from __future__ import annotations
 
 import argparse
 from pathlib import Path
-from typing import Iterator, Union
+from typing import Iterator, Optional, Union
 
 import polars as pl
 
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
+    AUX_NAMES,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.entrypoint import (
+    compute_maybe_token_aligner_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+    TokenAlignerPlan,
+)
 from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair
 from sglang.srt.debug_utils.comparator.bundle_matcher import (
     TensorBundleInfo,
@@ -37,16 +46,22 @@ def run(args: argparse.Namespace) -> None:
     warning_sink.set_output_format(args.output_format)
 
     dfs: Pair[pl.DataFrame] = _read_df(args)
+    token_aligner_plan = compute_maybe_token_aligner_plan(args, dfs)
+
+    dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES)))
 
     bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
         dfs=dfs,
-        skip_keys=_compute_skip_keys(args),
+        skip_keys=_compute_skip_keys(
+            args, has_token_aligner_plan=token_aligner_plan is not None
+        ),
     )
 
     comparison_records = _compare_bundle_pairs(
         bundle_info_pairs=bundle_info_pairs,
         baseline_path=Path(args.baseline_path),
         target_path=Path(args.target_path),
+        token_aligner_plan=token_aligner_plan,
         diff_threshold=args.diff_threshold,
     )
     _consume_comparison_records(
@@ -68,10 +83,12 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
     return Pair(x=df_baseline, y=df_target)
 
 
-def _compute_skip_keys(args: argparse.Namespace) -> set[str]:
+def _compute_skip_keys(args, *, has_token_aligner_plan: bool):
     skip_keys: set[str] = {"dump_index", "filename"}
     if args.grouping == "logical":
         skip_keys |= {"rank"}
+        if has_token_aligner_plan:
+            skip_keys |= {"step"}
     return skip_keys
 
 
@@ -80,6 +97,7 @@ def _compare_bundle_pairs(
     bundle_info_pairs: list[Pair[TensorBundleInfo]],
     baseline_path: Path,
     target_path: Path,
+    token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold: float,
 ) -> Iterator[Union[ComparisonRecord, SkipRecord]]:
     for bundle_info_pair in bundle_info_pairs:
@@ -95,6 +113,7 @@ def _compare_bundle_pairs(
             filenames_pair=filenames_pair,
             baseline_path=baseline_path,
             target_path=target_path,
+            token_aligner_plan=token_aligner_plan,
             diff_threshold=diff_threshold,
         )
 

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

@@ -151,8 +151,8 @@ class TestComputeAlignerPlan:
 
         ta_plan = TokenAlignerPlan(
             locators=Pair(
-                x=TokenLocator(token_index_in_step=[0]),
+                x=TokenLocator(steps=[0], token_index_in_step=[0]),
-                y=TokenLocator(token_index_in_step=[0]),
+                y=TokenLocator(steps=[0], token_index_in_step=[0]),
             ),
         )
 

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

@@ -28,12 +28,13 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
 
 
 class TestExecuteAlignment:
-    """Tests for token alignment execution (single-step)."""
+    """Tests for token alignment execution."""
 
     def test_thd_vs_thd_identity(self):
         """Two identical thd sides produce element-wise equal aligned tensors."""
         torch.manual_seed(42)
-        hidden = torch.randn(5, 8)  # 5 tokens, hidden_dim=8
+        hidden_step0 = torch.randn(5, 8)  # 5 tokens, hidden_dim=8
+        hidden_step1 = torch.randn(2, 8)  # 2 tokens
 
         aux = TokenAlignerStepAux(
             input_ids=[10, 20, 30, 40, 50],
@@ -41,9 +42,15 @@ class TestExecuteAlignment:
             seq_lens=[3, 2],
             seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
         )
+        aux_step1 = TokenAlignerStepAux(
+            input_ids=[31, 51],
+            positions=[3, 2],
+            seq_lens=[1, 1],
+            seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
+        )
 
         side_aux = TokenAlignerGlobalAux(
-            step_auxs={0: aux},
+            step_auxs={0: aux, 1: aux_step1},
             framework="sglang",
             layout="thd",
         )
@@ -51,12 +58,13 @@ class TestExecuteAlignment:
         index = build_seqs_info(side_aux)
         plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index, y=index))
 
+        tensors = {0: hidden_step0, 1: hidden_step1}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
-            plan=plan, tensor_pair=Pair(x=hidden, y=hidden)
+            plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors)
         )
 
         assert torch.equal(aligned.x, aligned.y)
-        assert aligned.x.shape[0] == len(plan.locators.x.token_index_in_step)
+        assert aligned.x.shape[0] == len(plan.locators.x.steps)
 
     def test_zero_matched_tokens(self):
         """Empty TokenAlignerPlan (no matched tokens) returns shape[0]==0 without crash."""
@@ -64,14 +72,14 @@ class TestExecuteAlignment:
 
         plan = TokenAlignerPlan(
             locators=Pair(
-                x=TokenLocator(token_index_in_step=[]),
+                x=TokenLocator(steps=[], token_index_in_step=[]),
-                y=TokenLocator(token_index_in_step=[]),
+                y=TokenLocator(steps=[], token_index_in_step=[]),
             ),
         )
 
-        tensor = torch.randn(5, 8)
+        tensors = {0: torch.randn(5, 8)}
         aligned: Pair[torch.Tensor] = execute_token_aligner(
-            plan=plan, tensor_pair=Pair(x=tensor, y=tensor)
+            plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors)
         )
 
         assert aligned.x.shape[0] == 0

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

@@ -4,6 +4,7 @@ import pytest
 
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
     _match_sequences,
+    compute_token_aligner_plan,
 )
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
     build_seqs_info,
@@ -48,6 +49,7 @@ class TestBuildTokenIndexSGLangThd:
         seq_a = index.sequences[SGLangSeqId(rid="A")]
         assert seq_a.input_ids == [10, 20, 30]
         assert seq_a.positions == [0, 1, 2]
+        assert seq_a.locator.steps == [0, 0, 0]
         assert seq_a.locator.token_index_in_step == [0, 1, 2]
 
         seq_b = index.sequences[SGLangSeqId(rid="B")]
@@ -55,6 +57,93 @@ class TestBuildTokenIndexSGLangThd:
         assert seq_b.positions == [0, 1]
         assert seq_b.locator.token_index_in_step == [3, 4]
 
+    def test_multi_step_prefill_decode(self):
+        """Prefill step followed by decode steps, sequences accumulate tokens."""
+        side_aux = 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")],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[31, 51],
+                    positions=[3, 2],
+                    seq_lens=[1, 1],
+                    seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+
+        index = build_seqs_info(side_aux)
+        assert len(index.sequences) == 2
+
+        seq_a = index.sequences[SGLangSeqId(rid="A")]
+        assert seq_a.input_ids == [10, 20, 30, 31]
+        assert seq_a.positions == [0, 1, 2, 3]
+        assert seq_a.locator.steps == [0, 0, 0, 1]
+
+        seq_b = index.sequences[SGLangSeqId(rid="B")]
+        assert seq_b.input_ids == [40, 50, 51]
+        assert seq_b.positions == [0, 1, 2]
+
+    def test_sequence_exit_and_join(self):
+        """Sequence A exits, new sequence D joins with different seq_id."""
+        side_aux = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30],
+                    positions=[0, 1, 2],
+                    seq_lens=[3],
+                    seq_ids=[SGLangSeqId(rid="A")],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[100, 200],
+                    positions=[0, 1],
+                    seq_lens=[2],
+                    seq_ids=[SGLangSeqId(rid="D")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+
+        index = build_seqs_info(side_aux)
+        assert len(index.sequences) == 2
+
+    def test_different_seq_ids_produce_separate_sequences(self):
+        """Different seq_ids at different steps → separate sequences."""
+        side_aux = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20],
+                    positions=[0, 1],
+                    seq_lens=[2],
+                    seq_ids=[SGLangSeqId(rid="A")],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[100, 200, 300],
+                    positions=[0, 1, 2],
+                    seq_lens=[3],
+                    seq_ids=[SGLangSeqId(rid="D")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+
+        index = build_seqs_info(side_aux)
+        assert len(index.sequences) == 2
+
+        all_input_ids = {
+            seq_id: rec.input_ids for seq_id, rec in index.sequences.items()
+        }
+        assert [10, 20] in all_input_ids.values()
+        assert [100, 200, 300] in all_input_ids.values()
+
 
 class TestBuildTokenIndexMegatronThd:
     """Tests for Megatron thd token index building."""
@@ -83,6 +172,7 @@ class TestBuildTokenIndexMegatronThd:
         seq0 = index.sequences[PositionalSeqId(step=0, seq_index=0)]
         assert seq0.input_ids == [10, 20, 30]
         assert seq0.positions == [0, 1, 2]
+        assert seq0.locator.steps == [0, 0, 0]
         assert seq0.locator.token_index_in_step == [0, 1, 2]
 
         seq1 = index.sequences[PositionalSeqId(step=0, seq_index=1)]
@@ -90,6 +180,44 @@ class TestBuildTokenIndexMegatronThd:
         assert seq1.positions == [0, 1]
         assert seq1.locator.token_index_in_step == [3, 4]
 
+    def test_multi_step_accumulation(self):
+        """Two steps with different seq_ids produce separate sequences."""
+        side_aux = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30, 40],
+                    positions=[0, 1, 0, 1],
+                    seq_lens=[2, 2],
+                    seq_ids=[
+                        PositionalSeqId(step=0, seq_index=0),
+                        PositionalSeqId(step=0, seq_index=1),
+                    ],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[50, 60, 70, 80],
+                    positions=[0, 1, 0, 1],
+                    seq_lens=[2, 2],
+                    seq_ids=[
+                        PositionalSeqId(step=1, seq_index=0),
+                        PositionalSeqId(step=1, seq_index=1),
+                    ],
+                ),
+            },
+            framework="megatron",
+            layout="thd",
+        )
+
+        index = build_seqs_info(side_aux)
+        assert len(index.sequences) == 4
+
+        seq0 = index.sequences[PositionalSeqId(step=0, seq_index=0)]
+        assert seq0.input_ids == [10, 20]
+        assert seq0.locator.steps == [0, 0]
+
+        seq2 = index.sequences[PositionalSeqId(step=1, seq_index=0)]
+        assert seq2.input_ids == [50, 60]
+        assert seq2.locator.steps == [1, 1]
+
 
 class TestMatchSequences:
     """Tests for _match_sequences: for each y, find matching x."""
@@ -240,6 +368,92 @@ class TestMatchSequences:
         assert matched == []
 
 
+class TestComputeAlignmentPlanCrossLayout:
+    """Tests for alignment plan across different step distributions."""
+
+    def test_thd_vs_thd_different_step_splits(self):
+        """Two thd sides with same tokens but different step distributions."""
+        side_aux_a = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20],
+                    positions=[0, 1],
+                    seq_lens=[2],
+                    seq_ids=[SGLangSeqId(rid="X")],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[30],
+                    positions=[2],
+                    seq_lens=[1],
+                    seq_ids=[SGLangSeqId(rid="X")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+        side_aux_b = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30],
+                    positions=[0, 1, 2],
+                    seq_lens=[3],
+                    seq_ids=[SGLangSeqId(rid="X")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+
+        index_a = build_seqs_info(side_aux_a)
+        index_b = build_seqs_info(side_aux_b)
+
+        plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b))
+        assert len(plan.locators.x.steps) == 3
+
+    def test_sglang_vs_megatron_thd(self):
+        """SGLang multi-step thd aligned with Megatron single-step thd."""
+        side_aux_a = 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")],
+                ),
+                1: TokenAlignerStepAux(
+                    input_ids=[31, 51],
+                    positions=[3, 2],
+                    seq_lens=[1, 1],
+                    seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
+                ),
+            },
+            framework="sglang",
+            layout="thd",
+        )
+        side_aux_b = TokenAlignerGlobalAux(
+            step_auxs={
+                0: TokenAlignerStepAux(
+                    input_ids=[10, 20, 30, 31, 40, 50, 51],
+                    positions=[0, 1, 2, 3, 0, 1, 2],
+                    seq_lens=[4, 3],
+                    seq_ids=[
+                        PositionalSeqId(step=0, seq_index=0),
+                        PositionalSeqId(step=0, seq_index=1),
+                    ],
+                ),
+            },
+            framework="megatron",
+            layout="thd",
+        )
+
+        index_a = build_seqs_info(side_aux_a)
+        index_b = build_seqs_info(side_aux_b)
+
+        plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b))
+
+        assert len(plan.locators.x.steps) == 7
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -263,6 +477,7 @@ def _make_index(
             input_ids=list(input_ids),
             positions=list(range(num_tokens)),
             locator=TokenLocator(
+                steps=[0] * num_tokens,
                 token_index_in_step=list(range(num_tokens)),
             ),
         )
@@ -280,6 +495,7 @@ def _make_seq_info_dict(
             input_ids=list(input_ids),
             positions=list(range(num_tokens)),
             locator=TokenLocator(
+                steps=[0] * num_tokens,
                 token_index_in_step=list(range(num_tokens)),
             ),
         )

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -11,8 +11,10 @@ from sglang.srt.debug_utils.comparator.output_types import (
     AnyRecord,
     ComparisonRecord,
     ConfigRecord,
+    GeneralWarning,
     SkipRecord,
     SummaryRecord,
+    WarningRecord,
     _OutputRecord,
     parse_record_json,
 )
@@ -1012,6 +1014,221 @@ class TestEntrypointReplicatedAxis:
         assert summary.passed == 0
 
 
+class TestEntrypointAlignment:
+    """Test `--grouping logical` with token alignment (aux tensors present)."""
+
+    def test_sglang_multi_step_alignment(self, tmp_path, capsys):
+        """SGLang multi-step dumps with aux tensors auto-trigger alignment."""
+        torch.manual_seed(42)
+        hidden_dim = 8
+
+        hidden_step0 = torch.randn(5, hidden_dim)
+        hidden_step1 = torch.randn(2, hidden_dim)
+
+        exp_paths: list[Path] = []
+        for side_dir in ["baseline", "target"]:
+            d = tmp_path / side_dir
+            d.mkdir()
+
+            dumper = _Dumper(
+                config=DumperConfig(
+                    enable=True,
+                    dir=str(d),
+                    exp_name=_FIXED_EXP_NAME,
+                    enable_http_server=False,
+                )
+            )
+
+            # Step 0: prefill with 2 sequences (3+2 tokens)
+            dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50]))
+            dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1]))
+            dumper.dump("seq_lens", torch.tensor([3, 2]))
+            dumper.dump("req_pool_indices", torch.tensor([7, 3]))
+            dumper.dump("rids", ["A", "B"])
+            dumper.dump("hidden_states", hidden_step0)
+            dumper.step()
+
+            # Step 1: decode (1 token per sequence)
+            dumper.dump("input_ids", torch.tensor([31, 51]))
+            dumper.dump("positions", torch.tensor([3, 2]))
+            dumper.dump("seq_lens", torch.tensor([1, 1]))
+            dumper.dump("req_pool_indices", torch.tensor([7, 3]))
+            dumper.dump("rids", ["A", "B"])
+            dumper.dump("hidden_states", hidden_step1)
+            dumper.step()
+
+            exp_paths.append(d / _FIXED_EXP_NAME)
+
+        args = _make_args(exp_paths[0], exp_paths[1], grouping="logical")
+        records = _run_and_parse(args, capsys)
+
+        comparisons = _get_comparisons(records)
+        # AUX_NAMES are filtered out after plan computation → only hidden_states remains
+        assert len(comparisons) == 1
+        assert comparisons[0].name == "hidden_states"
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.passed == 1
+        assert summary.failed == 0
+        assert summary.skipped == 0
+
+    def test_sglang_vs_megatron_cross_framework(self, tmp_path, capsys):
+        """SGLang 4-step thd baseline vs Megatron 1-step thd target align correctly."""
+        torch.manual_seed(42)
+        hidden_dim: int = 8
+
+        all_hiddens: torch.Tensor = torch.randn(11, hidden_dim)
+        seq_a_hiddens: torch.Tensor = all_hiddens[:6]
+        seq_b_hiddens: torch.Tensor = all_hiddens[6:]
+
+        # --- SGLang baseline: 1 prefill + 3 decode ---
+        sglang_dir: Path = tmp_path / "baseline"
+        sglang_dir.mkdir()
+        sglang_dumper = _Dumper(
+            config=DumperConfig(
+                enable=True,
+                dir=str(sglang_dir),
+                exp_name=_FIXED_EXP_NAME,
+                enable_http_server=False,
+            )
+        )
+
+        # Step 0: prefill — seq A (3 tokens) + seq B (2 tokens)
+        sglang_dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50]))
+        sglang_dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1]))
+        sglang_dumper.dump("seq_lens", torch.tensor([3, 2]))
+        sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3]))
+        sglang_dumper.dump("rids", ["A", "B"])
+        sglang_dumper.dump(
+            "hidden_states",
+            torch.stack(
+                [
+                    seq_a_hiddens[0],
+                    seq_a_hiddens[1],
+                    seq_a_hiddens[2],
+                    seq_b_hiddens[0],
+                    seq_b_hiddens[1],
+                ]
+            ),
+        )
+        sglang_dumper.step()
+
+        # Steps 1-3: decode — 1 token per sequence
+        decode_data: list[dict[str, object]] = [
+            {
+                "input_ids": torch.tensor([31, 51]),
+                "positions": torch.tensor([3, 2]),
+                "hidden": torch.stack([seq_a_hiddens[3], seq_b_hiddens[2]]),
+            },
+            {
+                "input_ids": torch.tensor([32, 52]),
+                "positions": torch.tensor([4, 3]),
+                "hidden": torch.stack([seq_a_hiddens[4], seq_b_hiddens[3]]),
+            },
+            {
+                "input_ids": torch.tensor([33, 53]),
+                "positions": torch.tensor([5, 4]),
+                "hidden": torch.stack([seq_a_hiddens[5], seq_b_hiddens[4]]),
+            },
+        ]
+        for step_data in decode_data:
+            sglang_dumper.dump("input_ids", step_data["input_ids"])
+            sglang_dumper.dump("positions", step_data["positions"])
+            sglang_dumper.dump("seq_lens", torch.tensor([1, 1]))
+            sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3]))
+            sglang_dumper.dump("rids", ["A", "B"])
+            sglang_dumper.dump("hidden_states", step_data["hidden"])
+            sglang_dumper.step()
+
+        # --- Megatron target: 1 step, thd [T, H] ---
+        megatron_dir: Path = tmp_path / "target"
+        megatron_dir.mkdir()
+        megatron_dumper = _Dumper(
+            config=DumperConfig(
+                enable=True,
+                dir=str(megatron_dir),
+                exp_name=_FIXED_EXP_NAME,
+                enable_http_server=False,
+            )
+        )
+
+        # THD flat: seq A (6 tokens) + seq B (5 tokens) = 11 tokens total
+        megatron_input_ids: torch.Tensor = torch.tensor(
+            [10, 20, 30, 31, 32, 33, 40, 50, 51, 52, 53]
+        )
+        megatron_cu_seqlens: torch.Tensor = torch.tensor([0, 6, 11])
+
+        megatron_hidden: torch.Tensor = torch.cat([seq_a_hiddens, seq_b_hiddens], dim=0)
+
+        megatron_dumper.dump("input_ids", megatron_input_ids)
+        megatron_dumper.dump("cu_seqlens_q", megatron_cu_seqlens)
+        megatron_dumper.dump("hidden_states", megatron_hidden)
+        megatron_dumper.step()
+
+        # --- Run comparison ---
+        args = _make_args(
+            sglang_dir / _FIXED_EXP_NAME,
+            megatron_dir / _FIXED_EXP_NAME,
+            grouping="logical",
+        )
+
+        records = _run_and_parse(args, capsys)
+
+        warning_records = [r for r in records if isinstance(r, WarningRecord)]
+        layout_warnings = [
+            w
+            for wr in warning_records
+            for w in wr.warnings
+            if isinstance(w, GeneralWarning)
+            and w.category == "layout_detection_fallback"
+        ]
+        assert len(layout_warnings) == 1
+
+        comparisons = _get_comparisons(records)
+        # AUX_NAMES filtered out → only hidden_states remains
+        assert len(comparisons) == 1
+        assert comparisons[0].name == "hidden_states"
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.passed == 1
+        assert summary.failed == 0
+        assert summary.skipped == 0
+
+    def test_alignment_fallback_when_no_aux(self, tmp_path, capsys):
+        """Without aux tensors, logical grouping skips alignment and compares per-step."""
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=2)
+        args = _make_args(
+            baseline_path, target_path, grouping="logical", diff_threshold=0.1
+        )
+
+        capsys.readouterr()
+        run(args)
+        captured = capsys.readouterr()
+        records = _parse_jsonl(captured.out)
+        warning_records = [r for r in records if isinstance(r, WarningRecord)]
+        aux_missing_warnings = [
+            w
+            for wr in warning_records
+            for w in wr.warnings
+            if isinstance(w, GeneralWarning) and w.category == "aux_tensors_missing"
+        ]
+        assert len(aux_missing_warnings) == 1
+
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 2
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.total == 2
+        assert summary.passed == 2
+
+
 # --------------------------- Assertion helpers -------------------
 
 

test/registered/debug_utils/comparator/test_model_validation.py

@@ -92,7 +92,7 @@ class TestTokenAlignerSeqInfo:
         info = TokenAlignerSeqInfo(
             input_ids=[10, 20, 30],
             positions=[0, 1, 2],
-            locator=TokenLocator(token_index_in_step=[0, 1, 0]),
+            locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
         )
         assert len(info.input_ids) == 3
 
@@ -101,7 +101,7 @@ class TestTokenAlignerSeqInfo:
             TokenAlignerSeqInfo(
                 input_ids=[10, 20, 30],
                 positions=[0, 1, 2],
-                locator=TokenLocator(token_index_in_step=[0, 1]),
+                locator=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1, 0]),
             )
 
     def test_positions_not_sequential(self):
@@ -109,7 +109,7 @@ class TestTokenAlignerSeqInfo:
             TokenAlignerSeqInfo(
                 input_ids=[10, 20, 30],
                 positions=[0, 2, 1],
-                locator=TokenLocator(token_index_in_step=[0, 1, 0]),
+                locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
             )
 
 
@@ -117,18 +117,18 @@ class TestTokenAlignerPlan:
     def test_valid(self):
         plan = TokenAlignerPlan(
             locators=Pair(
-                x=TokenLocator(token_index_in_step=[0, 1, 0]),
+                x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
-                y=TokenLocator(token_index_in_step=[0, 0, 1]),
+                y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
             ),
         )
-        assert len(plan.locators.x.token_index_in_step) == 3
+        assert len(plan.locators.x.steps) == 3
 
     def test_length_mismatch(self):
         with pytest.raises(ValidationError, match="Length mismatch"):
             TokenAlignerPlan(
                 locators=Pair(
-                    x=TokenLocator(token_index_in_step=[0, 1]),
+                    x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]),
-                    y=TokenLocator(token_index_in_step=[0, 0, 1]),
+                    y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
                 ),
             )
 

test/registered/debug_utils/test_dumper.py

@@ -2127,7 +2127,9 @@ class TestRegisterForwardHook:
 class TestPluginCoreFields:
     def test_sglang_core_fields(self):
         plugin = _SGLangPlugin()
-        assert plugin.core_fields() == frozenset({"input_ids", "positions", "seq_lens"})
+        assert plugin.core_fields() == frozenset(
+            {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
+        )
 
     def test_megatron_core_fields(self):
         plugin = _MegatronPlugin()

test/run_suite.py

@@ -192,7 +192,7 @@ def run_a_suite(args):
     files = [
         f
         for f in glob.glob("registered/**/*.py", recursive=True)
-        if not f.endswith("/conftest.py")
+        if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
     ]
     # Strict: all registered files must have proper registration
     sanity_check = True