Support concat mode in token aligner in dump comparator (#19599)

2026-03-01 10:35:50 +08:00
parent e78f1283f7
commit b0b26a7ef1
26 changed files with 1017 additions and 63 deletions
@@ -14,7 +14,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
    AlignerPerStepPlan,
    AlignerPlan,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
    TokenAlignerPlan,
    TokenLocator,
 )
@@ -234,6 +234,7 @@ class TestExecuteAlignerPlanWithTokenDim:
                x=[self._make_step_plan(step=0, indices=[0])],
                y=[self._make_step_plan(step=0, indices=[0])],
            ),
+            token_aligner_mode="smart",
            token_aligner_plan=token_plan,
        )

@@ -285,6 +286,7 @@ class TestExecuteAlignerPlanWithTokenDim:
                x=[self._make_step_plan(step=0, indices=[0])],
                y=[self._make_step_plan(step=0, indices=[0])],
            ),
+            token_aligner_mode="smart",
            token_aligner_plan=token_plan,
        )

@@ -143,6 +143,7 @@ class TestComputeAlignerPlan:

        plan: AlignerPlan = compute_aligner_plan(
            metas_pair=Pair(x=metas_x, y=metas_y),
+            token_aligner_mode=None,
            token_aligner_plan=None,
        )

@@ -151,7 +152,7 @@ class TestComputeAlignerPlan:
        assert plan.token_aligner_plan is None

    def test_preserves_token_aligner_plan(self) -> None:
-        from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+        from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
            TokenAlignerPlan,
            TokenLocator,
        )
@@ -166,10 +167,12 @@ class TestComputeAlignerPlan:

        plan: AlignerPlan = compute_aligner_plan(
            metas_pair=Pair(x=[_make_meta()], y=[_make_meta()]),
+            token_aligner_mode="smart",
            token_aligner_plan=ta_plan,
        )

        assert plan.token_aligner_plan is ta_plan
+        assert plan.token_aligner_mode == "smart"


 class TestComputePerStepSubPlansThd:
@@ -5,13 +5,13 @@ import polars as pl
 import pytest
 import torch

-from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader import (
    _detect_plugin,
    _ensure_dims_in_metas,
    _load_and_align_aux_tensor,
    _load_non_tensor_aux,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_plugins import (
    _MegatronPlugin,
    _SGLangPlugin,
 )
@@ -214,7 +214,7 @@ class TestLoadNonTensorAux:
            from unittest.mock import patch

            with patch(
-                "sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
+                "sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader.warning_sink",
                sink,
            ):
                result = _load_non_tensor_aux(
@@ -273,7 +273,7 @@ class TestLoadAndAlignAuxTensor:
            from unittest.mock import patch

            with patch(
-                "sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
+                "sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader.warning_sink",
                sink,
            ):
                result = _load_and_align_aux_tensor(
@@ -329,7 +329,7 @@ class TestLoadNonTensorAuxDp:
            from unittest.mock import patch

            with patch(
-                "sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
+                "sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_loader.warning_sink",
                sink,
            ):
                result = _load_non_tensor_aux(
@@ -3,12 +3,12 @@ import sys
 import pytest
 import torch

-from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_plugins import (
    _infer_positions,
    _MegatronPlugin,
    _SGLangPlugin,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
    PositionalSeqId,
    SGLangSeqId,
    TokenAlignerStepAux,
@@ -0,0 +1,84 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps import (
+    execute_token_aligner_concat_steps,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+class TestExecuteConcat:
+    def test_single_step_equal_length(self) -> None:
+        x = torch.tensor([1.0, 2.0, 3.0])
+        y = torch.tensor([4.0, 5.0, 6.0])
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(x={0: x}, y={0: y}),
+        )
+        assert torch.equal(result.x, x)
+        assert torch.equal(result.y, y)
+
+    def test_truncates_to_min(self) -> None:
+        x = torch.tensor([1.0, 2.0, 3.0, 4.0])
+        y = torch.tensor([5.0, 6.0])
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(x={0: x}, y={0: y}),
+        )
+        assert torch.equal(result.x, torch.tensor([1.0, 2.0]))
+        assert torch.equal(result.y, y)
+
+    def test_multi_step_sorted_concat(self) -> None:
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(
+                x={1: torch.tensor([3.0, 4.0]), 0: torch.tensor([1.0, 2.0])},
+                y={0: torch.tensor([5.0, 6.0, 7.0, 8.0])},
+            ),
+        )
+        assert torch.equal(result.x, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert torch.equal(result.y, torch.tensor([5.0, 6.0, 7.0, 8.0]))
+
+    def test_named_token_dim_nonzero(self) -> None:
+        """Token dim at dim=1 (not dim=0) — concat and truncate along correct dim."""
+        # shape [2, 3, 4]: dim0=batch, dim1=token, dim2=hidden
+        x_step0 = torch.randn(2, 3, 4).refine_names("b", "t", "h")
+        x_step1 = torch.randn(2, 5, 4).refine_names("b", "t", "h")
+        y_step0 = torch.randn(2, 6, 4).refine_names("b", "t", "h")
+
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(
+                x={0: x_step0, 1: x_step1},
+                y={0: y_step0},
+            ),
+        )
+
+        # x: 3+5=8 tokens; y: 6 tokens → truncate to 6
+        assert result.x.shape == (2, 6, 4)
+        assert result.y.shape == (2, 6, 4)
+
+    def test_named_dims_no_token_dim_fallback(self) -> None:
+        """Named dims without t or s → fallback to dim 0."""
+        x = torch.randn(4, 8).refine_names("b", "h")
+        y = torch.randn(3, 8).refine_names("b", "h")
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(x={0: x}, y={0: y}),
+        )
+        assert result.x.shape == (3, 8)
+        assert result.y.shape == (3, 8)
+
+    def test_seq_dim_fallback(self) -> None:
+        """Named dims with s but no t → uses s as token dim."""
+        x = torch.randn(2, 5, 4).refine_names("b", "s", "h")
+        y = torch.randn(2, 3, 4).refine_names("b", "s", "h")
+        result: Pair[torch.Tensor] = execute_token_aligner_concat_steps(
+            tensor_of_step_pair=Pair(x={0: x}, y={0: y}),
+        )
+        assert result.x.shape == (2, 3, 4)
+        assert result.y.shape == (2, 3, 4)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -5,16 +5,16 @@ import sys
 import pytest
 import torch

-from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.executor import (
    execute_token_aligner,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.planner import (
    compute_token_aligner_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.seq_info_builder import (
    build_seqs_info,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
    SGLangSeqId,
    TokenAlignerGlobalAux,
    TokenAlignerPlan,
@@ -2,14 +2,14 @@ import sys

 import pytest

-from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.planner import (
    _match_sequences,
    compute_token_aligner_plan,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.seq_info_builder import (
    build_seqs_info,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
    PositionalSeqId,
    SeqId,
    SGLangSeqId,
@@ -0,0 +1,172 @@
+import sys
+from pathlib import Path
+from unittest.mock import patch
+
+import polars as pl
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps.thd_seq_lens_loader import (
+    load_thd_seq_lens_only,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.aux_plugins import (
+    _SGLangPlugin,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+def _save_pt(
+    dump_path: Path,
+    *,
+    name: str,
+    step: int,
+    rank: int,
+    value: object,
+    meta: dict | None = None,
+) -> str:
+    filename: str = f"name={name}___step={step}___rank={rank}.pt"
+    payload: dict = {"value": value, "meta": meta or {}}
+    torch.save(payload, dump_path / filename)
+    return filename
+
+
+def _make_df_from_filenames(filenames: list[str]) -> pl.DataFrame:
+    rows: list[dict] = []
+    for fn in filenames:
+        parts: dict = {}
+        stem: str = fn.removesuffix(".pt")
+        for kv in stem.split("___"):
+            if "=" in kv:
+                k, v = kv.split("=", 1)
+                parts[k] = v
+        rows.append(
+            {
+                "filename": fn,
+                "name": parts["name"],
+                "step": int(parts["step"]),
+                "rank": int(parts["rank"]),
+            }
+        )
+    return pl.DataFrame(rows)
+
+
+class TestLoadThdSeqLensOnly:
+    """Tests for load_thd_seq_lens_only."""
+
+    def test_returns_none_when_no_plugin(self, tmp_path: Path) -> None:
+        """No recognized plugin → returns None."""
+        fn: str = _save_pt(
+            tmp_path, name="unrelated_tensor", step=0, rank=0, value=torch.tensor([1])
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is None
+
+    def test_returns_none_when_no_cp_sharded_names(self, tmp_path: Path) -> None:
+        """Plugin detected but cp_sharded_names is empty → returns None."""
+
+        class _NoCpPlugin(_SGLangPlugin):
+            @property
+            def cp_sharded_names(self) -> frozenset[str]:
+                return frozenset()
+
+        fn: str = _save_pt(
+            tmp_path,
+            name="seq_lens",
+            step=0,
+            rank=0,
+            value=torch.tensor([3, 5]),
+            meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        with patch(
+            "sglang.srt.debug_utils.comparator.aligner.token_aligner.concat_steps.thd_seq_lens_loader._detect_plugin",
+            return_value=_NoCpPlugin(),
+        ):
+            result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is None
+
+    def test_sglang_extracts_seq_lens(self, tmp_path: Path) -> None:
+        """SGLang format: seq_lens tensor present → extracts per-seq lengths."""
+        fn: str = _save_pt(
+            tmp_path,
+            name="seq_lens",
+            step=0,
+            rank=0,
+            value=torch.tensor([3, 5]),
+            meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is not None
+        assert result == {0: [3, 5]}
+
+    def test_megatron_extracts_from_cu_seqlens(self, tmp_path: Path) -> None:
+        """Megatron format: cu_seqlens_q tensor → derives seq_lens via diff."""
+        fn: str = _save_pt(
+            tmp_path,
+            name="cu_seqlens_q",
+            step=0,
+            rank=0,
+            value=torch.tensor([0, 3, 8], dtype=torch.int64),
+            meta={"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is not None
+        assert result == {0: [3, 5]}
+
+    def test_multi_step(self, tmp_path: Path) -> None:
+        """Two steps with different seq_lens → returns both in result dict."""
+        fn0: str = _save_pt(
+            tmp_path,
+            name="seq_lens",
+            step=0,
+            rank=0,
+            value=torch.tensor([3, 5]),
+            meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
+        )
+        fn1: str = _save_pt(
+            tmp_path,
+            name="seq_lens",
+            step=1,
+            rank=0,
+            value=torch.tensor([10, 20, 30]),
+            meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
+
+        result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is not None
+        assert result == {0: [3, 5], 1: [10, 20, 30]}
+
+    def test_returns_none_when_seq_lens_missing(self, tmp_path: Path) -> None:
+        """Plugin with cp_sharded_names but no seq_lens/cu_seqlens_q tensor → None."""
+        fn: str = _save_pt(
+            tmp_path,
+            name="cu_seqlens_kv",
+            step=0,
+            rank=0,
+            value=torch.tensor([0, 4], dtype=torch.int64),
+            meta={"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = load_thd_seq_lens_only(dump_path=tmp_path, df=df)
+
+        assert result is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -544,7 +544,7 @@ class TestEntrypointGroupingLogical:
        assert summary.skipped == 0

    def test_multi_step_tp(self, tmp_path, capsys):
-        """Two steps with TP=2 shards produce two per-step comparisons (no aux → no alignment)."""
+        """Two steps with TP=2 shards: concat mode merges into one comparison."""
        torch.manual_seed(42)
        full_tensor = torch.randn(4, 8)

@@ -574,14 +574,14 @@ class TestEntrypointGroupingLogical:

        records = _run_and_parse(args, capsys)
        comparisons = _get_comparisons(records)
-        assert len(comparisons) == 2
-        assert comparisons[0].baseline.shape == [4, 8]
-        assert comparisons[1].baseline.shape == [4, 8]
+        assert len(comparisons) == 1
+        # concat along dim 0 (fallback, no token dim) → 2 steps × [4, 8] = [8, 8]
+        assert comparisons[0].baseline.shape == [8, 8]

        summary = records[-1]
        assert isinstance(summary, SummaryRecord)
-        assert summary.total == 2
-        assert summary.passed == 2
+        assert summary.total == 1
+        assert summary.passed == 1

    def test_cp_axis_unshard(self, tmp_path, capsys):
        """CP-sharded tensors are correctly concatenated along the sequence dim."""
@@ -1046,6 +1046,447 @@ class TestEntrypointGroupingLogical:
        assert comp.name == "hidden"


+class TestEntrypointConcatMode:
+    """Test concat token-aligner mode through the full entrypoint pipeline."""
+
+    @staticmethod
+    def _make_dirs(tmp_path: Path) -> tuple[Path, Path]:
+        baseline_dir: Path = tmp_path / "baseline"
+        target_dir: Path = tmp_path / "target"
+        baseline_dir.mkdir()
+        target_dir.mkdir()
+        return baseline_dir, target_dir
+
+    @staticmethod
+    def _create_both_sides(
+        tmp_path: Path,
+        *,
+        baseline_steps: list[torch.Tensor],
+        target_steps: list[torch.Tensor],
+        name: str = "hidden",
+        dims: str | None = None,
+    ) -> tuple[Path, Path]:
+        """Create multi-step rank-0 dumps for both sides and return exp paths."""
+        baseline_dir, target_dir = TestEntrypointConcatMode._make_dirs(tmp_path)
+
+        for side_dir, steps in [
+            (baseline_dir, baseline_steps),
+            (target_dir, target_steps),
+        ]:
+            _create_multi_step_rank_dump(
+                side_dir,
+                rank=0,
+                name=name,
+                tensors_per_step=steps,
+                dims=dims,
+            )
+
+        return baseline_dir / _FIXED_EXP_NAME, target_dir / _FIXED_EXP_NAME
+
+    @staticmethod
+    def _run_concat(
+        tmp_path: Path,
+        capsys: pytest.CaptureFixture,
+        *,
+        baseline_steps: list[torch.Tensor],
+        target_steps: list[torch.Tensor],
+        name: str = "hidden",
+        dims: str | None = None,
+        diff_threshold: float = 0.01,
+    ) -> list[AnyRecord]:
+        """Create both-side dumps, run comparator, return parsed records."""
+        baseline_path, target_path = TestEntrypointConcatMode._create_both_sides(
+            tmp_path,
+            baseline_steps=baseline_steps,
+            target_steps=target_steps,
+            name=name,
+            dims=dims,
+        )
+        args: Namespace = _make_args(
+            baseline_path, target_path, diff_threshold=diff_threshold
+        )
+        return _run_and_parse(args, capsys)
+
+    def test_concat_multi_step_different_data(self, tmp_path, capsys):
+        """Multi-step concat with different data per step + truncation."""
+        torch.manual_seed(42)
+
+        # baseline: 2 steps [5,4] + [3,4] → concat → [8,4]
+        baseline_step0 = torch.randn(5, 4)
+        baseline_step1 = torch.randn(3, 4)
+        baseline_concat = torch.cat([baseline_step0, baseline_step1], dim=0)
+
+        # target: 1 step [6,4] — will be truncated to min(8,6)=6
+        target_step0 = baseline_concat[:6] + torch.randn(6, 4) * 0.0001
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[baseline_step0, baseline_step1],
+            target_steps=[target_step0],
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        # truncated to min(8,6) = 6 along concat dim
+        assert comparisons[0].baseline.shape == [6, 4]
+        assert comparisons[0].target.shape == [6, 4]
+
+    def test_concat_multi_step_tp_unshard(self, tmp_path, capsys):
+        """Multi-step different data + TP=2 unshard + concat."""
+        torch.manual_seed(42)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        # 2 steps: [4,8] each → concat → [8,8]
+        full_step0 = torch.randn(4, 8)
+        full_step1 = torch.randn(4, 8)
+
+        _create_multi_step_tp_sharded_dumps(
+            baseline_dir,
+            full_tensors_per_step=[full_step0, full_step1],
+            name="hidden",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="b h(tp)",
+        )
+        _create_multi_step_tp_sharded_dumps(
+            target_dir,
+            full_tensors_per_step=[
+                full_step0 + torch.randn(4, 8) * 0.0001,
+                full_step1 + torch.randn(4, 8) * 0.0001,
+            ],
+            name="hidden",
+            tp_size=2,
+            shard_dim=1,
+            dims_str="b h(tp)",
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        # 2 steps × [4, 8] concat along dim 0 (fallback) → [8, 8]
+        assert comparisons[0].baseline.shape == [8, 8]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_unequal_step_counts(self, tmp_path, capsys):
+        """Baseline 3 steps vs target 2 steps with truncation."""
+        torch.manual_seed(42)
+
+        # baseline: 3 steps [3]+[4]+[2] = 9 tokens along dim 0
+        b_step0 = torch.randn(3, 4)
+        b_step1 = torch.randn(4, 4)
+        b_step2 = torch.randn(2, 4)
+        b_concat = torch.cat([b_step0, b_step1, b_step2], dim=0)
+
+        # target: 2 steps [5]+[3] = 8 tokens along dim 0
+        t_step0 = b_concat[:5] + torch.randn(5, 4) * 0.0001
+        t_step1 = b_concat[5:8] + torch.randn(3, 4) * 0.0001
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1, b_step2],
+            target_steps=[t_step0, t_step1],
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        # truncated to min(9,8) = 8
+        assert comparisons[0].baseline.shape == [8, 4]
+        assert comparisons[0].target.shape == [8, 4]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_token_dim_nonzero(self, tmp_path, capsys):
+        """Token dim at dim=1 (dims='b t h') — concat along dim 1."""
+        torch.manual_seed(42)
+
+        # 2 steps: [2,5,4] + [2,3,4] → concat along dim 1 → [2,8,4]
+        b_step0 = torch.randn(2, 5, 4)
+        b_step1 = torch.randn(2, 3, 4)
+        b_concat = torch.cat([b_step0, b_step1], dim=1)
+
+        t_step0 = b_concat[:, :5, :] + torch.randn(2, 5, 4) * 0.0001
+        t_step1 = b_concat[:, 5:, :] + torch.randn(2, 3, 4) * 0.0001
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1],
+            target_steps=[t_step0, t_step1],
+            dims="b t h",
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        assert comparisons[0].baseline.shape == [2, 8, 4]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_seq_dim_fallback(self, tmp_path, capsys):
+        """No 't' dim but 's' dim present (dims='b s h') → concat along s."""
+        torch.manual_seed(42)
+
+        # 2 steps: [2,5,4] + [2,3,4] → concat along dim 1 (s) → [2,8,4]
+        b_step0 = torch.randn(2, 5, 4)
+        b_step1 = torch.randn(2, 3, 4)
+        b_concat = torch.cat([b_step0, b_step1], dim=1)
+
+        t_step0 = b_concat[:, :5, :] + torch.randn(2, 5, 4) * 0.0001
+        t_step1 = b_concat[:, 5:, :] + torch.randn(2, 3, 4) * 0.0001
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1],
+            target_steps=[t_step0, t_step1],
+            dims="b s h",
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        assert comparisons[0].baseline.shape == [2, 8, 4]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_no_dims_fallback(self, tmp_path, capsys):
+        """No dims annotation → fallback to concat along dim 0."""
+        torch.manual_seed(42)
+
+        # 2 steps: [5,4] + [3,4] → concat along dim 0 → [8,4]
+        b_step0 = torch.randn(5, 4)
+        b_step1 = torch.randn(3, 4)
+        b_concat = torch.cat([b_step0, b_step1], dim=0)
+
+        t_step0 = b_concat[:5] + torch.randn(5, 4) * 0.0001
+        t_step1 = b_concat[5:] + torch.randn(3, 4) * 0.0001
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1],
+            target_steps=[t_step0, t_step1],
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        assert comparisons[0].baseline.shape == [8, 4]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_preserves_step_order(self, tmp_path, capsys):
+        """Verify step0 data precedes step1 data in the concatenated result."""
+        # deterministic integer data: step0=[1,2,3], step1=[4,5]
+        b_step0 = torch.tensor([[1.0], [2.0], [3.0]])
+        b_step1 = torch.tensor([[4.0], [5.0]])
+
+        # target: same data, single step [1,2,3,4,5]
+        t_full = torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0]])
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1],
+            target_steps=[t_full],
+        )
+        comp = _assert_single_comparison_passed(records)
+        # if order were wrong, diff would not pass with exact integer data
+        assert comp.baseline.shape == [5, 1]
+        assert comp.diff is not None
+        assert comp.diff.max_abs_diff == 0.0
+
+    def test_concat_aux_tensors_not_filtered(self, tmp_path, capsys):
+        """Concat mode does not filter aux tensors — all participate in comparison."""
+        torch.manual_seed(42)
+
+        baseline_dir, target_dir = self._make_dirs(tmp_path)
+
+        hidden = torch.randn(4, 8)
+        input_ids = torch.randint(0, 100, (4,))
+        positions = torch.arange(4)
+
+        _create_rank_dump(
+            baseline_dir,
+            rank=0,
+            name="hidden_states",
+            tensor=hidden,
+            extra_dumps=[("input_ids", input_ids), ("positions", positions)],
+        )
+        _create_rank_dump(
+            target_dir,
+            rank=0,
+            name="hidden_states",
+            tensor=hidden + torch.randn(4, 8) * 0.0001,
+            extra_dumps=[("input_ids", input_ids), ("positions", positions)],
+        )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comparisons = _get_comparisons(records)
+        # all 3 tensors should be compared (not filtered out)
+        names = {c.name for c in comparisons}
+        assert "hidden_states" in names
+        assert "input_ids" in names
+        assert "positions" in names
+        assert len(comparisons) == 3
+
+    def test_concat_aligner_plan_fields(self, tmp_path, capsys):
+        """ComparisonRecord.aligner_plan reports mode='concat' with plan=None."""
+        torch.manual_seed(42)
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[torch.randn(3, 4), torch.randn(2, 4)],
+            target_steps=[torch.randn(3, 4), torch.randn(2, 4)],
+            diff_threshold=100.0,
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        plan = comparisons[0].aligner_plan
+        assert plan is not None
+        assert plan.token_aligner_mode == "concat_steps"
+        assert plan.token_aligner_plan is None
+
+    def test_concat_comparison_fails(self, tmp_path, capsys):
+        """Completely different data → comparison fails."""
+        torch.manual_seed(42)
+        b_step0 = torch.randn(4, 4)
+        b_step1 = torch.randn(3, 4)
+
+        # target: completely different random data
+        torch.manual_seed(99)
+        t_step0 = torch.randn(4, 4) * 100
+        t_step1 = torch.randn(3, 4) * 100
+
+        records = self._run_concat(
+            tmp_path,
+            capsys,
+            baseline_steps=[b_step0, b_step1],
+            target_steps=[t_step0, t_step1],
+            diff_threshold=1e-6,
+        )
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        assert comparisons[0].diff is not None
+        assert not comparisons[0].diff.passed
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.failed == 1
+        assert summary.passed == 0
+
+    def test_concat_multi_step_cp_unshard(self, tmp_path, capsys):
+        """Multi-step different data + CP=2 unshard along seq dim + concat."""
+        torch.manual_seed(42)
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        # 2 steps: [4,8,6] each → concat along seq dim (dim 1) → [4,16,6]
+        full_step0 = torch.randn(4, 8, 6)
+        full_step1 = torch.randn(4, 8, 6)
+
+        for side_dir, steps in [
+            (baseline_dir, [full_step0, full_step1]),
+            (
+                target_dir,
+                [
+                    full_step0 + torch.randn(4, 8, 6) * 0.0001,
+                    full_step1 + torch.randn(4, 8, 6) * 0.0001,
+                ],
+            ),
+        ]:
+            for cp_rank in range(2):
+                per_step_shards: list[torch.Tensor] = [
+                    t.chunk(2, dim=1)[cp_rank] for t in steps
+                ]
+                _create_multi_step_rank_dump(
+                    side_dir,
+                    rank=cp_rank,
+                    name="attn_out",
+                    tensors_per_step=per_step_shards,
+                    dims="b s(cp) h",
+                    parallel_info={"cp_rank": cp_rank, "cp_size": 2},
+                )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comparisons = _get_comparisons(records)
+        assert len(comparisons) == 1
+        # CP unshard: [4,4,6] × 2 ranks → [4,8,6] per step
+        # concat along seq dim (dim 1): 2 steps × [4,8,6] → [4,16,6]
+        assert comparisons[0].baseline.shape == [4, 16, 6]
+        assert comparisons[0].diff is not None
+        assert comparisons[0].diff.passed
+
+    def test_concat_thd_cp_zigzag(self, tmp_path: Path, capsys) -> None:
+        """Concat mode with THD CP=2 zigzag (Megatron format) — unshard + reorder works."""
+        torch.manual_seed(42)
+        cp_size: int = 2
+        seq_lens: list[int] = [100, 64]
+        total_tokens: int = sum(seq_lens)
+        total_per_rank: int = 128
+        num_steps: int = 2
+
+        full_tensor: torch.Tensor = torch.randn(total_tokens + 92)
+
+        baseline_dir: Path = tmp_path / "baseline"
+        target_dir: Path = tmp_path / "target"
+        baseline_dir.mkdir()
+        target_dir.mkdir()
+
+        baseline_path: Path = _create_thd_cp_zigzag_dumps(
+            baseline_dir,
+            full_tensor=full_tensor,
+            name="hidden_states",
+            seq_lens=seq_lens,
+            cp_size=cp_size,
+            total_per_rank=total_per_rank,
+            num_steps=num_steps,
+        )
+
+        target_tensor: torch.Tensor = full_tensor + torch.randn_like(full_tensor) * 1e-5
+        target_path: Path = _create_thd_cp_zigzag_dumps(
+            target_dir,
+            full_tensor=target_tensor,
+            name="hidden_states",
+            seq_lens=seq_lens,
+            cp_size=cp_size,
+            total_per_rank=total_per_rank,
+            num_steps=num_steps,
+        )
+
+        args: Namespace = _make_args(
+            baseline_path,
+            target_path,
+            token_aligner="concat_steps",
+            diff_threshold=1e-3,
+        )
+        records: list[AnyRecord] = _run_and_parse(args, capsys)
+
+        comparisons: list[ComparisonRecord] = _get_comparisons(records)
+        hidden_comparisons: list[ComparisonRecord] = [
+            c for c in comparisons if c.name == "hidden_states"
+        ]
+        assert len(hidden_comparisons) >= 1
+        assert all(c.diff is not None and c.diff.passed for c in hidden_comparisons)
+
+
 class TestEntrypointAxisAligner:
    """Test cross-framework dim reordering through the full entrypoint pipeline."""

@@ -1406,7 +1847,9 @@ class TestEntrypointAlignment:

            exp_paths.append(d / _FIXED_EXP_NAME)

-        args = _make_args(exp_paths[0], exp_paths[1], grouping="logical")
+        args = _make_args(
+            exp_paths[0], exp_paths[1], grouping="logical", token_aligner="smart"
+        )
        records = _run_and_parse(args, capsys)

        comparisons = _get_comparisons(records)
@@ -1518,6 +1961,7 @@ class TestEntrypointAlignment:
            sglang_dir / _FIXED_EXP_NAME,
            megatron_dir / _FIXED_EXP_NAME,
            grouping="logical",
+            token_aligner="smart",
        )

        records = _run_and_parse(args, capsys)
@@ -1549,7 +1993,11 @@ class TestEntrypointAlignment:
        """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
+            baseline_path,
+            target_path,
+            grouping="logical",
+            token_aligner="smart",
+            diff_threshold=0.1,
        )

        capsys.readouterr()
@@ -1885,6 +2333,7 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
        filter=None,
        output_format="json",
        grouping="logical",
+        token_aligner="concat_steps",
        viz_bundle_details=False,
        viz_output_dir="/tmp/comparator_viz/",
        visualize_per_token=None,
@@ -1948,6 +2397,44 @@ def _create_rank_dump(
    return directory / _FIXED_EXP_NAME


+def _create_multi_step_rank_dump(
+    directory: Path,
+    *,
+    rank: int,
+    name: str,
+    tensors_per_step: list[torch.Tensor],
+    dims: str | None = None,
+    parallel_info: dict | None = None,
+    framework: str = "sglang",
+) -> Path:
+    """Create a dump file with *different* tensors per step.
+
+    Unlike ``_create_rank_dump`` (which repeats the same tensor),
+    this helper accepts a list of tensors — one per step.
+    """
+    with pytest.MonkeyPatch.context() as mp:
+        mp.setattr(_dumper_module, "_get_rank", lambda: rank)
+
+        dumper = _Dumper(
+            config=DumperConfig(
+                enable=True,
+                dir=str(directory),
+                exp_name=_FIXED_EXP_NAME,
+            )
+        )
+
+        static_meta: dict = {"world_rank": rank, "world_size": 1}
+        if parallel_info is not None:
+            static_meta[f"{framework}_parallel_info"] = parallel_info
+        dumper.__dict__["_static_meta"] = static_meta
+
+        for tensor in tensors_per_step:
+            dumper.dump(name, tensor, dims=dims)
+            dumper.step()
+
+    return directory / _FIXED_EXP_NAME
+
+
 def _create_cp_tp_sharded_dumps(
    directory: Path,
    *,
@@ -2153,6 +2640,38 @@ def _create_tp_sharded_dumps(
    return directory / _FIXED_EXP_NAME


+def _create_multi_step_tp_sharded_dumps(
+    directory: Path,
+    *,
+    full_tensors_per_step: list[torch.Tensor],
+    name: str,
+    tp_size: int,
+    shard_dim: int,
+    dims_str: str,
+) -> Path:
+    """Create TP-sharded dump files with *different* tensors per step.
+
+    Each step's full tensor is chunked across TP ranks, then
+    ``_create_multi_step_rank_dump`` writes one file per rank.
+    """
+    shards_per_rank: list[list[torch.Tensor]] = [[] for _ in range(tp_size)]
+    for full_tensor in full_tensors_per_step:
+        shards = list(full_tensor.chunk(tp_size, dim=shard_dim))
+        for tp_rank in range(tp_size):
+            shards_per_rank[tp_rank].append(shards[tp_rank])
+
+    for tp_rank in range(tp_size):
+        _create_multi_step_rank_dump(
+            directory,
+            rank=tp_rank,
+            name=name,
+            tensors_per_step=shards_per_rank[tp_rank],
+            dims=dims_str,
+            parallel_info={"tp_rank": tp_rank, "tp_size": tp_size},
+        )
+    return directory / _FIXED_EXP_NAME
+
+
 def _create_tp_partial_dumps(
    directory: Path,
    *,
@@ -2470,6 +2989,7 @@ class TestEntrypointThdCpZigzag:
            sglang_dir / _FIXED_EXP_NAME,
            megatron_dir / _FIXED_EXP_NAME,
            grouping="logical",
+            token_aligner="smart",
            diff_threshold=1e-3,
        )
        records: list[AnyRecord] = _run_and_parse(args, capsys)
@@ -2517,7 +3037,11 @@ class TestEntrypointThdCpZigzag:
        )

        args: Namespace = _make_args(
-            baseline_path, target_path, grouping="logical", diff_threshold=1e-3
+            baseline_path,
+            target_path,
+            grouping="logical",
+            token_aligner="smart",
+            diff_threshold=1e-3,
        )
        records: list[AnyRecord] = _run_and_parse(args, capsys)

@@ -8,7 +8,7 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
    AlignerPerStepPlan,
    AlignerPlan,
 )
-from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
    PositionalSeqId,
    TokenAlignerPlan,
    TokenAlignerSeqInfo,