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

Support token align with packed CP data in dump comparator (#19463)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:12:54 +08:00
committed by GitHub
parent 695e93b91f
commit 8293a914a6
6 changed files with 254 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,7 +1,7 @@
from __future__ import annotations
from pathlib import Path
from typing import Iterable, Optional, Tuple
from typing import Any, Optional
import polars as pl
import torch
@@ -24,7 +24,12 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
from sglang.srt.debug_utils.comparator.aligner.unsharder.parallel_info import (
normalize_parallel_info,
)
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.dims import (
ParallelAxis,
TokenLayout,
apply_dim_names,
parse_dim_names,
)
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.srt.debug_utils.dump_loader import ValueWithMeta, filter_rows
@@ -47,19 +52,20 @@ def load_and_normalize_aux(
non_tensor_names: set[str] = available_names & plugin.non_tensor_names
steps_data: dict[int, dict[str, object]] = {}
thd_seq_lens_by_step: dict[int, list[int]] = {}
for step in steps:
step_data = dict(
_load_step_data(
step=step,
tensor_names=tensor_names,
non_tensor_names=non_tensor_names,
df=df,
dump_path=dump_path,
plugin=plugin,
)
step_data, thd_seq_lens = _load_step_data(
step=step,
tensor_names=tensor_names,
non_tensor_names=non_tensor_names,
df=df,
dump_path=dump_path,
plugin=plugin,
)
if step_data:
steps_data[step] = step_data
if thd_seq_lens is not None:
thd_seq_lens_by_step[step] = thd_seq_lens
layout: TokenLayout = plugin.detect_layout(steps_data)
@@ -69,7 +75,10 @@ def load_and_normalize_aux(
}
return TokenAlignerGlobalAux(
step_auxs=step_auxs, framework=plugin.name, layout=layout
step_auxs=step_auxs,
framework=plugin.name,
layout=layout,
thd_seq_lens_by_step=thd_seq_lens_by_step or None,
)
@@ -104,19 +113,50 @@ def _load_step_data(
df: pl.DataFrame,
dump_path: Path,
plugin: _AuxFrameworkPlugin,
) -> Iterable[Tuple[str, object]]:
"""Load all tensor and non-tensor aux values for a single step."""
) -> tuple[dict[str, object], Optional[list[int]]]:
"""Load all tensor and non-tensor aux values for a single step.
Two-pass loading: non-CP-sharded tensors first (to obtain cu_seqlens_q
for seq_lens), then CP-sharded tensors with seq_lens for THD unshard/reorder.
Returns (step_data, thd_global_seq_lens).
"""
result: dict[str, object] = {}
# Pass 1: non-tensor values
for name in non_tensor_names:
value = _load_non_tensor_aux(name=name, step=step, df=df, dump_path=dump_path)
if value is not None:
yield name, value
result[name] = value
for name in tensor_names:
# Pass 1: non-CP-sharded tensors (e.g. cu_seqlens_q, seq_lens)
non_cp_tensor_names: set[str] = tensor_names - plugin.cp_sharded_names
cp_tensor_names: set[str] = tensor_names & plugin.cp_sharded_names
for name in non_cp_tensor_names:
tensor = _load_and_align_aux_tensor(
name=name, step=step, df=df, dump_path=dump_path, plugin=plugin
)
if tensor is not None:
yield name, tensor
result[name] = tensor
# Derive global seq_lens for THD unshard (framework-specific extraction)
thd_global_seq_lens: Optional[list[int]] = plugin.extract_global_seq_lens(result)
# Pass 2: CP-sharded tensors (input_ids, position_ids, etc.)
for name in cp_tensor_names:
tensor = _load_and_align_aux_tensor(
name=name,
step=step,
df=df,
dump_path=dump_path,
plugin=plugin,
thd_global_seq_lens=thd_global_seq_lens,
)
if tensor is not None:
result[name] = tensor
return result, thd_global_seq_lens
def _load_non_tensor_aux(
@@ -156,6 +196,7 @@ def _load_and_align_aux_tensor(
df: pl.DataFrame,
dump_path: Path,
plugin: _AuxFrameworkPlugin,
thd_global_seq_lens: Optional[list[int]] = None,
) -> Optional[torch.Tensor]:
"""Load an auxiliary tensor for (name, step), align if needed."""
rows = filter_rows(df, conditions={"name": name, "step": step})
@@ -175,14 +216,26 @@ def _load_and_align_aux_tensor(
if len(tensors) == 1:
return tensors[0]
metas: list[dict] = [item.meta for item in loaded]
metas = _ensure_dims_in_metas(name=name, plugin=plugin, metas=metas)
metas: list[dict[str, Any]] = [item.meta for item in loaded]
metas = _ensure_dims_in_metas(
name=name, plugin=plugin, metas=metas, ndim=tensors[0].ndim
)
sub_plans = compute_per_step_sub_plans(metas=metas)
sub_plans = compute_per_step_sub_plans(
metas=metas,
thd_global_seq_lens=(
thd_global_seq_lens if name in plugin.cp_sharded_names else None
),
)
if sub_plans:
dims_str: Optional[str] = metas[0].get("dims")
if dims_str is not None:
dim_names: list[str] = parse_dim_names(dims_str)
tensors = [apply_dim_names(t, dim_names) for t in tensors]
result = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert result is not None
return result
return result.rename(None) # strip named dims before returning to plugin
warning_sink.add(
GeneralWarning(
@@ -197,13 +250,16 @@ def _load_and_align_aux_tensor(
def _ensure_dims_in_metas(
*, name: str, plugin: _AuxFrameworkPlugin, metas: list[dict]
) -> list[dict]:
*,
name: str,
plugin: _AuxFrameworkPlugin,
metas: list[dict[str, Any]],
ndim: int,
) -> list[dict[str, Any]]:
"""Inject inferred dims into metas if not already present.
Returns metas unchanged if dims is already set, or a new list with dims
injected if inference succeeds. Raises if the tensor is CP-sharded
(not yet supported).
injected if inference succeeds for CP-sharded tensors.
"""
if metas[0].get("dims") is not None:
return metas
@@ -214,10 +270,7 @@ def _ensure_dims_in_metas(
return metas
if name in plugin.cp_sharded_names:
raise NotImplementedError(
f"Aux tensor '{name}' is CP-sharded but reorderer does not yet support "
f"zigzag reordering on the 't' dimension. "
f"Pass explicit dims= at dump time or wait for t-dim zigzag support."
)
inferred_dims: str = plugin.infer_cp_sharded_dims(name=name, ndim=ndim)
return [{**m, "dims": inferred_dims} for m in metas]
return metas

66
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/aux_plugins.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Optional
import torch
@@ -56,6 +57,21 @@ class _AuxFrameworkPlugin(ABC):
def all_names(self) -> frozenset[str]:
return self.tensor_names | self.non_tensor_names
def extract_global_seq_lens(
self, step_data: dict[str, object]
) -> Optional[list[int]]:
"""Extract per-seq token counts from loaded step data.
Returns None if this framework doesn't support THD / no relevant data available.
"""
return None
def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
"""Infer dims string for a CP-sharded aux tensor based on its ndim."""
raise NotImplementedError(
f"infer_cp_sharded_dims not implemented for {type(self).__name__}"
)
# ── sglang plugin ─────────────────────────────────────────────────
@@ -87,6 +103,30 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
return TokenLayout.T
def extract_global_seq_lens(
self, step_data: dict[str, object]
) -> Optional[list[int]]:
if not self.cp_sharded_names:
return None
seq_lens = step_data.get("seq_lens")
if not isinstance(seq_lens, torch.Tensor):
return None
return seq_lens.tolist()
def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
"""Infer dims for CP-sharded aux tensors.
NOTE: assumes zigzag ordering — natural-order CP without explicit dims
will be mishandled. Callers should set dims explicitly for non-zigzag CP.
"""
if ndim == 1:
return "t(cp,zigzag)"
raise ValueError(
f"SGLang: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
)
def compute_step_aux(
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux:
@@ -149,6 +189,32 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
def has_required_names(self, names: set[str]) -> bool:
return "input_ids" in names
def extract_global_seq_lens(
self, step_data: dict[str, object]
) -> Optional[list[int]]:
if not self.cp_sharded_names:
return None
cu_seqlens_q = step_data.get("cu_seqlens_q")
if not isinstance(cu_seqlens_q, torch.Tensor):
return None
return (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).tolist()
def infer_cp_sharded_dims(self, name: str, ndim: int) -> str:
"""Infer dims for CP-sharded aux tensors.
NOTE: assumes zigzag ordering — natural-order CP without explicit dims
will be mishandled. Callers should set dims explicitly for non-zigzag CP.
"""
if ndim == 1:
return "t(cp,zigzag)"
if ndim == 2:
return "b s(cp,zigzag)"
raise ValueError(
f"Megatron: cannot infer dims for CP-sharded '{name}' with ndim={ndim}"
)
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
for step_data in raw.values():
if (qkv_format := step_data.get("qkv_format")) is not None:

49
python/sglang/srt/debug_utils/comparator/aligner/token_aligner/entrypoint.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
@@ -26,10 +27,18 @@ 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(
@dataclass(frozen=True)
class TokenAlignerResult:
"""Result of token aligner computation, bundling the plan with THD metadata."""
plan: Optional[TokenAlignerPlan]
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]]
def compute_maybe_token_aligner_result(
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
) -> TokenAlignerResult:
if args.grouping == "logical":
if not (has_aux_tensors(dfs.x) and has_aux_tensors(dfs.y)):
warning_sink.add(
@@ -38,23 +47,34 @@ def compute_maybe_token_aligner_plan(
message="Aux tensors missing, skipping token alignment",
)
)
return None
return TokenAlignerResult(
plan=None, thd_seq_lens_by_step_pair=Pair(x=None, y=None)
)
return _build_token_aligner_plan(args=args, dfs=dfs)
return _build_token_aligner_result(args=args, dfs=dfs)
return None
return TokenAlignerResult(plan=None, thd_seq_lens_by_step_pair=Pair(x=None, y=None))
def _build_token_aligner_plan(
def _build_token_aligner_result(
*,
args: argparse.Namespace,
dfs: Pair[pl.DataFrame],
) -> Optional[TokenAlignerPlan]:
) -> TokenAlignerResult:
"""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)
baseline_aux: Optional[TokenAlignerGlobalAux] = load_and_normalize_aux(
dump_path=dump_paths.x, df=dfs.x
)
target_aux: Optional[TokenAlignerGlobalAux] = load_and_normalize_aux(
dump_path=dump_paths.y, df=dfs.y
)
thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
x=baseline_aux.thd_seq_lens_by_step if baseline_aux is not None else None,
y=target_aux.thd_seq_lens_by_step if target_aux is not None else None,
)
if baseline_aux is None or target_aux is None:
warning_sink.add(
@@ -63,7 +83,9 @@ def _build_token_aligner_plan(
message="Framework detection failed, skipping token alignment",
)
)
return None
return TokenAlignerResult(
plan=None, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair
)
global_aux: Pair[TokenAlignerGlobalAux] = Pair(
x=baseline_aux,
@@ -72,4 +94,9 @@ def _build_token_aligner_plan(
seqs_info: Pair[TokenAlignerSeqsInfo] = global_aux.map(build_seqs_info)
return compute_token_aligner_plan(seqs_info_pair=seqs_info)
plan: Optional[TokenAlignerPlan] = compute_token_aligner_plan(
seqs_info_pair=seqs_info
)
return TokenAlignerResult(
plan=plan, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair
)

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

@@ -1,7 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import NamedTuple, Union
from dataclasses import dataclass, field
from typing import NamedTuple, Optional, Union
from pydantic import model_validator
@@ -52,6 +52,7 @@ class TokenAlignerGlobalAux:
step_auxs: dict[int, TokenAlignerStepAux]
framework: str # "sglang" | "megatron"
layout: TokenLayout
thd_seq_lens_by_step: Optional[dict[int, list[int]]] = field(default=None)
class TokenLocator(_FrozenBase):