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,
)