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

Support multi-step alignment and pipeline integration in dump comparator (#19378)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-26 10:23:22 +08:00
committed by GitHub
parent 4e843f1216
commit 265eb56d44
22 changed files with 535 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/debug_utils/comparator/aligner/__init__.py
View File
test/registered/debug_utils/comparator/aligner/entrypoint/__init__.py
View File
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
+2 -2
View File
@@ -151,8 +151,8 @@ class TestComputeAlignerPlan:
ta_plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(token_index_in_step=[0]),
y=TokenLocator(token_index_in_step=[0]),
x=TokenLocator(steps=[0], token_index_in_step=[0]),
y=TokenLocator(steps=[0], token_index_in_step=[0]),
),
)
test/registered/debug_utils/comparator/aligner/reorderer/__init__.py
View File
test/registered/debug_utils/comparator/aligner/token_aligner/__init__.py
View File
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py
+17 -9
View File
@@ -28,12 +28,13 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
class TestExecuteAlignment:
"""Tests for token alignment execution (single-step)."""
"""Tests for token alignment execution."""
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
hidden_step0 = torch.randn(5, 8) # 5 tokens, hidden_dim=8
hidden_step1 = torch.randn(2, 8) # 2 tokens
aux = TokenAlignerStepAux(
input_ids=[10, 20, 30, 40, 50],
@@ -41,9 +42,15 @@ class TestExecuteAlignment:
seq_lens=[3, 2],
seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
)
aux_step1 = TokenAlignerStepAux(
input_ids=[31, 51],
positions=[3, 2],
seq_lens=[1, 1],
seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
)
side_aux = TokenAlignerGlobalAux(
step_auxs={0: aux},
step_auxs={0: aux, 1: aux_step1},
framework="sglang",
layout="thd",
)
@@ -51,12 +58,13 @@ class TestExecuteAlignment:
index = build_seqs_info(side_aux)
plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index, y=index))
tensors = {0: hidden_step0, 1: hidden_step1}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan, tensor_pair=Pair(x=hidden, y=hidden)
plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors)
)
assert torch.equal(aligned.x, aligned.y)
assert aligned.x.shape[0] == len(plan.locators.x.token_index_in_step)
assert aligned.x.shape[0] == len(plan.locators.x.steps)
def test_zero_matched_tokens(self):
"""Empty TokenAlignerPlan (no matched tokens) returns shape[0]==0 without crash."""
@@ -64,14 +72,14 @@ class TestExecuteAlignment:
plan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(token_index_in_step=[]),
y=TokenLocator(token_index_in_step=[]),
x=TokenLocator(steps=[], token_index_in_step=[]),
y=TokenLocator(steps=[], token_index_in_step=[]),
),
)
tensor = torch.randn(5, 8)
tensors = {0: torch.randn(5, 8)}
aligned: Pair[torch.Tensor] = execute_token_aligner(
plan=plan, tensor_pair=Pair(x=tensor, y=tensor)
plan=plan, tensor_of_step_pair=Pair(x=tensors, y=tensors)
)
assert aligned.x.shape[0] == 0
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py
+216
View File
@@ -4,6 +4,7 @@ import pytest
from sglang.srt.debug_utils.comparator.aligner.token_aligner.planner import (
_match_sequences,
compute_token_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.seq_info_builder import (
build_seqs_info,
@@ -48,6 +49,7 @@ class TestBuildTokenIndexSGLangThd:
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.steps == [0, 0, 0]
assert seq_a.locator.token_index_in_step == [0, 1, 2]
seq_b = index.sequences[SGLangSeqId(rid="B")]
@@ -55,6 +57,93 @@ class TestBuildTokenIndexSGLangThd:
assert seq_b.positions == [0, 1]
assert seq_b.locator.token_index_in_step == [3, 4]
def test_multi_step_prefill_decode(self):
"""Prefill step followed by decode steps, sequences accumulate tokens."""
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")],
),
1: TokenAlignerStepAux(
input_ids=[31, 51],
positions=[3, 2],
seq_lens=[1, 1],
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, 31]
assert seq_a.positions == [0, 1, 2, 3]
assert seq_a.locator.steps == [0, 0, 0, 1]
seq_b = index.sequences[SGLangSeqId(rid="B")]
assert seq_b.input_ids == [40, 50, 51]
assert seq_b.positions == [0, 1, 2]
def test_sequence_exit_and_join(self):
"""Sequence A exits, new sequence D joins with different seq_id."""
side_aux = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20, 30],
positions=[0, 1, 2],
seq_lens=[3],
seq_ids=[SGLangSeqId(rid="A")],
),
1: TokenAlignerStepAux(
input_ids=[100, 200],
positions=[0, 1],
seq_lens=[2],
seq_ids=[SGLangSeqId(rid="D")],
),
},
framework="sglang",
layout="thd",
)
index = build_seqs_info(side_aux)
assert len(index.sequences) == 2
def test_different_seq_ids_produce_separate_sequences(self):
"""Different seq_ids at different steps → separate sequences."""
side_aux = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20],
positions=[0, 1],
seq_lens=[2],
seq_ids=[SGLangSeqId(rid="A")],
),
1: TokenAlignerStepAux(
input_ids=[100, 200, 300],
positions=[0, 1, 2],
seq_lens=[3],
seq_ids=[SGLangSeqId(rid="D")],
),
},
framework="sglang",
layout="thd",
)
index = build_seqs_info(side_aux)
assert len(index.sequences) == 2
all_input_ids = {
seq_id: rec.input_ids for seq_id, rec in index.sequences.items()
}
assert [10, 20] in all_input_ids.values()
assert [100, 200, 300] in all_input_ids.values()
class TestBuildTokenIndexMegatronThd:
"""Tests for Megatron thd token index building."""
@@ -83,6 +172,7 @@ class TestBuildTokenIndexMegatronThd:
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.steps == [0, 0, 0]
assert seq0.locator.token_index_in_step == [0, 1, 2]
seq1 = index.sequences[PositionalSeqId(step=0, seq_index=1)]
@@ -90,6 +180,44 @@ class TestBuildTokenIndexMegatronThd:
assert seq1.positions == [0, 1]
assert seq1.locator.token_index_in_step == [3, 4]
def test_multi_step_accumulation(self):
"""Two steps with different seq_ids produce separate sequences."""
side_aux = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20, 30, 40],
positions=[0, 1, 0, 1],
seq_lens=[2, 2],
seq_ids=[
PositionalSeqId(step=0, seq_index=0),
PositionalSeqId(step=0, seq_index=1),
],
),
1: TokenAlignerStepAux(
input_ids=[50, 60, 70, 80],
positions=[0, 1, 0, 1],
seq_lens=[2, 2],
seq_ids=[
PositionalSeqId(step=1, seq_index=0),
PositionalSeqId(step=1, seq_index=1),
],
),
},
framework="megatron",
layout="thd",
)
index = build_seqs_info(side_aux)
assert len(index.sequences) == 4
seq0 = index.sequences[PositionalSeqId(step=0, seq_index=0)]
assert seq0.input_ids == [10, 20]
assert seq0.locator.steps == [0, 0]
seq2 = index.sequences[PositionalSeqId(step=1, seq_index=0)]
assert seq2.input_ids == [50, 60]
assert seq2.locator.steps == [1, 1]
class TestMatchSequences:
"""Tests for _match_sequences: for each y, find matching x."""
@@ -240,6 +368,92 @@ class TestMatchSequences:
assert matched == []
class TestComputeAlignmentPlanCrossLayout:
"""Tests for alignment plan across different step distributions."""
def test_thd_vs_thd_different_step_splits(self):
"""Two thd sides with same tokens but different step distributions."""
side_aux_a = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20],
positions=[0, 1],
seq_lens=[2],
seq_ids=[SGLangSeqId(rid="X")],
),
1: TokenAlignerStepAux(
input_ids=[30],
positions=[2],
seq_lens=[1],
seq_ids=[SGLangSeqId(rid="X")],
),
},
framework="sglang",
layout="thd",
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20, 30],
positions=[0, 1, 2],
seq_lens=[3],
seq_ids=[SGLangSeqId(rid="X")],
),
},
framework="sglang",
layout="thd",
)
index_a = build_seqs_info(side_aux_a)
index_b = build_seqs_info(side_aux_b)
plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b))
assert len(plan.locators.x.steps) == 3
def test_sglang_vs_megatron_thd(self):
"""SGLang multi-step thd aligned with Megatron single-step thd."""
side_aux_a = 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")],
),
1: TokenAlignerStepAux(
input_ids=[31, 51],
positions=[3, 2],
seq_lens=[1, 1],
seq_ids=[SGLangSeqId(rid="A"), SGLangSeqId(rid="B")],
),
},
framework="sglang",
layout="thd",
)
side_aux_b = TokenAlignerGlobalAux(
step_auxs={
0: TokenAlignerStepAux(
input_ids=[10, 20, 30, 31, 40, 50, 51],
positions=[0, 1, 2, 3, 0, 1, 2],
seq_lens=[4, 3],
seq_ids=[
PositionalSeqId(step=0, seq_index=0),
PositionalSeqId(step=0, seq_index=1),
],
),
},
framework="megatron",
layout="thd",
)
index_a = build_seqs_info(side_aux_a)
index_b = build_seqs_info(side_aux_b)
plan = compute_token_aligner_plan(seqs_info_pair=Pair(x=index_a, y=index_b))
assert len(plan.locators.x.steps) == 7
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@@ -263,6 +477,7 @@ def _make_index(
input_ids=list(input_ids),
positions=list(range(num_tokens)),
locator=TokenLocator(
steps=[0] * num_tokens,
token_index_in_step=list(range(num_tokens)),
),
)
@@ -280,6 +495,7 @@ def _make_seq_info_dict(
input_ids=list(input_ids),
positions=list(range(num_tokens)),
locator=TokenLocator(
steps=[0] * num_tokens,
token_index_in_step=list(range(num_tokens)),
),
)
test/registered/debug_utils/comparator/aligner/unsharder/__init__.py
View File