Support loading token aligner data in dump comparator (#19376)

2026-02-26 10:03:56 +08:00
parent e8dd14519d
commit d34d5aca07
10 changed files with 1182 additions and 43 deletions
@@ -0,0 +1,278 @@
+import sys
+from pathlib import Path
+
+import polars as pl
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.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 (
+    _MegatronPlugin,
+    _SGLangPlugin,
+)
+from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
+from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+_sglang_plugin = _SGLangPlugin()
+_megatron_plugin = _MegatronPlugin()
+
+
+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 TestEnsureDimsInMetas:
+    """Tests for _ensure_dims_in_metas."""
+
+    def _make_meta(self, *, cp_size: int = 1, cp_rank: int = 0) -> dict:
+        return {
+            "sglang_parallel_info": {
+                "tp_rank": 0,
+                "tp_size": 1,
+                "cp_rank": cp_rank,
+                "cp_size": cp_size,
+            }
+        }
+
+    def test_no_cp_returns_metas_unchanged(self):
+        """Without CP parallelism, metas are returned as-is."""
+        metas: list[dict] = [self._make_meta(cp_size=1)]
+        result = _ensure_dims_in_metas(
+            name="input_ids", plugin=_sglang_plugin, metas=metas
+        )
+        assert result is metas
+
+    def test_dims_already_present_returns_metas_unchanged(self):
+        """If dims is already in meta, metas are returned as-is."""
+        metas: list[dict] = [{**self._make_meta(cp_size=2, cp_rank=0), "dims": "t"}]
+        result = _ensure_dims_in_metas(
+            name="input_ids", plugin=_sglang_plugin, metas=metas
+        )
+        assert result is metas
+
+    def test_cp_sharded_sglang_input_ids_raises(self):
+        """CP + input_ids in sglang raises NotImplementedError."""
+        metas: list[dict] = [
+            self._make_meta(cp_size=2, cp_rank=0),
+            self._make_meta(cp_size=2, cp_rank=1),
+        ]
+        with pytest.raises(NotImplementedError, match="CP-sharded"):
+            _ensure_dims_in_metas(name="input_ids", plugin=_sglang_plugin, metas=metas)
+
+    def test_cp_sharded_sglang_positions_raises(self):
+        """CP + positions in sglang raises NotImplementedError."""
+        metas: list[dict] = [
+            self._make_meta(cp_size=2, cp_rank=0),
+            self._make_meta(cp_size=2, cp_rank=1),
+        ]
+        with pytest.raises(NotImplementedError, match="CP-sharded"):
+            _ensure_dims_in_metas(name="positions", plugin=_sglang_plugin, metas=metas)
+
+    def test_cp_sharded_megatron_input_ids_raises(self):
+        """CP + input_ids in megatron raises NotImplementedError."""
+        metas: list[dict] = [
+            {"megatron_parallel_info": {"cp_rank": 0, "cp_size": 2}},
+            {"megatron_parallel_info": {"cp_rank": 1, "cp_size": 2}},
+        ]
+        with pytest.raises(NotImplementedError, match="CP-sharded"):
+            _ensure_dims_in_metas(
+                name="input_ids", plugin=_megatron_plugin, metas=metas
+            )
+
+    def test_cp_non_sharded_name_returns_metas_unchanged(self):
+        """CP + non-sharded tensor name (seq_lens) returns metas as-is."""
+        metas: list[dict] = [
+            self._make_meta(cp_size=2, cp_rank=0),
+            self._make_meta(cp_size=2, cp_rank=1),
+        ]
+        result = _ensure_dims_in_metas(
+            name="seq_lens", plugin=_sglang_plugin, metas=metas
+        )
+        assert result is metas
+
+    def test_unknown_plugin_returns_metas_unchanged(self):
+        """CP + plugin with empty cp_sharded_names returns metas as-is."""
+
+        class _DummyPlugin(_SGLangPlugin):
+            @property
+            def cp_sharded_names(self) -> frozenset[str]:
+                return frozenset()
+
+        metas: list[dict] = [
+            self._make_meta(cp_size=2, cp_rank=0),
+            self._make_meta(cp_size=2, cp_rank=1),
+        ]
+        result = _ensure_dims_in_metas(
+            name="input_ids", plugin=_DummyPlugin(), metas=metas
+        )
+        assert result is metas
+
+
+class TestDetectPlugin:
+    def test_discriminating_names_sglang(self, tmp_path: Path) -> None:
+        fn: str = _save_pt(
+            tmp_path, name="seq_lens", step=0, rank=0, value=torch.tensor([3])
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = _detect_plugin(df, dump_path=tmp_path)
+
+        assert result is not None
+        assert result.name == "sglang"
+
+    def test_fallback_to_meta_based_detection(self, tmp_path: Path) -> None:
+        fn: str = _save_pt(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            value=torch.tensor([1, 2, 3]),
+            meta={"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}},
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn])
+
+        result = _detect_plugin(df, dump_path=tmp_path)
+
+        assert result is not None
+        assert result.name == "sglang"
+
+    def test_returns_none_no_match(self, tmp_path: Path) -> 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 = _detect_plugin(df, dump_path=tmp_path)
+
+        assert result is None
+
+
+class TestLoadNonTensorAux:
+    def test_multi_rank_mismatch_warning(self, tmp_path: Path) -> None:
+        fn0: str = _save_pt(tmp_path, name="rids", step=0, rank=0, value=["req_A"])
+        fn1: str = _save_pt(tmp_path, name="rids", step=0, rank=1, value=["req_B"])
+        df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
+
+        sink = WarningSink()
+        with sink.context() as warnings:
+            from unittest.mock import patch
+
+            with patch(
+                "sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
+                sink,
+            ):
+                result = _load_non_tensor_aux(
+                    name="rids", step=0, df=df, dump_path=tmp_path
+                )
+
+        assert result == ["req_A"]
+        assert len(warnings) == 1
+        assert isinstance(warnings[0], GeneralWarning)
+        assert "rids_mismatch" in warnings[0].category
+
+    def test_no_rows_returns_none(self, tmp_path: Path) -> None:
+        df: pl.DataFrame = _make_df_from_filenames([])
+
+        result = _load_non_tensor_aux(name="rids", step=0, df=df, dump_path=tmp_path)
+
+        assert result is None
+
+
+class TestLoadAndAlignAuxTensor:
+    def test_multi_rank_no_dims_emits_warning(self, tmp_path: Path) -> None:
+        fn0: str = _save_pt(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            value=torch.tensor([1, 2, 3]),
+            meta={
+                "sglang_parallel_info": {
+                    "tp_rank": 0,
+                    "tp_size": 2,
+                    "cp_rank": 0,
+                    "cp_size": 1,
+                }
+            },
+        )
+        fn1: str = _save_pt(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=1,
+            value=torch.tensor([4, 5, 6]),
+            meta={
+                "sglang_parallel_info": {
+                    "tp_rank": 1,
+                    "tp_size": 2,
+                    "cp_rank": 0,
+                    "cp_size": 1,
+                }
+            },
+        )
+        df: pl.DataFrame = _make_df_from_filenames([fn0, fn1])
+
+        sink = WarningSink()
+        with sink.context() as warnings:
+            from unittest.mock import patch
+
+            with patch(
+                "sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader.warning_sink",
+                sink,
+            ):
+                result = _load_and_align_aux_tensor(
+                    name="input_ids",
+                    step=0,
+                    df=df,
+                    dump_path=tmp_path,
+                    plugin=_sglang_plugin,
+                )
+
+        assert result is not None
+        assert torch.equal(result, torch.tensor([1, 2, 3]))
+        assert len(warnings) == 1
+        assert isinstance(warnings[0], GeneralWarning)
+        assert "aux_no_dims" in warnings[0].category
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -0,0 +1,144 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_plugins import (
+    _infer_positions,
+    _MegatronPlugin,
+    _SGLangPlugin,
+)
+from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+    PositionalSeqId,
+    SGLangSeqId,
+    TokenAlignerStepAux,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+_sglang_plugin = _SGLangPlugin()
+_megatron_plugin = _MegatronPlugin()
+
+
+class TestNormalizeSGLang:
+    """Tests for SGLang aux tensor normalization."""
+
+    def test_with_rids(self):
+        """SGLang tensors with rids produce string seq_ids."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30]),
+            "positions": torch.tensor([0, 1, 2]),
+            "seq_lens": torch.tensor([3]),
+            "rids": ["A"],
+        }
+
+        result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
+            step_data, layout="thd", step=0
+        )
+
+        assert result.input_ids == [10, 20, 30]
+        assert result.positions == [0, 1, 2]
+        assert result.seq_lens == [3]
+        assert result.seq_ids == [SGLangSeqId(rid="A")]
+
+    def test_rids_none_fallback(self):
+        """Missing rids results in (step, index) fallback seq_ids."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20]),
+            "positions": torch.tensor([0, 1]),
+            "seq_lens": torch.tensor([2]),
+        }
+
+        result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
+            step_data, layout="thd", step=3
+        )
+        assert result.seq_ids == [PositionalSeqId(step=3, seq_index=0)]
+
+    def test_multiple_seqs_with_rids(self):
+        """Multiple sequences with rids."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30, 40, 50]),
+            "positions": torch.tensor([0, 1, 2, 0, 1]),
+            "seq_lens": torch.tensor([3, 2]),
+            "rids": ["A", "B"],
+        }
+
+        result: TokenAlignerStepAux = _sglang_plugin.compute_step_aux(
+            step_data, layout="thd", step=0
+        )
+        assert result.seq_ids == [SGLangSeqId(rid="A"), SGLangSeqId(rid="B")]
+
+
+class TestNormalizeMegatron:
+    """Tests for Megatron aux tensor normalization."""
+
+    def test_cu_seqlens_to_seq_lens(self):
+        """cu_seqlens_q is converted to seq_lens via differencing."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30, 40, 50]),
+            "cu_seqlens_q": torch.tensor([0, 3, 5]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout="thd", step=0
+        )
+
+        assert result.seq_lens == [3, 2]
+
+    def test_positions_inferred_thd(self):
+        """Positions inferred from seq_lens in thd layout."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30, 40, 50]),
+            "cu_seqlens_q": torch.tensor([0, 3, 5]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout="thd", step=0
+        )
+
+        assert result.positions == [0, 1, 2, 0, 1]
+
+    def test_position_ids_passthrough(self):
+        """Explicit position_ids used directly instead of inference."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30, 40, 50]),
+            "position_ids": torch.tensor([5, 6, 7, 8, 9]),
+            "cu_seqlens_q": torch.tensor([0, 5]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout="thd", step=0
+        )
+
+        assert result.positions == [5, 6, 7, 8, 9]
+
+    def test_seq_ids_are_step_index_tuples(self):
+        """Megatron seq_ids are (step, seq_index) tuples."""
+        step_data: dict = {
+            "input_ids": torch.tensor([10, 20, 30, 40, 50]),
+            "cu_seqlens_q": torch.tensor([0, 3, 5]),
+        }
+
+        result: TokenAlignerStepAux = _megatron_plugin.compute_step_aux(
+            step_data, layout="thd", step=5
+        )
+        assert result.seq_ids == [
+            PositionalSeqId(step=5, seq_index=0),
+            PositionalSeqId(step=5, seq_index=1),
+        ]
+
+
+class TestInferPositions:
+    """Tests for position inference helper."""
+
+    def test_thd_multiple_sequences(self):
+        """thd: positions reset to 0 for each sequence."""
+        result = _infer_positions(
+            seq_lens=torch.tensor([2, 3]),
+        )
+        assert torch.equal(result, torch.tensor([0, 1, 0, 1, 2]))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -6,9 +6,11 @@ import pytest
 from sglang.srt.debug_utils.comparator.output_types import (
    ComparisonRecord,
    ConfigRecord,
+    GeneralWarning,
    ReplicatedMismatchWarning,
    SkipRecord,
    SummaryRecord,
+    WarningRecord,
    parse_record_json,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -104,7 +106,7 @@ class TestRecordTypes:
                    "diff_threshold": 1e-3,
                    "start_step": 0,
                    "end_step": 100,
-                },
+                }
            ),
            SkipRecord(name="attn", reason="no_baseline"),
            ComparisonRecord(
@@ -115,6 +117,9 @@ class TestRecordTypes:
                shape_mismatch=False,
            ),
            SummaryRecord(total=10, passed=8, failed=1, skipped=1),
+            WarningRecord(
+                warnings=[GeneralWarning(category="test", message="test warning")],
+            ),
        ]:
            restored = parse_record_json(record.model_dump_json())
            assert type(restored) is type(record)
@@ -186,6 +191,36 @@ class TestWarnings:
        assert restored_warning.baseline_index == 0
        assert restored_warning.max_abs_diff == pytest.approx(0.42)

+    def test_any_warning_discriminated_union_round_trip(self):
+        """All AnyWarning variants survive JSON round-trip via a WarningRecord."""
+        all_warnings = [
+            ReplicatedMismatchWarning(
+                axis="tp",
+                group_index=0,
+                differing_index=1,
+                baseline_index=0,
+                max_abs_diff=0.1,
+            ),
+            GeneralWarning(
+                category="aux_tensors_missing",
+                message="Aux tensors missing, skipping token alignment",
+            ),
+            GeneralWarning(
+                category="rids_mismatch",
+                message="rids mismatch across ranks: rank 0 has [1,2,3], "
+                "rank 1 has [4,5,6]",
+            ),
+        ]
+
+        record = WarningRecord(warnings=all_warnings)
+        restored = parse_record_json(record.model_dump_json())
+        assert isinstance(restored, WarningRecord)
+        assert len(restored.warnings) == len(all_warnings)
+
+        for original, parsed in zip(all_warnings, restored.warnings):
+            assert type(parsed) is type(original)
+            assert parsed == original
+

 if __name__ == "__main__":
    sys.exit(pytest.main([__file__]))
@@ -44,14 +44,18 @@ class TestMatchBundles:
        assert results[0].y[0].name == "t_a"

    def test_multiple_names_separate_bundles(self) -> None:
-        target_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a"),
-            _make_row(name="t_b"),
-        ])
-        baseline_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a"),
-            _make_row(name="t_b"),
-        ])
+        target_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a"),
+                _make_row(name="t_b"),
+            ]
+        )
+        baseline_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a"),
+                _make_row(name="t_b"),
+            ]
+        )

        results: list[Pair[TensorBundleInfo]] = match_bundles(
            dfs=Pair(x=baseline_df, y=target_df),
@@ -64,14 +68,18 @@ class TestMatchBundles:
        assert "t_b" in result_names

    def test_skip_rank_groups_across_ranks(self) -> None:
-        target_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a", rank=0),
-            _make_row(name="t_a", rank=1),
-        ])
-        baseline_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a", rank=0),
-            _make_row(name="t_a", rank=1),
-        ])
+        target_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a", rank=0),
+                _make_row(name="t_a", rank=1),
+            ]
+        )
+        baseline_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a", rank=0),
+                _make_row(name="t_a", rank=1),
+            ]
+        )

        results: list[Pair[TensorBundleInfo]] = match_bundles(
            dfs=Pair(x=baseline_df, y=target_df),
@@ -82,13 +90,17 @@ class TestMatchBundles:
        assert len(results[0].y) == 2

    def test_baseline_missing_tensor(self) -> None:
-        target_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a"),
-            _make_row(name="t_extra"),
-        ])
-        baseline_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a"),
-        ])
+        target_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a"),
+                _make_row(name="t_extra"),
+            ]
+        )
+        baseline_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a"),
+            ]
+        )

        results: list[Pair[TensorBundleInfo]] = match_bundles(
            dfs=Pair(x=baseline_df, y=target_df),
@@ -113,14 +125,18 @@ class TestMatchBundles:
        assert results == []

    def test_skip_step_groups_across_steps(self) -> None:
-        target_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a", step=0),
-            _make_row(name="t_a", step=1),
-        ])
-        baseline_df: pl.DataFrame = _make_df([
-            _make_row(name="t_a", step=0),
-            _make_row(name="t_a", step=1),
-        ])
+        target_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a", step=0),
+                _make_row(name="t_a", step=1),
+            ]
+        )
+        baseline_df: pl.DataFrame = _make_df(
+            [
+                _make_row(name="t_a", step=0),
+                _make_row(name="t_a", step=1),
+            ]
+        )

        results: list[Pair[TensorBundleInfo]] = match_bundles(
            dfs=Pair(x=baseline_df, y=target_df),
@@ -3,6 +3,13 @@ import sys
 import pytest
 from pydantic import ValidationError

+from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
+    PositionalSeqId,
+    TokenAlignerPlan,
+    TokenAlignerSeqInfo,
+    TokenAlignerStepAux,
+    TokenLocator,
+)
 from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
 from sglang.srt.debug_utils.comparator.output_types import (
    ComparisonRecord,
@@ -15,7 +22,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
    TensorInfo,
    TensorStats,
 )
-from sglang.srt.debug_utils.comparator.utils import _check_equal_lengths
+from sglang.srt.debug_utils.comparator.utils import Pair, _check_equal_lengths
 from sglang.test.ci.ci_register import register_cpu_ci

 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -33,6 +40,109 @@ class TestCheckEqualLengths:
            _check_equal_lengths(a=[1, 2], b=[3])


+class TestTokenAlignerStepAux:
+    def test_valid(self):
+        aux = TokenAlignerStepAux(
+            input_ids=[10, 20, 30],
+            positions=[0, 1, 2],
+            seq_lens=[2, 1],
+            seq_ids=[
+                PositionalSeqId(step=0, seq_index=0),
+                PositionalSeqId(step=0, seq_index=1),
+            ],
+        )
+        assert len(aux.input_ids) == 3
+
+    def test_token_length_mismatch(self):
+        with pytest.raises(ValueError, match="Length mismatch"):
+            TokenAlignerStepAux(
+                input_ids=[10, 20, 30],
+                positions=[0, 1],
+                seq_lens=[2, 1],
+                seq_ids=[
+                    PositionalSeqId(step=0, seq_index=0),
+                    PositionalSeqId(step=0, seq_index=1),
+                ],
+            )
+
+    def test_seq_length_mismatch(self):
+        with pytest.raises(ValueError, match="Length mismatch"):
+            TokenAlignerStepAux(
+                input_ids=[10, 20, 30],
+                positions=[0, 1, 2],
+                seq_lens=[2, 1],
+                seq_ids=[PositionalSeqId(step=0, seq_index=0)],
+            )
+
+    def test_sum_seq_lens_mismatch(self):
+        with pytest.raises(ValueError, match="sum\\(seq_lens\\)"):
+            TokenAlignerStepAux(
+                input_ids=[10, 20, 30],
+                positions=[0, 1, 2],
+                seq_lens=[1, 1],
+                seq_ids=[
+                    PositionalSeqId(step=0, seq_index=0),
+                    PositionalSeqId(step=0, seq_index=1),
+                ],
+            )
+
+
+class TestTokenAlignerSeqInfo:
+    def test_valid(self):
+        info = TokenAlignerSeqInfo(
+            input_ids=[10, 20, 30],
+            positions=[0, 1, 2],
+            locator=TokenLocator(token_index_in_step=[0, 1, 0]),
+        )
+        assert len(info.input_ids) == 3
+
+    def test_length_mismatch(self):
+        with pytest.raises(ValidationError):
+            TokenAlignerSeqInfo(
+                input_ids=[10, 20, 30],
+                positions=[0, 1, 2],
+                locator=TokenLocator(token_index_in_step=[0, 1]),
+            )
+
+    def test_positions_not_sequential(self):
+        with pytest.raises(ValidationError, match="positions must be"):
+            TokenAlignerSeqInfo(
+                input_ids=[10, 20, 30],
+                positions=[0, 2, 1],
+                locator=TokenLocator(token_index_in_step=[0, 1, 0]),
+            )
+
+
+class TestTokenAlignerPlan:
+    def test_valid(self):
+        plan = TokenAlignerPlan(
+            locators=Pair(
+                x=TokenLocator(token_index_in_step=[0, 1, 0]),
+                y=TokenLocator(token_index_in_step=[0, 0, 1]),
+            ),
+        )
+        assert len(plan.locators.x.token_index_in_step) == 3
+
+    def test_length_mismatch(self):
+        with pytest.raises(ValidationError, match="Length mismatch"):
+            TokenAlignerPlan(
+                locators=Pair(
+                    x=TokenLocator(token_index_in_step=[0, 1]),
+                    y=TokenLocator(token_index_in_step=[0, 0, 1]),
+                ),
+            )
+
+
+class TestSummaryRecord:
+    def test_valid(self):
+        record = SummaryRecord(total=10, passed=7, failed=2, skipped=1)
+        assert record.total == 10
+
+    def test_total_mismatch(self):
+        with pytest.raises(ValidationError, match="total=10"):
+            SummaryRecord(total=10, passed=5, failed=2, skipped=1)
+
+
 class TestAxisInfo:
    def test_valid(self):
        info = AxisInfo(axis_rank=0, axis_size=4)
@@ -59,16 +169,6 @@ class TestAxisInfo:
        assert info.axis_rank == 3


-class TestSummaryRecord:
-    def test_valid(self):
-        record = SummaryRecord(total=10, passed=7, failed=2, skipped=1)
-        assert record.total == 10
-
-    def test_total_mismatch(self):
-        with pytest.raises(ValidationError, match="total=10"):
-            SummaryRecord(total=10, passed=5, failed=2, skipped=1)
-
-
 def _make_tensor_info() -> TensorInfo:
    return TensorInfo(
        shape=[4, 4],