chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

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