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

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

@@ -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,
         )
 

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

@@ -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:

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

@@ -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(

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

@@ -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,

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

@@ -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__]))

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

@@ -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,

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

@@ -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,

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

@@ -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__]))