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/__init__.py
View File
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
test/registered/debug_utils/comparator/tensor_comparator/__init__.py
View File
test/registered/debug_utils/comparator/test_entrypoint.py
+217
View File
@@ -11,8 +11,10 @@ from sglang.srt.debug_utils.comparator.output_types import (
AnyRecord,
ComparisonRecord,
ConfigRecord,
GeneralWarning,
SkipRecord,
SummaryRecord,
WarningRecord,
_OutputRecord,
parse_record_json,
)
@@ -1012,6 +1014,221 @@ class TestEntrypointReplicatedAxis:
assert summary.passed == 0
class TestEntrypointAlignment:
"""Test `--grouping logical` with token alignment (aux tensors present)."""
def test_sglang_multi_step_alignment(self, tmp_path, capsys):
"""SGLang multi-step dumps with aux tensors auto-trigger alignment."""
torch.manual_seed(42)
hidden_dim = 8
hidden_step0 = torch.randn(5, hidden_dim)
hidden_step1 = torch.randn(2, hidden_dim)
exp_paths: list[Path] = []
for side_dir in ["baseline", "target"]:
d = tmp_path / side_dir
d.mkdir()
dumper = _Dumper(
config=DumperConfig(
enable=True,
dir=str(d),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
# Step 0: prefill with 2 sequences (3+2 tokens)
dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50]))
dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1]))
dumper.dump("seq_lens", torch.tensor([3, 2]))
dumper.dump("req_pool_indices", torch.tensor([7, 3]))
dumper.dump("rids", ["A", "B"])
dumper.dump("hidden_states", hidden_step0)
dumper.step()
# Step 1: decode (1 token per sequence)
dumper.dump("input_ids", torch.tensor([31, 51]))
dumper.dump("positions", torch.tensor([3, 2]))
dumper.dump("seq_lens", torch.tensor([1, 1]))
dumper.dump("req_pool_indices", torch.tensor([7, 3]))
dumper.dump("rids", ["A", "B"])
dumper.dump("hidden_states", hidden_step1)
dumper.step()
exp_paths.append(d / _FIXED_EXP_NAME)
args = _make_args(exp_paths[0], exp_paths[1], grouping="logical")
records = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
# AUX_NAMES are filtered out after plan computation → only hidden_states remains
assert len(comparisons) == 1
assert comparisons[0].name == "hidden_states"
assert comparisons[0].diff is not None
assert comparisons[0].diff.passed
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 1
assert summary.failed == 0
assert summary.skipped == 0
def test_sglang_vs_megatron_cross_framework(self, tmp_path, capsys):
"""SGLang 4-step thd baseline vs Megatron 1-step thd target align correctly."""
torch.manual_seed(42)
hidden_dim: int = 8
all_hiddens: torch.Tensor = torch.randn(11, hidden_dim)
seq_a_hiddens: torch.Tensor = all_hiddens[:6]
seq_b_hiddens: torch.Tensor = all_hiddens[6:]
# --- SGLang baseline: 1 prefill + 3 decode ---
sglang_dir: Path = tmp_path / "baseline"
sglang_dir.mkdir()
sglang_dumper = _Dumper(
config=DumperConfig(
enable=True,
dir=str(sglang_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
# Step 0: prefill — seq A (3 tokens) + seq B (2 tokens)
sglang_dumper.dump("input_ids", torch.tensor([10, 20, 30, 40, 50]))
sglang_dumper.dump("positions", torch.tensor([0, 1, 2, 0, 1]))
sglang_dumper.dump("seq_lens", torch.tensor([3, 2]))
sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3]))
sglang_dumper.dump("rids", ["A", "B"])
sglang_dumper.dump(
"hidden_states",
torch.stack(
[
seq_a_hiddens[0],
seq_a_hiddens[1],
seq_a_hiddens[2],
seq_b_hiddens[0],
seq_b_hiddens[1],
]
),
)
sglang_dumper.step()
# Steps 1-3: decode — 1 token per sequence
decode_data: list[dict[str, object]] = [
{
"input_ids": torch.tensor([31, 51]),
"positions": torch.tensor([3, 2]),
"hidden": torch.stack([seq_a_hiddens[3], seq_b_hiddens[2]]),
},
{
"input_ids": torch.tensor([32, 52]),
"positions": torch.tensor([4, 3]),
"hidden": torch.stack([seq_a_hiddens[4], seq_b_hiddens[3]]),
},
{
"input_ids": torch.tensor([33, 53]),
"positions": torch.tensor([5, 4]),
"hidden": torch.stack([seq_a_hiddens[5], seq_b_hiddens[4]]),
},
]
for step_data in decode_data:
sglang_dumper.dump("input_ids", step_data["input_ids"])
sglang_dumper.dump("positions", step_data["positions"])
sglang_dumper.dump("seq_lens", torch.tensor([1, 1]))
sglang_dumper.dump("req_pool_indices", torch.tensor([7, 3]))
sglang_dumper.dump("rids", ["A", "B"])
sglang_dumper.dump("hidden_states", step_data["hidden"])
sglang_dumper.step()
# --- Megatron target: 1 step, thd [T, H] ---
megatron_dir: Path = tmp_path / "target"
megatron_dir.mkdir()
megatron_dumper = _Dumper(
config=DumperConfig(
enable=True,
dir=str(megatron_dir),
exp_name=_FIXED_EXP_NAME,
enable_http_server=False,
)
)
# THD flat: seq A (6 tokens) + seq B (5 tokens) = 11 tokens total
megatron_input_ids: torch.Tensor = torch.tensor(
[10, 20, 30, 31, 32, 33, 40, 50, 51, 52, 53]
)
megatron_cu_seqlens: torch.Tensor = torch.tensor([0, 6, 11])
megatron_hidden: torch.Tensor = torch.cat([seq_a_hiddens, seq_b_hiddens], dim=0)
megatron_dumper.dump("input_ids", megatron_input_ids)
megatron_dumper.dump("cu_seqlens_q", megatron_cu_seqlens)
megatron_dumper.dump("hidden_states", megatron_hidden)
megatron_dumper.step()
# --- Run comparison ---
args = _make_args(
sglang_dir / _FIXED_EXP_NAME,
megatron_dir / _FIXED_EXP_NAME,
grouping="logical",
)
records = _run_and_parse(args, capsys)
warning_records = [r for r in records if isinstance(r, WarningRecord)]
layout_warnings = [
w
for wr in warning_records
for w in wr.warnings
if isinstance(w, GeneralWarning)
and w.category == "layout_detection_fallback"
]
assert len(layout_warnings) == 1
comparisons = _get_comparisons(records)
# AUX_NAMES filtered out → only hidden_states remains
assert len(comparisons) == 1
assert comparisons[0].name == "hidden_states"
assert comparisons[0].diff is not None
assert comparisons[0].diff.passed
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 1
assert summary.failed == 0
assert summary.skipped == 0
def test_alignment_fallback_when_no_aux(self, tmp_path, capsys):
"""Without aux tensors, logical grouping skips alignment and compares per-step."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=2)
args = _make_args(
baseline_path, target_path, grouping="logical", diff_threshold=0.1
)
capsys.readouterr()
run(args)
captured = capsys.readouterr()
records = _parse_jsonl(captured.out)
warning_records = [r for r in records if isinstance(r, WarningRecord)]
aux_missing_warnings = [
w
for wr in warning_records
for w in wr.warnings
if isinstance(w, GeneralWarning) and w.category == "aux_tensors_missing"
]
assert len(aux_missing_warnings) == 1
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 2
assert summary.passed == 2
# --------------------------- Assertion helpers -------------------
test/registered/debug_utils/comparator/test_model_validation.py
+8 -8
View File
@@ -92,7 +92,7 @@ class TestTokenAlignerSeqInfo:
info = TokenAlignerSeqInfo(
input_ids=[10, 20, 30],
positions=[0, 1, 2],
locator=TokenLocator(token_index_in_step=[0, 1, 0]),
locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
)
assert len(info.input_ids) == 3
@@ -101,7 +101,7 @@ class TestTokenAlignerSeqInfo:
TokenAlignerSeqInfo(
input_ids=[10, 20, 30],
positions=[0, 1, 2],
locator=TokenLocator(token_index_in_step=[0, 1]),
locator=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1, 0]),
)
def test_positions_not_sequential(self):
@@ -109,7 +109,7 @@ class TestTokenAlignerSeqInfo:
TokenAlignerSeqInfo(
input_ids=[10, 20, 30],
positions=[0, 2, 1],
locator=TokenLocator(token_index_in_step=[0, 1, 0]),
locator=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
)
@@ -117,18 +117,18 @@ 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]),
x=TokenLocator(steps=[0, 0, 1], token_index_in_step=[0, 1, 0]),
y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
),
)
assert len(plan.locators.x.token_index_in_step) == 3
assert len(plan.locators.x.steps) == 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]),
x=TokenLocator(steps=[0, 0], token_index_in_step=[0, 1]),
y=TokenLocator(steps=[0, 1, 1], token_index_in_step=[0, 0, 1]),
),
)
test/registered/debug_utils/test_dumper.py
+3 -1
View File
@@ -2127,7 +2127,9 @@ class TestRegisterForwardHook:
class TestPluginCoreFields:
def test_sglang_core_fields(self):
plugin = _SGLangPlugin()
assert plugin.core_fields() == frozenset({"input_ids", "positions", "seq_lens"})
assert plugin.core_fields() == frozenset(
{"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
)
def test_megatron_core_fields(self):
plugin = _MegatronPlugin()