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

Update token layout and cleanup printer in dump comparator (#19457)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-27 08:09:15 +08:00
committed by GitHub
parent 8ac64e1487
commit eb0e905fc3
15 changed files with 56 additions and 393 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,7 +24,7 @@ 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
from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
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
@@ -61,7 +61,7 @@ def load_and_normalize_aux(
if step_data:
steps_data[step] = step_data
layout: str = plugin.detect_layout(steps_data)
layout: TokenLayout = plugin.detect_layout(steps_data)
step_auxs: dict[int, TokenAlignerStepAux] = {
step: plugin.compute_step_aux(step_data, layout=layout, step=step)

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

@@ -10,6 +10,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
SGLangSeqId,
TokenAlignerStepAux,
)
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
@@ -45,11 +46,11 @@ class _AuxFrameworkPlugin(ABC):
return frozenset()
@abstractmethod
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str: ...
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout: ...
@abstractmethod
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux: ...
@abstractmethod
@@ -89,11 +90,11 @@ class _SGLangPlugin(_AuxFrameworkPlugin):
def has_required_names(self, names: set[str]) -> bool:
return "input_ids" in names and "seq_lens" in names
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
return "thd"
def detect_layout(self, raw: dict[int, dict[str, object]]) -> TokenLayout:
return TokenLayout.T
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux:
input_ids = step_data["input_ids"]
positions = step_data["positions"]
@@ -154,13 +155,13 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
def has_required_names(self, names: set[str]) -> bool:
return "input_ids" in names and "cu_seqlens_q" in names
def detect_layout(self, raw: dict[int, dict[str, object]]) -> str:
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:
fmt = qkv_format if isinstance(qkv_format, str) else str(qkv_format)
if "bshd" in fmt.lower():
raise NotImplementedError(_BSHD_NOT_SUPPORTED_MSG)
return "thd"
return TokenLayout.T
input_ids = step_data.get("input_ids")
if isinstance(input_ids, torch.Tensor) and input_ids.ndim == 2:
@@ -175,10 +176,10 @@ class _MegatronPlugin(_AuxFrameworkPlugin):
),
)
)
return "thd"
return TokenLayout.T
def compute_step_aux(
self, step_data: dict[str, object], *, layout: str, step: int
self, step_data: dict[str, object], *, layout: TokenLayout, step: int
) -> TokenAlignerStepAux:
input_ids: torch.Tensor = step_data["input_ids"]

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

@@ -5,6 +5,7 @@ from typing import NamedTuple, Union
from pydantic import model_validator
from sglang.srt.debug_utils.comparator.dims import TokenLayout
from sglang.srt.debug_utils.comparator.utils import (
Pair,
_check_equal_lengths,
@@ -50,7 +51,7 @@ class TokenAlignerGlobalAux:
step_auxs: dict[int, TokenAlignerStepAux]
framework: str # "sglang" | "megatron"
layout: str # "thd"
layout: TokenLayout
class TokenLocator(_FrozenBase):
@@ -106,7 +107,7 @@ class TokenAlignerSeqsInfo(_FrozenBase):
"""All sequences for one side across all steps."""
sequences: dict[SeqId, TokenAlignerSeqInfo]
layout: str
layout: TokenLayout
class TokenAlignerPlan(_FrozenBase):

5
python/sglang/srt/debug_utils/comparator/dims.py
View File

@@ -8,6 +8,11 @@ BATCH_DIM_NAME: str = "b"
SEQ_DIM_NAME: str = "s"
class TokenLayout(Enum):
T = "t" # single flat token dim
BS = "bs" # separate batch + seq dims, need collapse
class ParallelAxis(Enum):
TP = "tp"
CP = "cp"

1
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
View File

@@ -125,6 +125,7 @@ def _compute_diff(
max_diff_coord=list(max_diff_coord),
baseline_at_max=x_baseline[max_diff_coord].item(),
target_at_max=x_target[max_diff_coord].item(),
diff_threshold=diff_threshold,
passed=(
rel_diff <= diff_threshold
and max_abs_diff <= diff_threshold

5
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
View File

@@ -69,13 +69,10 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list
def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
marker = "✅" if diff.passed else "❌"
return [
prefix_text
+ marker
+ " "
+ "\t".join(
f"{name}={value}"
f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
for name, value in [
("rel_diff", diff.rel_diff),
("max_abs_diff", diff.max_abs_diff),

86
python/sglang/srt/debug_utils/comparator/tensor_comparator/printer.py
View File

@@ -1,86 +0,0 @@
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorStats,
)
def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None:
baseline = info.baseline
target = info.target
dtype_marker = "" if baseline.dtype == target.dtype else "🟠"
print(
f"Raw "
f"[shape] {baseline.shape} vs {target.shape}\t"
f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}"
)
if info.unified_shape != baseline.shape:
print(
f"Unify shape: {baseline.shape} -> {info.unified_shape} "
f"(to match {target.shape})"
)
print(
f"After unify "
f"[shape] {info.unified_shape} vs {target.shape}\t"
f"[dtype] {baseline.dtype} vs {target.dtype}"
)
_print_stats_comparison(baseline=baseline.stats, target=target.stats)
if info.shape_mismatch:
print("⚠️ Shape mismatch")
return
if info.diff is not None:
_print_diff(
diff=info.diff,
diff_threshold=diff_threshold,
)
if info.diff_downcast is not None and info.downcast_dtype is not None:
_print_diff(
diff=info.diff_downcast,
diff_threshold=diff_threshold,
prefix_text=f"When downcast to {info.downcast_dtype}: ",
)
if baseline.sample is not None:
print(f"x_baseline(sample)={baseline.sample}")
if target.sample is not None:
print(f"x_target(sample)={target.sample}")
def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None:
stat_names = list(TensorStats.model_fields.keys())
for stat_name in stat_names:
value_baseline = getattr(baseline, stat_name)
value_target = getattr(target, stat_name)
if value_baseline is None or value_target is None:
continue
print(
f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
f"(diff: {value_target - value_baseline:.4f})"
)
def _print_diff(diff: DiffInfo, diff_threshold: float, prefix_text: str = "") -> None:
print(
prefix_text
+ "\t".join(
f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
for name, value in [
("rel_diff", diff.rel_diff),
("max_abs_diff", diff.max_abs_diff),
("mean_abs_diff", diff.mean_abs_diff),
]
)
)
print(
f"max_abs_diff happens at coord={diff.max_diff_coord} with "
f"baseline={diff.baseline_at_max} "
f"target={diff.target_at_max}"
)

1
python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py
View File

@@ -28,6 +28,7 @@ class DiffInfo(_StrictBase):
max_diff_coord: list[int]
baseline_at_max: float
target_at_max: float
diff_threshold: float
passed: bool