chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support token aligner planning and execution in dump comparator (#19377)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-26 10:04:33 +08:00
committed by GitHub
parent d34d5aca07
commit f9a2f0398f
13 changed files with 770 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/executor.py
View File

@@ -14,6 +14,9 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
execute_reorderer_plan,
)
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
execute_token_aligner,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
execute_unsharder_plan,
)
@@ -32,8 +35,9 @@ def execute_aligner_plan(
tensors_pair: Pair[list[torch.Tensor]],
plan: AlignerPlan,
) -> AlignerResult:
"""Execute unified unshard/reorder per side, then combine."""
"""Execute unified unshard/reorder + token-align."""
# Per-side: unshard + reorder -> dict[step, tensor]
step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans(
tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
)
@@ -45,11 +49,22 @@ def execute_aligner_plan(
failed_side_xy: str = "x" if not step_tensors_x else "y"
return AlignerResult(tensors=None, failed_side_xy=failed_side_xy)
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
)
# 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(
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
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
)
return AlignerResult(tensors=combined, failed_side_xy=None)

7
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import Any
from typing import Any, Optional
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
@@ -10,6 +10,9 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
compute_reorderer_plans,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
normalize_parallel_info,
)
@@ -23,11 +26,13 @@ from sglang.srt.debug_utils.comparator.utils import Pair
def compute_aligner_plan(
*,
metas_pair: Pair[list[dict[str, Any]]],
token_aligner_plan: Optional[TokenAlignerPlan],
) -> AlignerPlan:
return AlignerPlan(
per_step_plans=metas_pair.map(
lambda metas: _compute_per_step_plans(metas=metas)
),
token_aligner_plan=token_aligner_plan,
)

6
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/types.py
View File

@@ -1,9 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Union
from typing import Optional, Union
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.utils import Pair
@@ -20,3 +23,4 @@ class AlignerPerStepPlan:
@dataclass(frozen=True)
class AlignerPlan:
per_step_plans: Pair[list[AlignerPerStepPlan]]
token_aligner_plan: Optional[TokenAlignerPlan]

75
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py Normal file
View File

@@ -0,0 +1,75 @@
from __future__ import annotations
import argparse
from pathlib import Path
from typing import Optional
import polars as pl
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
has_aux_tensors,
load_and_normalize_aux,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
compute_token_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
build_seqs_info,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerGlobalAux,
TokenAlignerPlan,
TokenAlignerSeqsInfo,
)
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
def compute_maybe_token_aligner_plan(
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
if args.grouping == "logical":
if not (has_aux_tensors(dfs.x) and has_aux_tensors(dfs.y)):
warning_sink.add(
GeneralWarning(
category="aux_tensors_missing",
message="Aux tensors missing, skipping token alignment",
)
)
return None
return _build_token_aligner_plan(args=args, dfs=dfs)
return None
def _build_token_aligner_plan(
*,
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
"""Load aux tensors, build token indices, and compute the alignment plan."""
dump_paths: Pair[Path] = Pair(x=Path(args.baseline_path), y=Path(args.target_path))
baseline_aux = load_and_normalize_aux(dump_path=dump_paths.x, df=dfs.x)
target_aux = load_and_normalize_aux(dump_path=dump_paths.y, df=dfs.y)
if baseline_aux is None or target_aux is None:
warning_sink.add(
GeneralWarning(
category="framework_detection_failed",
message="Framework detection failed, skipping token alignment",
)
)
return None
global_aux: Pair[TokenAlignerGlobalAux] = Pair(
x=baseline_aux,
y=target_aux,
)
seqs_info: Pair[TokenAlignerSeqsInfo] = global_aux.map(build_seqs_info)
return compute_token_aligner_plan(seqs_info_pair=seqs_info)

23
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py Normal file
View File

@@ -0,0 +1,23 @@
from __future__ import annotations
import torch
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def execute_token_aligner(
plan: TokenAlignerPlan,
tensor_pair: Pair[torch.Tensor],
) -> Pair[torch.Tensor]:
if not plan.locators.x.token_index_in_step:
empty_shape: list[int] = [0] + list(tensor_pair.x.shape[1:])
empty: torch.Tensor = torch.empty(empty_shape, dtype=tensor_pair.x.dtype)
return Pair(x=empty, y=empty.clone())
return Pair(
x=tensor_pair.x[plan.locators.x.token_index_in_step],
y=tensor_pair.y[plan.locators.y.token_index_in_step],
)

130
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py Normal file
View File

@@ -0,0 +1,130 @@
from __future__ import annotations
from collections import defaultdict
from typing import NamedTuple, Optional
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SeqId,
TokenAlignerPlan,
TokenAlignerSeqInfo,
TokenAlignerSeqsInfo,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def compute_token_aligner_plan(
seqs_info_pair: Pair[TokenAlignerSeqsInfo],
) -> TokenAlignerPlan:
"""Compute a token alignment plan from two side token seqs_info_pair."""
matched_pairs: list[tuple[SeqId, SeqId]] = _match_sequences(
seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences)
)
_empty = TokenLocator(token_index_in_step=[])
locator_x: TokenLocator = _empty
locator_y: TokenLocator = _empty
for seq_id_x, seq_id_y in matched_pairs:
rec: Pair[TokenAlignerSeqInfo] = Pair(
x=seqs_info_pair.x.sequences[seq_id_x],
y=seqs_info_pair.y.sequences[seq_id_y],
)
# positions is validated to be [0, 1, ..., N-1], so position == index
# and the common range is simply [0, min(len_x, len_y)).
common_len: int = min(len(rec.x.positions), len(rec.y.positions))
x_ids = rec.x.input_ids[:common_len]
y_ids = rec.y.input_ids[:common_len]
assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}"
locator_x = locator_x + TokenLocator(
token_index_in_step=rec.x.locator.token_index_in_step[:common_len],
)
locator_y = locator_y + TokenLocator(
token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
)
return TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y))
# -------------------- Sequence matcher --------------------
def _match_sequences(
seqs: Pair[dict[SeqId, TokenAlignerSeqInfo]],
) -> list[tuple[SeqId, SeqId]]:
"""For each y (target) sequence, find a matching x (baseline) sequence.
Two-pass: exact match first, then prefix match for remaining.
"""
x_lookup: dict[tuple[int, ...], list[SeqId]] = defaultdict(list)
for seq_id, rec in seqs.x.items():
x_lookup[tuple(rec.input_ids)].append(seq_id)
claimed_x_ids: set[SeqId] = set()
matched_seq_id_pairs: list[tuple[SeqId, SeqId]] = []
for seq_id_y in sorted(seqs.y.keys()):
seq_y: TokenAlignerSeqInfo = seqs.y[seq_id_y]
matched_x: Optional[SeqId] = _find_matching_x_exact(
seq_y=seq_y, x_lookup=x_lookup, claimed_x_ids=claimed_x_ids
)
if matched_x is None:
matched_x = _find_matching_x_prefix(
seq_y=seq_y, x_seqs=seqs.x, claimed_x_ids=claimed_x_ids
)
if matched_x is not None:
matched_seq_id_pairs.append((matched_x, seq_id_y))
claimed_x_ids.add(matched_x)
return matched_seq_id_pairs
def _find_matching_x_exact(
*,
seq_y: TokenAlignerSeqInfo,
x_lookup: dict[tuple[int, ...], list[SeqId]],
claimed_x_ids: set[SeqId],
) -> Optional[SeqId]:
"""Find an x sequence with identical input_ids."""
ids_y_key: tuple[int, ...] = tuple(seq_y.input_ids)
candidates: list[SeqId] = x_lookup.get(ids_y_key, [])
for candidate in candidates:
if candidate not in claimed_x_ids:
return candidate
return None
class _PrefixCandidate(NamedTuple):
seq_id_x: SeqId
overlap_len: int
def _find_matching_x_prefix(
*,
seq_y: TokenAlignerSeqInfo,
x_seqs: dict[SeqId, TokenAlignerSeqInfo],
claimed_x_ids: set[SeqId],
) -> Optional[SeqId]:
"""Find the x sequence with the longest prefix relationship to y."""
ids_y: list[int] = seq_y.input_ids
candidates: list[_PrefixCandidate] = [
_PrefixCandidate(
seq_id_x=seq_id_x, overlap_len=min(len(seq_x.input_ids), len(ids_y))
)
for seq_id_x, seq_x in x_seqs.items()
if seq_id_x not in claimed_x_ids and _is_prefix_pair(seq_x.input_ids, ids_y)
]
if not candidates:
return None
return max(candidates, key=lambda c: c.overlap_len).seq_id_x
def _is_prefix_pair(a: list[int], b: list[int]) -> bool:
"""True if a is a prefix of b, or b is a prefix of a."""
shorter_len: int = min(len(a), len(b))
return a[:shorter_len] == b[:shorter_len]

76
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py Normal file
View File

@@ -0,0 +1,76 @@
from __future__ import annotations
from dataclasses import dataclass, field
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SeqId,
TokenAlignerGlobalAux,
TokenAlignerSeqInfo,
TokenAlignerSeqsInfo,
TokenAlignerStepAux,
TokenLocator,
)
@dataclass
class _SeqInfoAccumulator:
"""Mutable accumulator for building TokenAlignerSeqInfo without per-step validation."""
input_ids: list[int] = field(default_factory=list)
positions: list[int] = field(default_factory=list)
token_index_in_step: list[int] = field(default_factory=list)
def extend(
self,
*,
input_ids: list[int],
positions: list[int],
token_index_in_step: list[int],
) -> None:
self.input_ids.extend(input_ids)
self.positions.extend(positions)
self.token_index_in_step.extend(token_index_in_step)
def build(self) -> TokenAlignerSeqInfo:
return TokenAlignerSeqInfo(
input_ids=self.input_ids,
positions=self.positions,
locator=TokenLocator(
token_index_in_step=self.token_index_in_step,
),
)
def build_seqs_info(global_aux: TokenAlignerGlobalAux) -> TokenAlignerSeqsInfo:
"""Build sequence info for one side from its auxiliary tensors."""
return TokenAlignerSeqsInfo(
sequences=_build_token_aligner_seq_infos(global_aux),
layout=global_aux.layout,
)
def _build_token_aligner_seq_infos(
global_aux: TokenAlignerGlobalAux,
) -> dict[SeqId, TokenAlignerSeqInfo]:
"""Build token index for any framework/layout using seq_ids for identity tracking."""
accum: dict[SeqId, _SeqInfoAccumulator] = {}
for step in sorted(global_aux.step_auxs.keys()):
aux: TokenAlignerStepAux = global_aux.step_auxs[step]
offset: int = 0
for seq_index, seq_len in enumerate(aux.seq_lens):
seq_id: SeqId = aux.seq_ids[seq_index]
if seq_id not in accum:
accum[seq_id] = _SeqInfoAccumulator()
accum[seq_id].extend(
input_ids=aux.input_ids[offset : offset + seq_len],
positions=aux.positions[offset : offset + seq_len],
token_index_in_step=list(range(offset, offset + seq_len)),
)
offset += seq_len
return {seq_id: acc.build() for seq_id, acc in accum.items()}

4
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -71,7 +71,9 @@ def _compare_bundle_pair_raw(
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
lambda items: [it.meta for it in items]
)
plan: AlignerPlan = compute_aligner_plan(metas_pair=metas_pair)
plan: AlignerPlan = compute_aligner_plan(
metas_pair=metas_pair, token_aligner_plan=None
)
# 3. Execute (tensor + plan only, no meta)
tensors_pair: Pair[list[torch.Tensor]] = valid_pair.map(

6
test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py
View File

@@ -119,6 +119,7 @@ class TestExecuteAlignerPlan:
x=[self._make_step_plan(step=0, indices=[0, 1])],
y=[self._make_step_plan(step=0, indices=[0])],
),
token_aligner_plan=None,
)
tensors_pair: Pair[list[torch.Tensor]] = Pair(
@@ -139,6 +140,7 @@ class TestExecuteAlignerPlan:
x=[self._make_step_plan(step=0, indices=[0])],
y=[self._make_step_plan(step=0, indices=[0, 1])],
),
token_aligner_plan=None,
)
tensors_pair: Pair[list[torch.Tensor]] = Pair(
@@ -153,12 +155,13 @@ class TestExecuteAlignerPlan:
assert result.tensors is None
assert result.failed_side_xy == "y"
def test_single_step(self) -> None:
def test_no_token_aligner_single_step(self) -> None:
plan = AlignerPlan(
per_step_plans=Pair(
x=[self._make_step_plan(step=0, indices=[0])],
y=[self._make_step_plan(step=0, indices=[0])],
),
token_aligner_plan=None,
)
t_x: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -180,6 +183,7 @@ class TestExecuteAlignerPlan:
x=[self._make_step_plan(step=0, indices=[0])],
y=[self._make_step_plan(step=0, indices=[0])],
),
token_aligner_plan=None,
)
tensors_pair: Pair[list[torch.Tensor]] = Pair(

22
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
View File

@@ -136,10 +136,32 @@ class TestComputeAlignerPlan:
plan: AlignerPlan = compute_aligner_plan(
metas_pair=Pair(x=metas_x, y=metas_y),
token_aligner_plan=None,
)
assert len(plan.per_step_plans.x) == 1
assert len(plan.per_step_plans.y) == 1
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 (
TokenAlignerPlan,
TokenLocator,
)
ta_plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(token_index_in_step=[0]),
y=TokenLocator(token_index_in_step=[0]),
),
)
plan: AlignerPlan = compute_aligner_plan(
metas_pair=Pair(x=[_make_meta()], y=[_make_meta()]),
token_aligner_plan=ta_plan,
)
assert plan.token_aligner_plan is ta_plan
if __name__ == "__main__":

84
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py Normal file
View File

@@ -0,0 +1,84 @@
from __future__ import annotations
import sys
import pytest
import torch
from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
execute_token_aligner,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
compute_token_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
build_seqs_info,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerGlobalAux,
TokenAlignerPlan,
TokenAlignerStepAux,
TokenLocator,
)
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 TestExecuteAlignment:
"""Tests for token alignment execution (single-step)."""
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
aux = 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")],
)
side_aux = TokenAlignerGlobalAux(
step_auxs={0: aux},
framework="sglang",
layout="thd",
)
index = build_seqs_info(side_aux)
plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index, y=index))
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan, tensor_pair=Pair(x=hidden, y=hidden)
)
assert torch.equal(aligned.x, aligned.y)
assert aligned.x.shape[0] == len(plan.locators.x.token_index_in_step)
def test_zero_matched_tokens(self):
"""Empty TokenAlignerPlan (no matched tokens) returns shape[0]==0 without crash."""
torch.manual_seed(42)
plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(token_index_in_step=[]),
y=TokenLocator(token_index_in_step=[]),
),
)
tensor = torch.randn(5, 8)
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan, tensor_pair=Pair(x=tensor, y=tensor)
)
assert aligned.x.shape[0] == 0
assert aligned.y.shape[0] == 0
assert aligned.x.shape[1:] == (8,)
assert aligned.y.shape[1:] == (8,)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

306
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py Normal file
View File

@@ -0,0 +1,306 @@
import sys
import pytest
from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
_match_sequences,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
build_seqs_info,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
PositionalSeqId,
SeqId,
SGLangSeqId,
TokenAlignerGlobalAux,
TokenAlignerSeqInfo,
TokenAlignerSeqsInfo,
TokenAlignerStepAux,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=30, suite="default", nightly=True)
class TestBuildTokenIndexSGLangThd:
"""Tests for SGLang thd token index building."""
def test_single_step_prefill(self):
"""Single prefill step with two sequences."""
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")],
),
},
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]
assert seq_a.positions == [0, 1, 2]
assert seq_a.locator.token_index_in_step == [0, 1, 2]
seq_b = index.sequences[SGLangSeqId(rid="B")]
assert seq_b.input_ids == [40, 50]
assert seq_b.positions == [0, 1]
assert seq_b.locator.token_index_in_step == [3, 4]
class TestBuildTokenIndexMegatronThd:
"""Tests for Megatron thd token index building."""
def test_single_step_two_sequences(self):
"""Single step with two sequences in thd layout."""
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=[
PositionalSeqId(step=0, seq_index=0),
PositionalSeqId(step=0, seq_index=1),
],
),
},
framework="megatron",
layout="thd",
)
index = build_seqs_info(side_aux)
assert len(index.sequences) == 2
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.token_index_in_step == [0, 1, 2]
seq1 = index.sequences[PositionalSeqId(step=0, seq_index=1)]
assert seq1.input_ids == [40, 50]
assert seq1.positions == [0, 1]
assert seq1.locator.token_index_in_step == [3, 4]
class TestMatchSequences:
"""Tests for _match_sequences: for each y, find matching x."""
def test_exact_match_simple(self):
"""Identical input_ids on both sides → all matched."""
matched = _match_seqs(
x={0: (10, 20, 30), 1: (40, 50)},
y={0: (10, 20, 30), 1: (40, 50)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0)), (S(1), S(1))}
def test_exact_match_different_order(self):
"""Sequences in different order still match by content."""
matched = _match_seqs(
x={0: (10, 20), 1: (40, 50)},
y={0: (40, 50), 1: (10, 20)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(1), S(0)), (S(0), S(1))}
def test_exact_match_different_seq_ids(self):
"""Seq IDs don't need to correspond — matching is by content."""
matched = _match_seqs(
x={5: (10, 20), 9: (30, 40)},
y={2: (30, 40), 7: (10, 20)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(9), S(2)), (S(5), S(7))}
def test_no_match(self):
"""Completely different input_ids → no matches."""
matched = _match_seqs(
x={0: (10, 20)},
y={0: (99, 88)},
)
assert matched == []
def test_empty_sides(self):
"""Empty x or y → no matches."""
assert _match_seqs(x={}, y={0: (10,)}) == []
assert _match_seqs(x={0: (10,)}, y={}) == []
assert _match_seqs(x={}, y={}) == []
def test_x_has_more_sequences(self):
"""Extra x sequences are ignored (no y needs them)."""
matched = _match_seqs(
x={0: (10, 20), 1: (30, 40), 2: (50, 60)},
y={0: (30, 40)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(1), S(0))}
def test_y_has_more_sequences(self):
"""Extra y sequences remain unmatched."""
matched = _match_seqs(
x={0: (10, 20)},
y={0: (10, 20), 1: (30, 40), 2: (50, 60)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0))}
def test_one_x_not_reused(self):
"""Each x can only be claimed once, even if multiple y want it."""
matched = _match_seqs(
x={0: (10, 20)},
y={0: (10, 20), 1: (10, 20)},
)
assert len(matched) == 1
def test_ambiguous_all_matched(self):
"""Multiple identical sequences on both sides → all paired (greedy 1:1)."""
matched = _match_seqs(
x={0: (10, 20), 1: (10, 20), 2: (10, 20)},
y={0: (10, 20), 1: (10, 20), 2: (10, 20)},
)
S = _int_to_seq_id
assert len(matched) == 3
x_ids = {m[0] for m in matched}
y_ids = {m[1] for m in matched}
assert x_ids == {S(0), S(1), S(2)}
assert y_ids == {S(0), S(1), S(2)}
def test_prefix_x_shorter(self):
"""x has fewer tokens (prefix of y) → prefix match."""
matched = _match_seqs(
x={0: (10, 20)},
y={0: (10, 20, 30)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0))}
def test_prefix_y_shorter(self):
"""y has fewer tokens (prefix of x) → prefix match."""
matched = _match_seqs(
x={0: (10, 20, 30)},
y={0: (10, 20)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0))}
def test_prefix_picks_longest(self):
"""Among multiple prefix candidates, picks the one with longest overlap."""
matched = _match_seqs(
x={0: (10,), 1: (10, 20, 30)},
y={0: (10, 20, 30, 40)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(1), S(0))}
def test_exact_preferred_over_prefix(self):
"""Exact match is tried first, even if a longer prefix candidate exists."""
matched = _match_seqs(
x={0: (10, 20), 1: (10, 20, 30)},
y={0: (10, 20)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0))}
def test_prefix_fallback_after_exact(self):
"""Exact matches consume sequences, remaining use prefix match."""
matched = _match_seqs(
x={0: (10, 20, 30), 1: (40, 50)},
y={0: (10, 20, 30), 1: (40, 50, 60)},
)
S = _int_to_seq_id
assert len(matched) == 2
matched_set = _matched_ids(matched)
assert (S(0), S(0)) in matched_set
assert (S(1), S(1)) in matched_set
def test_single_token_sequences(self):
"""Single-token sequences can match."""
matched = _match_seqs(
x={0: (42,)},
y={0: (42,)},
)
S = _int_to_seq_id
assert _matched_ids(matched) == {(S(0), S(0))}
def test_no_partial_overlap_without_prefix(self):
"""Overlapping content that isn't a prefix → no match."""
matched = _match_seqs(
x={0: (10, 20, 30)},
y={0: (20, 30, 40)},
)
assert matched == []
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _int_to_seq_id(k: int) -> SeqId:
"""Convert an int key to a SeqId for test convenience."""
return SGLangSeqId(rid=str(k))
def _make_index(
*,
sequences: dict[int, tuple[int, ...]],
layout: str = "thd",
) -> TokenAlignerSeqsInfo:
"""Create a TokenAlignerSeqsInfo from simplified input_ids-only specification."""
records: dict[SeqId, TokenAlignerSeqInfo] = {}
for k, input_ids in sequences.items():
num_tokens = len(input_ids)
records[_int_to_seq_id(k)] = TokenAlignerSeqInfo(
input_ids=list(input_ids),
positions=list(range(num_tokens)),
locator=TokenLocator(
token_index_in_step=list(range(num_tokens)),
),
)
return TokenAlignerSeqsInfo(sequences=records, layout=layout)
def _make_seq_info_dict(
sequences: dict[int, tuple[int, ...]],
) -> dict[SeqId, TokenAlignerSeqInfo]:
"""Create a dict of TokenAlignerSeqInfo from {int_key: input_ids_tuple}."""
result: dict[SeqId, TokenAlignerSeqInfo] = {}
for k, input_ids in sequences.items():
num_tokens = len(input_ids)
result[_int_to_seq_id(k)] = TokenAlignerSeqInfo(
input_ids=list(input_ids),
positions=list(range(num_tokens)),
locator=TokenLocator(
token_index_in_step=list(range(num_tokens)),
),
)
return result
def _match_seqs(
*,
x: dict[int, tuple[int, ...]],
y: dict[int, tuple[int, ...]],
) -> list[tuple[SeqId, SeqId]]:
"""Shorthand: build SeqInfo dicts and call _match_sequences."""
return _match_sequences(
seqs=Pair(x=_make_seq_info_dict(x), y=_make_seq_info_dict(y))
)
def _matched_ids(matched: list[tuple[SeqId, SeqId]]) -> set[tuple[SeqId, SeqId]]:
"""Convert matched pairs list to set for order-independent comparison."""
return set(matched)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

28
test/registered/debug_utils/comparator/aligner/unsharder/test_executor.py
View File

@@ -94,11 +94,11 @@ class TestExecuteUnsharderPlan:
for tp_rank in range(4):
tensors.append(source[tp_rank])
with warning_sink.context() as _warnings:
with warning_sink.context():
intermediate = execute_unsharder_plan(plans[0], tensors)
assert len(intermediate) == 4
with warning_sink.context() as _warnings:
with warning_sink.context():
final = execute_unsharder_plan(plans[1], intermediate)
assert len(final) == 1
@@ -126,8 +126,8 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 2
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
@@ -168,8 +168,8 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 2
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
@@ -222,8 +222,8 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 3
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
@@ -271,8 +271,8 @@ class TestExecuteUnsharderPlan:
assert len(plans) == 3
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
@@ -357,8 +357,8 @@ class TestPickOperation:
assert len(plans) == 2
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1
@@ -387,8 +387,8 @@ class TestPickOperation:
assert all(isinstance(p.params, PickParams) for p in plans)
current = tensors
for plan in plans:
with warning_sink.context() as _warnings:
with warning_sink.context():
for plan in plans:
current = execute_unsharder_plan(plan, current)
assert len(current) == 1