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

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

@@ -14,6 +14,9 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.executor import (
execute_reorderer_plan,
)
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.token_aligner.executor import (
execute_token_aligner,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
execute_unsharder_plan,
)
@@ -32,8 +35,9 @@ def execute_aligner_plan(
tensors_pair: Pair[list[torch.Tensor]],
plan: AlignerPlan,
) -> AlignerResult:
"""Execute unified unshard/reorder per side, then combine."""
"""Execute unified unshard/reorder + token-align."""
# Per-side: unshard + reorder -> dict[step, tensor]
step_tensors_x: dict[int, torch.Tensor] = _execute_step_plans(
tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
)
@@ -45,11 +49,22 @@ def execute_aligner_plan(
failed_side_xy: str = "x" if not step_tensors_x else "y"
return AlignerResult(tensors=None, failed_side_xy=failed_side_xy)
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
)
# 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],
),
)
else:
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
)
return AlignerResult(tensors=combined, failed_side_xy=None)

7
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/planner.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import Any
from typing import Any, Optional
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
@@ -10,6 +10,9 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
from sglang.srt.debug_utils.comparator.aligner.reorderer.planner import (
compute_reorderer_plans,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
normalize_parallel_info,
)
@@ -23,11 +26,13 @@ from sglang.srt.debug_utils.comparator.utils import Pair
def compute_aligner_plan(
*,
metas_pair: Pair[list[dict[str, Any]]],
token_aligner_plan: Optional[TokenAlignerPlan],
) -> AlignerPlan:
return AlignerPlan(
per_step_plans=metas_pair.map(
lambda metas: _compute_per_step_plans(metas=metas)
),
token_aligner_plan=token_aligner_plan,
)

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

@@ -1,9 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Union
from typing import Optional, Union
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.utils import Pair
@@ -20,3 +23,4 @@ class AlignerPerStepPlan:
@dataclass(frozen=True)
class AlignerPlan:
per_step_plans: Pair[list[AlignerPerStepPlan]]
token_aligner_plan: Optional[TokenAlignerPlan]

75
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py Normal file
View File

@@ -0,0 +1,75 @@
from __future__ import annotations
import argparse
from pathlib import Path
from typing import Optional
import polars as pl
from sglang.srt.debug_utils.comparator.aligner.token_aligner.aux_loader import (
has_aux_tensors,
load_and_normalize_aux,
)
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 (
TokenAlignerGlobalAux,
TokenAlignerPlan,
TokenAlignerSeqsInfo,
)
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
def compute_maybe_token_aligner_plan(
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
if args.grouping == "logical":
if not (has_aux_tensors(dfs.x) and has_aux_tensors(dfs.y)):
warning_sink.add(
GeneralWarning(
category="aux_tensors_missing",
message="Aux tensors missing, skipping token alignment",
)
)
return None
return _build_token_aligner_plan(args=args, dfs=dfs)
return None
def _build_token_aligner_plan(
*,
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
"""Load aux tensors, build token indices, and compute the alignment plan."""
dump_paths: Pair[Path] = Pair(x=Path(args.baseline_path), y=Path(args.target_path))
baseline_aux = load_and_normalize_aux(dump_path=dump_paths.x, df=dfs.x)
target_aux = load_and_normalize_aux(dump_path=dump_paths.y, df=dfs.y)
if baseline_aux is None or target_aux is None:
warning_sink.add(
GeneralWarning(
category="framework_detection_failed",
message="Framework detection failed, skipping token alignment",
)
)
return None
global_aux: Pair[TokenAlignerGlobalAux] = Pair(
x=baseline_aux,
y=target_aux,
)
seqs_info: Pair[TokenAlignerSeqsInfo] = global_aux.map(build_seqs_info)
return compute_token_aligner_plan(seqs_info_pair=seqs_info)

23
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/executor.py Normal file
View File

@@ -0,0 +1,23 @@
from __future__ import annotations
import torch
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
TokenAlignerPlan,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def execute_token_aligner(
plan: TokenAlignerPlan,
tensor_pair: Pair[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)
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],
)

130
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/planner.py Normal file
View File

@@ -0,0 +1,130 @@
from __future__ import annotations
from collections import defaultdict
from typing import NamedTuple, Optional
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SeqId,
TokenAlignerPlan,
TokenAlignerSeqInfo,
TokenAlignerSeqsInfo,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def compute_token_aligner_plan(
seqs_info_pair: Pair[TokenAlignerSeqsInfo],
) -> TokenAlignerPlan:
"""Compute a token alignment plan from two side token seqs_info_pair."""
matched_pairs: list[tuple[SeqId, SeqId]] = _match_sequences(
seqs=Pair(x=seqs_info_pair.x.sequences, y=seqs_info_pair.y.sequences)
)
_empty = TokenLocator(token_index_in_step=[])
locator_x: TokenLocator = _empty
locator_y: TokenLocator = _empty
for seq_id_x, seq_id_y in matched_pairs:
rec: Pair[TokenAlignerSeqInfo] = Pair(
x=seqs_info_pair.x.sequences[seq_id_x],
y=seqs_info_pair.y.sequences[seq_id_y],
)
# positions is validated to be [0, 1, ..., N-1], so position == index
# and the common range is simply [0, min(len_x, len_y)).
common_len: int = min(len(rec.x.positions), len(rec.y.positions))
x_ids = rec.x.input_ids[:common_len]
y_ids = rec.y.input_ids[:common_len]
assert x_ids == y_ids, f"{seq_id_x=} {seq_id_y=} {x_ids=} {y_ids=}"
locator_x = locator_x + TokenLocator(
token_index_in_step=rec.x.locator.token_index_in_step[:common_len],
)
locator_y = locator_y + TokenLocator(
token_index_in_step=rec.y.locator.token_index_in_step[:common_len],
)
return TokenAlignerPlan(locators=Pair(x=locator_x, y=locator_y))
# -------------------- Sequence matcher --------------------
def _match_sequences(
seqs: Pair[dict[SeqId, TokenAlignerSeqInfo]],
) -> list[tuple[SeqId, SeqId]]:
"""For each y (target) sequence, find a matching x (baseline) sequence.
Two-pass: exact match first, then prefix match for remaining.
"""
x_lookup: dict[tuple[int, ...], list[SeqId]] = defaultdict(list)
for seq_id, rec in seqs.x.items():
x_lookup[tuple(rec.input_ids)].append(seq_id)
claimed_x_ids: set[SeqId] = set()
matched_seq_id_pairs: list[tuple[SeqId, SeqId]] = []
for seq_id_y in sorted(seqs.y.keys()):
seq_y: TokenAlignerSeqInfo = seqs.y[seq_id_y]
matched_x: Optional[SeqId] = _find_matching_x_exact(
seq_y=seq_y, x_lookup=x_lookup, claimed_x_ids=claimed_x_ids
)
if matched_x is None:
matched_x = _find_matching_x_prefix(
seq_y=seq_y, x_seqs=seqs.x, claimed_x_ids=claimed_x_ids
)
if matched_x is not None:
matched_seq_id_pairs.append((matched_x, seq_id_y))
claimed_x_ids.add(matched_x)
return matched_seq_id_pairs
def _find_matching_x_exact(
*,
seq_y: TokenAlignerSeqInfo,
x_lookup: dict[tuple[int, ...], list[SeqId]],
claimed_x_ids: set[SeqId],
) -> Optional[SeqId]:
"""Find an x sequence with identical input_ids."""
ids_y_key: tuple[int, ...] = tuple(seq_y.input_ids)
candidates: list[SeqId] = x_lookup.get(ids_y_key, [])
for candidate in candidates:
if candidate not in claimed_x_ids:
return candidate
return None
class _PrefixCandidate(NamedTuple):
seq_id_x: SeqId
overlap_len: int
def _find_matching_x_prefix(
*,
seq_y: TokenAlignerSeqInfo,
x_seqs: dict[SeqId, TokenAlignerSeqInfo],
claimed_x_ids: set[SeqId],
) -> Optional[SeqId]:
"""Find the x sequence with the longest prefix relationship to y."""
ids_y: list[int] = seq_y.input_ids
candidates: list[_PrefixCandidate] = [
_PrefixCandidate(
seq_id_x=seq_id_x, overlap_len=min(len(seq_x.input_ids), len(ids_y))
)
for seq_id_x, seq_x in x_seqs.items()
if seq_id_x not in claimed_x_ids and _is_prefix_pair(seq_x.input_ids, ids_y)
]
if not candidates:
return None
return max(candidates, key=lambda c: c.overlap_len).seq_id_x
def _is_prefix_pair(a: list[int], b: list[int]) -> bool:
"""True if a is a prefix of b, or b is a prefix of a."""
shorter_len: int = min(len(a), len(b))
return a[:shorter_len] == b[:shorter_len]

76
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/seq_info_builder.py Normal file
View File

@@ -0,0 +1,76 @@
from __future__ import annotations
from dataclasses import dataclass, field
from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SeqId,
TokenAlignerGlobalAux,
TokenAlignerSeqInfo,
TokenAlignerSeqsInfo,
TokenAlignerStepAux,
TokenLocator,
)
@dataclass
class _SeqInfoAccumulator:
"""Mutable accumulator for building TokenAlignerSeqInfo without per-step validation."""
input_ids: list[int] = field(default_factory=list)
positions: list[int] = field(default_factory=list)
token_index_in_step: list[int] = field(default_factory=list)
def extend(
self,
*,
input_ids: list[int],
positions: list[int],
token_index_in_step: list[int],
) -> None:
self.input_ids.extend(input_ids)
self.positions.extend(positions)
self.token_index_in_step.extend(token_index_in_step)
def build(self) -> TokenAlignerSeqInfo:
return TokenAlignerSeqInfo(
input_ids=self.input_ids,
positions=self.positions,
locator=TokenLocator(
token_index_in_step=self.token_index_in_step,
),
)
def build_seqs_info(global_aux: TokenAlignerGlobalAux) -> TokenAlignerSeqsInfo:
"""Build sequence info for one side from its auxiliary tensors."""
return TokenAlignerSeqsInfo(
sequences=_build_token_aligner_seq_infos(global_aux),
layout=global_aux.layout,
)
def _build_token_aligner_seq_infos(
global_aux: TokenAlignerGlobalAux,
) -> dict[SeqId, TokenAlignerSeqInfo]:
"""Build token index for any framework/layout using seq_ids for identity tracking."""
accum: dict[SeqId, _SeqInfoAccumulator] = {}
for step in sorted(global_aux.step_auxs.keys()):
aux: TokenAlignerStepAux = global_aux.step_auxs[step]
offset: int = 0
for seq_index, seq_len in enumerate(aux.seq_lens):
seq_id: SeqId = aux.seq_ids[seq_index]
if seq_id not in accum:
accum[seq_id] = _SeqInfoAccumulator()
accum[seq_id].extend(
input_ids=aux.input_ids[offset : offset + seq_len],
positions=aux.positions[offset : offset + seq_len],
token_index_in_step=list(range(offset, offset + seq_len)),
)
offset += seq_len
return {seq_id: acc.build() for seq_id, acc in accum.items()}

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

@@ -71,7 +71,9 @@ def _compare_bundle_pair_raw(
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
lambda items: [it.meta for it in items]
)
plan: AlignerPlan = compute_aligner_plan(metas_pair=metas_pair)
plan: AlignerPlan = compute_aligner_plan(
metas_pair=metas_pair, token_aligner_plan=None
)
# 3. Execute (tensor + plan only, no meta)
tensors_pair: Pair[list[torch.Tensor]] = valid_pair.map(