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

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

@@ -51,13 +51,9 @@ def execute_aligner_plan(
# 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],
),
tensor_of_step_pair=Pair(x=step_tensors_x, y=step_tensors_y),
)
else:
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1

1
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_loader.py
View File

@@ -43,7 +43,6 @@ def load_and_normalize_aux(
available_names: set[str] = set(df["name"].unique().to_list()) & plugin.all_names
steps: list[int] = sorted(df["step"].unique().to_list())
assert len(steps) == 1, f"Multi-step not yet supported, got {len(steps)} steps"
tensor_names: set[str] = available_names & plugin.tensor_names
non_tensor_names: set[str] = available_names & plugin.non_tensor_names

31
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py
View File

@@ -4,20 +4,39 @@ import torch
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def execute_token_aligner(
plan: TokenAlignerPlan,
tensor_pair: Pair[torch.Tensor],
tensor_of_step_pair: Pair[dict[int, 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)
if not plan.locators.x.steps:
dummy: torch.Tensor = next(iter(tensor_of_step_pair.x.values()))
empty_shape: list[int] = [0] + list(dummy.shape[1:])
empty: torch.Tensor = torch.empty(empty_shape, dtype=dummy.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],
x=_extract_and_stack_tokens(
tensor_of_step=tensor_of_step_pair.x,
locator=plan.locators.x,
),
y=_extract_and_stack_tokens(
tensor_of_step=tensor_of_step_pair.y,
locator=plan.locators.y,
),
)
def _extract_and_stack_tokens(
*,
tensor_of_step: dict[int, torch.Tensor],
locator: TokenLocator,
) -> torch.Tensor:
tokens: list[torch.Tensor] = [
tensor_of_step[s][i] for s, i in zip(locator.steps, locator.token_index_in_step)
]
return torch.stack(tokens)

4
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py
View File

@@ -21,7 +21,7 @@ def compute_token_aligner_plan(
seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences)
)
_empty = TokenLocator(token_index_in_step=[])
_empty = TokenLocator(steps=[], token_index_in_step=[])
locator_x: TokenLocator = _empty
locator_y: TokenLocator = _empty
@@ -40,9 +40,11 @@ def compute_token_aligner_plan(
assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}"
locator_x = locator_x + TokenLocator(
steps=rec.x.locator.steps[:common_len],
token_index_in_step=rec.x.locator.token_index_in_step[:common_len],
)
locator_y = locator_y + TokenLocator(
steps=rec.y.locator.steps[:common_len],
token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
)

5
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py
View File

@@ -18,6 +18,7 @@ class _SeqInfoAccumulator:
input_ids: list[int] = field(default_factory=list)
positions: list[int] = field(default_factory=list)
steps: list[int] = field(default_factory=list)
token_index_in_step: list[int] = field(default_factory=list)
def extend(
@@ -25,10 +26,12 @@ class _SeqInfoAccumulator:
*,
input_ids: list[int],
positions: list[int],
steps: list[int],
token_index_in_step: list[int],
) -> None:
self.input_ids.extend(input_ids)
self.positions.extend(positions)
self.steps.extend(steps)
self.token_index_in_step.extend(token_index_in_step)
def build(self) -> TokenAlignerSeqInfo:
@@ -36,6 +39,7 @@ class _SeqInfoAccumulator:
input_ids=self.input_ids,
positions=self.positions,
locator=TokenLocator(
steps=self.steps,
token_index_in_step=self.token_index_in_step,
),
)
@@ -68,6 +72,7 @@ def _build_token_aligner_seq_infos(
accum[seq_id].extend(
input_ids=aux.input_ids[offset : offset + seq_len],
positions=aux.positions[offset : offset + seq_len],
steps=[step] * seq_len,
token_index_in_step=list(range(offset, offset + seq_len)),
)

9
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/types.py
View File

@@ -54,15 +54,17 @@ class TokenAlignerGlobalAux:
class TokenLocator(_FrozenBase):
"""Locates tokens within a single-step tensor.
"""Locates tokens within a multi-step tensor store.
token i is at tensor[token_index_in_step[i]].
token i is at tensor_of_step[steps[i]][token_index_in_step[i]].
"""
steps: list[int]
token_index_in_step: list[int]
def __add__(self, other: TokenLocator) -> TokenLocator:
return TokenLocator(
steps=self.steps + other.steps,
token_index_in_step=self.token_index_in_step + other.token_index_in_step,
)
@@ -81,6 +83,7 @@ class TokenAlignerSeqInfo(_FrozenBase):
_check_equal_lengths(
input_ids=self.input_ids,
positions=self.positions,
locator_steps=self.locator.steps,
locator_token_index_in_step=self.locator.token_index_in_step,
)
@@ -114,7 +117,9 @@ class TokenAlignerPlan(_FrozenBase):
@model_validator(mode="after")
def _validate_fields(self) -> TokenAlignerPlan:
_check_equal_lengths(
locators_x_steps=self.locators.x.steps,
locators_x_token_index_in_step=self.locators.x.token_index_in_step,
locators_y_steps=self.locators.y.steps,
locators_y_token_index_in_step=self.locators.y.token_index_in_step,
)
return self

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

@@ -3,7 +3,7 @@
from __future__ import annotations
from pathlib import Path
from typing import Any, Union
from typing import Any, Optional, Union
import torch
@@ -15,6 +15,9 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.planner import (
compute_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
SkipRecord,
@@ -35,6 +38,7 @@ def compare_bundle_pair(
filenames_pair: Pair[list[str]],
baseline_path: Path,
target_path: Path,
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
) -> Union[ComparisonRecord, SkipRecord]:
with warning_sink.context() as collected_warnings:
@@ -43,6 +47,7 @@ def compare_bundle_pair(
filenames_pair=filenames_pair,
baseline_path=baseline_path,
target_path=target_path,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
)
@@ -55,6 +60,7 @@ def _compare_bundle_pair_raw(
filenames_pair: Pair[list[str]],
baseline_path: Path,
target_path: Path,
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
) -> Union[ComparisonRecord, SkipRecord]:
# 1. Load (tensor + meta, ungrouped)
@@ -72,7 +78,7 @@ def _compare_bundle_pair_raw(
lambda items: [it.meta for it in items]
)
plan: AlignerPlan = compute_aligner_plan(
metas_pair=metas_pair, token_aligner_plan=None
metas_pair=metas_pair, token_aligner_plan=token_aligner_plan
)
# 3. Execute (tensor + plan only, no meta)

25
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -2,10 +2,19 @@ from __future__ import annotations
import argparse
from pathlib import Path
from typing import Iterator, Union
from typing import Iterator, Optional, Union
import polars as pl
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
AUX_NAMES,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.entrypoint import (
compute_maybe_token_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.bundle_comparator import compare_bundle_pair
from sglang.srt.debug_utils.comparator.bundle_matcher import (
TensorBundleInfo,
@@ -37,16 +46,22 @@ def run(args: argparse.Namespace) -> None:
warning_sink.set_output_format(args.output_format)
dfs: Pair[pl.DataFrame] = _read_df(args)
token_aligner_plan = compute_maybe_token_aligner_plan(args, dfs)
dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES)))
bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
dfs=dfs,
skip_keys=_compute_skip_keys(args),
skip_keys=_compute_skip_keys(
args, has_token_aligner_plan=token_aligner_plan is not None
),
)
comparison_records = _compare_bundle_pairs(
bundle_info_pairs=bundle_info_pairs,
baseline_path=Path(args.baseline_path),
target_path=Path(args.target_path),
token_aligner_plan=token_aligner_plan,
diff_threshold=args.diff_threshold,
)
_consume_comparison_records(
@@ -68,10 +83,12 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
return Pair(x=df_baseline, y=df_target)
def _compute_skip_keys(args: argparse.Namespace) -> set[str]:
def _compute_skip_keys(args, *, has_token_aligner_plan: bool):
skip_keys: set[str] = {"dump_index", "filename"}
if args.grouping == "logical":
skip_keys |= {"rank"}
if has_token_aligner_plan:
skip_keys |= {"step"}
return skip_keys
@@ -80,6 +97,7 @@ def _compare_bundle_pairs(
bundle_info_pairs: list[Pair[TensorBundleInfo]],
baseline_path: Path,
target_path: Path,
token_aligner_plan: Optional[TokenAlignerPlan],
diff_threshold: float,
) -> Iterator[Union[ComparisonRecord, SkipRecord]]:
for bundle_info_pair in bundle_info_pairs:
@@ -95,6 +113,7 @@ def _compare_bundle_pairs(
filenames_pair=filenames_pair,
baseline_path=baseline_path,
target_path=target_path,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
)

0
test/registered/debug_utils/comparator/__init__.py Normal file
View File

0
test/registered/debug_utils/comparator/aligner/__init__.py Normal file
View File

0
test/registered/debug_utils/comparator/aligner/entrypoint/__init__.py Normal file
View File

4
test/registered/debug_utils/comparator/aligner/entrypoint/test_planner.py
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]),
),
)

0
test/registered/debug_utils/comparator/aligner/reorderer/__init__.py Normal file
View File

0
test/registered/debug_utils/comparator/aligner/token_aligner/__init__.py Normal file
View File

26
test/registered/debug_utils/comparator/aligner/token_aligner/test_executor.py
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

216
test/registered/debug_utils/comparator/aligner/token_aligner/test_planner.py
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)),
),
)

0
test/registered/debug_utils/comparator/aligner/unsharder/__init__.py Normal file
View File

0
test/registered/debug_utils/comparator/tensor_comparator/__init__.py Normal file
View File

217
test/registered/debug_utils/comparator/test_entrypoint.py
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 -------------------

16
test/registered/debug_utils/comparator/test_model_validation.py
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]),
),
)

4
test/registered/debug_utils/test_dumper.py
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()

2
test/run_suite.py
View File

@@ -192,7 +192,7 @@ def run_a_suite(args):
files = [
f
for f in glob.glob("registered/**/*.py", recursive=True)
if not f.endswith("/conftest.py")
if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
]
# Strict: all registered files must have proper registration
sanity_check = True