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

Trace execution information in dump comparator (#19682)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:46:27 +08:00
committed by GitHub
parent 3ebd85bf1c
commit 5bf3deb4bc
17 changed files with 1180 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/debug_utils/comparator/__init__.py
View File

@@ -1,3 +1,6 @@
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( # noqa: F401
TracedAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401
AlignerPlan,
)

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

@@ -1,13 +1,19 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
from typing import NamedTuple, Optional
import torch
from sglang.srt.debug_utils.comparator.aligner.axis_aligner import (
execute_axis_aligner_plan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSidePlan,
TracedStepPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
AlignerPerStepSubPlan,
@@ -28,15 +34,31 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.executor import (
execute_unsharder_plan,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.srt.debug_utils.comparator.output_types import (
ReplicatedCheckResult,
ShapeSnapshot,
)
from sglang.srt.debug_utils.comparator.utils import Pair
class StepPlansResult(NamedTuple):
tensors: dict[int, torch.Tensor]
checks: list[ReplicatedCheckResult]
traced_side: TracedSidePlan
class SubPlansResult(NamedTuple):
tensor: Optional[torch.Tensor]
checks: list[ReplicatedCheckResult]
snapshots: list[ShapeSnapshot]
@dataclass(frozen=True)
class AlignerResult:
tensors: Optional[Pair[torch.Tensor]]
failed_side_xy: Optional[str] # "x" or "y"; None if success
replicated_checks: list[ReplicatedCheckResult] = field(default_factory=list)
traced_plan: Optional[TracedAlignerPlan] = None
def execute_aligner_plan(
@@ -48,26 +70,34 @@ def execute_aligner_plan(
all_checks: list[ReplicatedCheckResult] = []
# Per-side: unshard + reorder -> dict[step, tensor]
step_tensors_x, checks_x = _execute_step_plans(
result_x: StepPlansResult = _execute_step_plans(
tensors=tensors_pair.x, step_plans=plan.per_step_plans.x
)
all_checks.extend(checks_x)
all_checks.extend(result_x.checks)
step_tensors_y, checks_y = _execute_step_plans(
result_y: StepPlansResult = _execute_step_plans(
tensors=tensors_pair.y, step_plans=plan.per_step_plans.y
)
all_checks.extend(checks_y)
all_checks.extend(result_y.checks)
if not step_tensors_x or not step_tensors_y:
failed_side_xy: str = "x" if not step_tensors_x else "y"
traced_plan: TracedAlignerPlan = TracedAlignerPlan(
plan=plan,
per_side=Pair(x=result_x.traced_side, y=result_y.traced_side),
)
if not result_x.tensors or not result_y.tensors:
failed_side_xy: str = "x" if not result_x.tensors else "y"
return AlignerResult(
tensors=None,
failed_side_xy=failed_side_xy,
replicated_checks=all_checks,
traced_plan=traced_plan,
)
# Cross-side: token alignment (or direct extraction for single-step)
step_pair: Pair[dict[int, torch.Tensor]] = Pair(x=step_tensors_x, y=step_tensors_y)
step_pair: Pair[dict[int, torch.Tensor]] = Pair(
x=result_x.tensors, y=result_y.tensors
)
combined: Pair[torch.Tensor]
if plan.token_aligner_mode == "concat_steps":
combined = execute_token_aligner_concat_steps(tensor_of_step_pair=step_pair)
@@ -78,10 +108,10 @@ def execute_aligner_plan(
tensor_of_step_pair=step_pair,
)
else:
assert len(step_tensors_x) == 1 and len(step_tensors_y) == 1
assert len(result_x.tensors) == 1 and len(result_y.tensors) == 1
combined = Pair(
x=list(step_tensors_x.values())[0],
y=list(step_tensors_y.values())[0],
x=list(result_x.tensors.values())[0],
y=list(result_y.tensors.values())[0],
)
# Cross-side: axis alignment (squeeze singletons + rearrange dim order)
@@ -95,50 +125,78 @@ def execute_aligner_plan(
tensors=combined,
failed_side_xy=None,
replicated_checks=all_checks,
traced_plan=traced_plan,
)
def _execute_step_plans(
tensors: list[torch.Tensor],
step_plans: list[AlignerPerStepPlan],
) -> tuple[dict[int, torch.Tensor], list[ReplicatedCheckResult]]:
) -> StepPlansResult:
result: dict[int, torch.Tensor] = {}
all_checks: list[ReplicatedCheckResult] = []
traced_steps: list[TracedStepPlan] = []
for step_plan in step_plans:
step_tensors: list[torch.Tensor] = [
tensors[i] for i in step_plan.input_object_indices
]
tensor, checks = execute_sub_plans(
sub_result: SubPlansResult = execute_sub_plans(
tensors=step_tensors, plans=step_plan.sub_plans
)
all_checks.extend(checks)
if tensor is not None:
result[step_plan.step] = tensor
all_checks.extend(sub_result.checks)
return result, all_checks
traced_subs: list[TracedSubPlan] = [
TracedSubPlan(plan=sub_plan, snapshot=snapshot)
for sub_plan, snapshot in zip(step_plan.sub_plans, sub_result.snapshots)
]
traced_steps.append(
TracedStepPlan(
step=step_plan.step,
input_object_indices=step_plan.input_object_indices,
sub_plans=traced_subs,
)
)
if sub_result.tensor is not None:
result[step_plan.step] = sub_result.tensor
return StepPlansResult(
tensors=result,
checks=all_checks,
traced_side=TracedSidePlan(step_plans=traced_steps),
)
def execute_sub_plans(
tensors: list[torch.Tensor],
plans: list[AlignerPerStepSubPlan],
) -> tuple[Optional[torch.Tensor], list[ReplicatedCheckResult]]:
) -> SubPlansResult:
if not tensors:
return None, []
return SubPlansResult(tensor=None, checks=[], snapshots=[])
if not plans:
if len(tensors) != 1:
return None, []
return tensors[0], []
return SubPlansResult(tensor=None, checks=[], snapshots=[])
return SubPlansResult(tensor=tensors[0], checks=[], snapshots=[])
current: list[torch.Tensor] = tensors
all_checks: list[ReplicatedCheckResult] = []
all_snapshots: list[ShapeSnapshot] = []
for plan in plans:
input_shapes: list[list[int]] = [list(t.shape) for t in current]
current, checks = execute_sub_plan(tensors=current, plan=plan)
output_shapes: list[list[int]] = [list(t.shape) for t in current]
all_checks.extend(checks)
all_snapshots.append(
ShapeSnapshot(
input_shapes=input_shapes,
output_shapes=output_shapes,
)
)
assert len(current) == 1
return current[0], all_checks
return SubPlansResult(tensor=current[0], checks=all_checks, snapshots=all_snapshots)
def execute_sub_plan(

37
python/sglang/srt/debug_utils/comparator/aligner/entrypoint/traced_types.py Normal file
View File

@@ -0,0 +1,37 @@
"""Traced wrapper types that embed execution traces (ShapeSnapshots) into plan nodes.
These types are created *after* execution, pairing each sub-plan with its
observed shape snapshot so that downstream formatters never need to manually
zip plan + trace by index.
"""
from __future__ import annotations
from typing import Optional
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepSubPlan,
AlignerPlan,
)
from sglang.srt.debug_utils.comparator.output_types import ShapeSnapshot
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
class TracedSubPlan(_StrictBase):
plan: AlignerPerStepSubPlan
snapshot: Optional[ShapeSnapshot] = None
class TracedStepPlan(_StrictBase):
step: int
input_object_indices: list[int]
sub_plans: list[TracedSubPlan]
class TracedSidePlan(_StrictBase):
step_plans: list[TracedStepPlan]
class TracedAlignerPlan(_StrictBase):
plan: AlignerPlan
per_side: Pair[TracedSidePlan]

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

@@ -240,9 +240,11 @@ def _load_and_align_aux_tensor(
dim_names: list[str] = resolve_dim_names(dims_str)
tensors = [apply_dim_names(t, dim_names) for t in tensors]
result, _replicated_checks = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert result is not None
return result.rename(None) # strip named dims before returning to plugin
sub_result = execute_sub_plans(tensors=tensors, plans=sub_plans)
assert sub_result.tensor is not None
return sub_result.tensor.rename(
None
) # strip named dims before returning to plugin
log_sink.add(
InfoLog(

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

@@ -29,6 +29,8 @@ from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_ra
from sglang.srt.debug_utils.comparator.log_sink import log_sink
from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
from sglang.srt.debug_utils.comparator.output_types import (
BundleFileInfo,
BundleSideInfo,
ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
@@ -44,6 +46,37 @@ from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
def _collect_bundle_side_info(
items: list[ValueWithMeta],
metas: list[dict[str, Any]],
) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import (
PARALLEL_INFO_KEYS,
extract_parallel_info,
)
files: list[BundleFileInfo] = []
for item, meta in zip(items, metas):
assert isinstance(item.value, torch.Tensor)
tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {}
for key in PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append(
BundleFileInfo(
shape=list(tensor.shape),
dtype=str(tensor.dtype),
rank=meta.get("rank"),
parallel_info=parallel_info if parallel_info else None,
)
)
dims: Optional[str] = metas[0].get("dims") if metas else None
return BundleSideInfo(num_files=len(files), files=files, dims=dims)
def compare_bundle_pair(
*,
name: str,
@@ -186,6 +219,12 @@ def _compare_bundle_pair_tensor_type(
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
)
# Collect raw bundle info before alignment
raw_bundle_info: Pair[BundleSideInfo] = Pair(
x=_collect_bundle_side_info(items=valid_pair.x, metas=metas_pair.x),
y=_collect_bundle_side_info(items=valid_pair.y, metas=metas_pair.y),
)
# Apply dim names to tensors, then execute
tensors_pair: Pair[list[torch.Tensor]] = Pair(
x=_apply_dim_names_from_meta(
@@ -226,8 +265,9 @@ def _compare_bundle_pair_tensor_type(
)
record = TensorComparisonRecord(
**info.model_dump(),
aligner_plan=plan,
traced_plan=aligner_result.traced_plan,
replicated_checks=replicated_checks,
raw_bundle_info=raw_bundle_info,
)
if viz_output_dir is not None:

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

@@ -10,11 +10,11 @@ import polars as pl
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
RankInfoRecord,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
PARALLEL_INFO_KEYS: list[str] = ["sglang_parallel_info", "megatron_parallel_info"]
def emit_display_records(
@@ -68,8 +68,8 @@ def _collect_rank_info(
meta: dict[str, Any] = ValueWithMeta.load(dump_dir / row["filename"]).meta
row_data: dict[str, Any] = {"rank": row["rank"]}
for key in _PARALLEL_INFO_KEYS:
_extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
for key in PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
table_rows.append(row_data)
return table_rows or None
@@ -119,7 +119,7 @@ def _collect_input_ids_and_positions(
return table_rows or None
def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
def extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
if not info or info.get("error"):
return

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

@@ -31,12 +31,12 @@ from sglang.srt.debug_utils.comparator.output_types import (
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
report_sink,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.srt.debug_utils.comparator.utils import (
Pair,
auto_descend_dir,
@@ -53,7 +53,11 @@ def main() -> None:
def run(args: argparse.Namespace) -> int:
report_sink.configure(output_format=args.output_format, report_path=None)
report_sink.configure(
output_format=args.output_format,
report_path=None,
verbosity=args.verbosity,
)
dir_pair: Pair[Path] = Pair(
x=auto_descend_dir(Path(args.baseline_path), label="baseline_path"),
@@ -69,6 +73,16 @@ def run(args: argparse.Namespace) -> int:
Path(args.override_config) if args.override_config else None
)
report_path: Optional[Path] = _resolve_report_path(
target_path=dir_pair.y,
report_path_arg=args.report_path,
)
report_sink.configure(
output_format=args.output_format,
report_path=report_path,
verbosity=args.verbosity,
)
report_path: Optional[Path] = _resolve_report_path(
target_path=dir_pair.y,
report_path_arg=args.report_path,
@@ -294,6 +308,14 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
default="text",
help="Output format: text (default) or json (JSONL, one JSON object per line)",
)
parser.add_argument(
"--verbosity",
type=str,
choices=["minimal", "normal", "verbose"],
default="normal",
help="Output verbosity: minimal (1 line per tensor), normal (compact lifecycle), "
"verbose (full detail). Default: normal",
)
parser.add_argument(
"--preset",
type=str,

2
python/sglang/srt/debug_utils/comparator/log_sink.py
View File

@@ -27,8 +27,8 @@ class LogSink:
from sglang.srt.debug_utils.comparator.output_types import (
LogRecord,
_split_logs,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
errors, infos = _split_logs([log])
report_sink.add(LogRecord(errors=errors, infos=infos))

158
python/sglang/srt/debug_utils/comparator/output_types.py
View File

@@ -1,12 +1,12 @@
from __future__ import annotations
import sys
from abc import abstractmethod
from pathlib import Path
from typing import IO, TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
import polars as pl
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
from rich.console import RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison,
@@ -16,12 +16,15 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
)
from sglang.srt.debug_utils.comparator.utils import _StrictBase
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPlan,
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.report_sink import Verbosity
class BaseLog(_StrictBase):
@@ -59,6 +62,26 @@ class ReplicatedCheckResult(_StrictBase):
diff: Optional[DiffInfo] = None
class BundleFileInfo(_StrictBase):
"""Per-file info within a bundle (one rank's raw tensor)."""
shape: list[int]
dtype: str
rank: Optional[int] = None
parallel_info: Optional[dict[str, str]] = None # e.g. {"tp": "0/4", "ep": "1/2"}
class BundleSideInfo(_StrictBase):
num_files: int
files: list[BundleFileInfo]
dims: Optional[str] = None # e.g. "b s h(tp) d"
class ShapeSnapshot(_StrictBase):
input_shapes: list[list[int]]
output_shapes: list[list[int]]
class _OutputRecord(_StrictBase):
errors: list[ErrorLog] = Field(default_factory=list)
infos: list[InfoLog] = Field(default_factory=list)
@@ -66,6 +89,12 @@ class _OutputRecord(_StrictBase):
@abstractmethod
def _format_body(self) -> str: ...
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return self._format_body()
def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
return self._format_body()
def to_text(self) -> str:
body = self._format_body()
if self.errors:
@@ -87,6 +116,11 @@ class _BaseComparisonRecord(_OutputRecord):
return f"[step={self.location.step}] "
return ""
def _format_location_prefix_rich(self) -> str:
if self.location.step is not None:
return escape(f"[step={self.location.step}]") + " "
return ""
def _format_location_suffix(self) -> str:
if self.location.step is not None:
return f" (step={self.location.step})"
@@ -149,8 +183,9 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
model_config = ConfigDict(extra="forbid", defer_build=True)
type: Literal["comparison"] = "comparison"
aligner_plan: Optional[AlignerPlan] = None
traced_plan: Optional[TracedAlignerPlan] = None
replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
raw_bundle_info: Optional[Pair[BundleSideInfo]] = None
@property
def category(self) -> str:
@@ -164,8 +199,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
body: str = self._format_location_prefix() + format_comparison(self)
if self.replicated_checks:
body += "\n" + format_replicated_checks(self.replicated_checks)
if self.aligner_plan is not None:
body += "\n" + _format_aligner_plan(self.aligner_plan)
if self.traced_plan is not None:
body += "\n" + _format_aligner_plan(self.traced_plan)
return body
@@ -225,26 +260,48 @@ class LogRecord(_OutputRecord):
return ""
def _format_aligner_plan(plan: AlignerPlan) -> str:
def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
lines: list[str] = ["Aligner Plan:"]
for side_label, side_plans in [
("baseline", plan.per_step_plans.x),
("target", plan.per_step_plans.y),
for side_label, traced_side in [
("baseline", traced_plan.per_side.x),
("target", traced_plan.per_side.y),
]:
if not side_plans:
if not traced_side.step_plans:
lines.append(f" {side_label}: (no steps)")
continue
step_summaries: list[str] = []
for step_plan in side_plans:
sub_strs: list[str] = []
for sub in step_plan.sub_plans:
sub_strs.append(f"{sub.type}")
for traced_step in traced_side.step_plans:
sub_strs: list[str] = [
_format_sub_plan_text(traced_sub)
for traced_sub in traced_step.sub_plans
]
summary: str = ", ".join(sub_strs) if sub_strs else "passthrough"
step_summaries.append(f"step={step_plan.step}: {summary}")
step_summaries.append(f"step={traced_step.step}: {summary}")
lines.append(f" {side_label}: [{'; '.join(step_summaries)}]")
lines.extend(_format_cross_side_plan_text(traced_plan.plan))
return "\n".join(lines)
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot
in_count: int = len(snap.input_shapes)
out_count: int = len(snap.output_shapes)
in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?"
out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?"
sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}"
return sub_desc
def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]:
lines: list[str] = []
if plan.token_aligner_plan is not None:
num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
lines.append(f" token_aligner: {num_tokens} tokens aligned")
@@ -257,7 +314,7 @@ def _format_aligner_plan(plan: AlignerPlan) -> str:
parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
lines.append(f" axis_aligner: {', '.join(parts)}")
return "\n".join(lines)
return lines
AnyRecord = Annotated[
@@ -281,64 +338,3 @@ def _get_any_record_adapter() -> TypeAdapter:
def parse_record_json(json_str: str | bytes) -> AnyRecord:
return _get_any_record_adapter().validate_json(json_str)
def _print_to_stdout(record: _OutputRecord, *, output_format: str) -> None:
if output_format == "json":
print(record.model_dump_json())
else:
print(record.to_text())
class ReportSink:
"""Unified entry point for all record output."""
def __init__(self) -> None:
self._output_format: str = "text"
self._report_file: Optional[IO[str]] = None
self._report_path: Optional[Path] = None
def configure(
self,
*,
output_format: str = "text",
report_path: Optional[Path] = None,
) -> None:
self._output_format = output_format
if report_path is not None:
try:
report_path.parent.mkdir(parents=True, exist_ok=True)
self._report_file = open(report_path, "w", encoding="utf-8")
self._report_path = report_path
except OSError as exc:
print(
f"Warning: cannot open report file {report_path}: {exc}",
file=sys.stderr,
)
def add(self, record: _OutputRecord) -> None:
_print_to_stdout(record, output_format=self._output_format)
if self._report_file is not None:
self._report_file.write(record.model_dump_json())
self._report_file.write("\n")
self._report_file.flush()
def close(self) -> None:
if self._report_file is not None:
self._report_file.close()
self._report_file = None
@property
def report_path(self) -> Optional[Path]:
return self._report_path
def _reset(self) -> None:
"""Reset state for test isolation."""
self.close()
self._output_format = "text"
self._report_path = None
report_sink = ReportSink()

87
python/sglang/srt/debug_utils/comparator/report_sink.py Normal file
View File

@@ -0,0 +1,87 @@
from __future__ import annotations
import sys
from pathlib import Path
from typing import IO, Literal, Optional
from rich.console import Console
from sglang.srt.debug_utils.comparator.output_types import _OutputRecord
Verbosity = Literal["minimal", "normal", "verbose"]
class ReportSink:
"""Unified entry point for all record output."""
def __init__(self) -> None:
self._output_format: str = "text"
self._verbosity: Verbosity = "normal"
self._report_file: Optional[IO[str]] = None
self._report_path: Optional[Path] = None
self._console: Optional[Console] = None
@property
def verbosity(self) -> Verbosity:
return self._verbosity
def configure(
self,
*,
output_format: str = "text",
report_path: Optional[Path] = None,
verbosity: Verbosity = "normal",
) -> None:
self._output_format = output_format
self._verbosity = verbosity
if report_path is not None:
try:
report_path.parent.mkdir(parents=True, exist_ok=True)
self._report_file = open(report_path, "w", encoding="utf-8")
self._report_path = report_path
except OSError as exc:
print(
f"Warning: cannot open report file {report_path}: {exc}",
file=sys.stderr,
)
def add(self, record: _OutputRecord) -> None:
self._print_to_stdout(record)
if self._report_file is not None:
self._report_file.write(record.model_dump_json())
self._report_file.write("\n")
self._report_file.flush()
def close(self) -> None:
if self._report_file is not None:
self._report_file.close()
self._report_file = None
@property
def report_path(self) -> Optional[Path]:
return self._report_path
def _reset(self) -> None:
self.close()
self._output_format = "text"
self._verbosity = "normal"
self._report_path = None
self._console = None
def _get_console(self) -> Console:
if self._console is None:
self._console = Console()
return self._console
def _print_to_stdout(self, record: _OutputRecord) -> None:
if self._output_format == "json":
print(record.model_dump_json())
else:
console: Console = self._get_console()
console.print(record.to_rich(verbosity=self._verbosity))
console.print() # blank line between records
report_sink = ReportSink()

55
test/registered/debug_utils/comparator/aligner/entrypoint/test_executor.py
View File

@@ -5,6 +5,8 @@ import torch
from sglang.srt.debug_utils.comparator.aligner.entrypoint.executor import (
AlignerResult,
StepPlansResult,
SubPlansResult,
_execute_step_plans,
execute_aligner_plan,
execute_sub_plan,
@@ -31,25 +33,28 @@ register_cpu_ci(est_time=15, suite="default", nightly=True)
class TestExecuteSubPlans:
def test_empty_tensors_returns_none(self) -> None:
result, checks = execute_sub_plans(tensors=[], plans=[])
assert result is None
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=[], plans=[])
assert r.tensor is None
assert r.checks == []
assert r.snapshots == []
def test_no_plans_single_tensor_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
result, checks = execute_sub_plans(tensors=[tensor], plans=[])
assert result is not None
assert torch.equal(result, tensor)
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=[tensor], plans=[])
assert r.tensor is not None
assert torch.equal(r.tensor, tensor)
assert r.checks == []
assert r.snapshots == []
def test_no_plans_multiple_tensors_returns_none(self) -> None:
tensors: list[torch.Tensor] = [
torch.tensor([1.0]),
torch.tensor([2.0]),
]
result, checks = execute_sub_plans(tensors=tensors, plans=[])
assert result is None
assert checks == []
r: SubPlansResult = execute_sub_plans(tensors=tensors, plans=[])
assert r.tensor is None
assert r.checks == []
assert r.snapshots == []
def test_with_unsharder_plan(self) -> None:
t0: torch.Tensor = torch.tensor([[1.0, 2.0]]).refine_names("b", "h")
@@ -61,12 +66,13 @@ class TestExecuteSubPlans:
groups=[[0, 1]],
)
result, checks = execute_sub_plans(tensors=[t0, t1], plans=[plan])
r: SubPlansResult = execute_sub_plans(tensors=[t0, t1], plans=[plan])
assert result is not None
assert r.tensor is not None
expected: torch.Tensor = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
assert torch.equal(result.rename(None), expected)
assert checks == []
assert torch.equal(r.tensor.rename(None), expected)
assert r.checks == []
assert len(r.snapshots) == 1
class TestExecuteSubPlan:
@@ -91,10 +97,13 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result, checks = _execute_step_plans(tensors=tensors, step_plans=[step_plan])
r: StepPlansResult = _execute_step_plans(
tensors=tensors, step_plans=[step_plan]
)
assert result == {}
assert checks == []
assert r.tensors == {}
assert r.checks == []
assert len(r.traced_side.step_plans) == 1
def test_single_step_passthrough(self) -> None:
tensor: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -105,11 +114,15 @@ class TestExecuteStepPlans:
sub_plans=[],
)
result, checks = _execute_step_plans(tensors=[tensor], step_plans=[step_plan])
r: StepPlansResult = _execute_step_plans(
tensors=[tensor], step_plans=[step_plan]
)
assert 5 in result
assert torch.equal(result[5], tensor)
assert checks == []
assert 5 in r.tensors
assert torch.equal(r.tensors[5], tensor)
assert r.checks == []
assert len(r.traced_side.step_plans) == 1
assert r.traced_side.step_plans[0].step == 5
class TestExecuteAlignerPlan:

27
test/registered/debug_utils/comparator/conftest.py
View File

@@ -1,12 +1,37 @@
import sys
import warnings
from pathlib import Path
warnings.filterwarnings(
"ignore", message="builtin type Swig.*", category=DeprecationWarning
)
# Add the test root to sys.path so `registered.debug_utils.comparator.testing_helpers`
# can be imported by test modules.
_TEST_ROOT: Path = Path(__file__).resolve().parents[3]
if str(_TEST_ROOT) not in sys.path:
sys.path.insert(0, str(_TEST_ROOT))
import pytest
from sglang.srt.debug_utils.comparator.output_types import report_sink
from sglang.srt.debug_utils.comparator.report_sink import report_sink
collect_ignore_glob: list[str] = []
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line(
"filterwarnings",
"ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning",
)
config.addinivalue_line(
"filterwarnings",
"ignore:builtin type Swig.*:DeprecationWarning",
)
config.addinivalue_line(
"filterwarnings",
"ignore:Named tensors and all their associated APIs:UserWarning",
)
collect_ignore_glob: list[str] = []

10
test/registered/debug_utils/comparator/test_display.py
View File

@@ -9,8 +9,8 @@ import torch
from sglang.srt.debug_utils.comparator.display import (
_collect_input_ids_and_positions,
_collect_rank_info,
_extract_parallel_info,
_render_polars_as_text,
extract_parallel_info,
)
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
@@ -359,26 +359,26 @@ class TestExtractParallelInfo:
"pp_size": 2,
}
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info=info)
extract_parallel_info(row_data=row_data, info=info)
assert row_data["tp"] == "1/4"
assert row_data["pp"] == "0/2"
def test_skips_error_info(self) -> None:
row_data: dict = {}
_extract_parallel_info(
extract_parallel_info(
row_data=row_data, info={"error": True, "tp_rank": 0, "tp_size": 1}
)
assert row_data == {}
def test_skips_empty_info(self) -> None:
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info={})
extract_parallel_info(row_data=row_data, info={})
assert row_data == {}
def test_ignores_rank_without_size(self) -> None:
row_data: dict = {}
_extract_parallel_info(row_data=row_data, info={"tp_rank": 0})
extract_parallel_info(row_data=row_data, info={"tp_rank": 0})
assert "tp" not in row_data

106
test/registered/debug_utils/comparator/test_entrypoint.py
View File

@@ -1458,7 +1458,7 @@ class TestEntrypointConcatMode:
assert len(comparisons) == 3
def test_concat_aligner_plan_fields(self, tmp_path, capsys):
"""TensorComparisonRecord.aligner_plan reports mode='concat' with plan=None."""
"""TensorComparisonRecord.traced_plan reports mode='concat' with plan=None."""
torch.manual_seed(42)
records = self._run_concat(
@@ -1470,8 +1470,9 @@ class TestEntrypointConcatMode:
)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
plan = comparisons[0].aligner_plan
assert plan is not None
traced_plan = comparisons[0].traced_plan
assert traced_plan is not None
plan = traced_plan.plan
assert plan.token_aligner_mode == "concat_steps"
assert plan.token_aligner_plan is None
@@ -4250,7 +4251,7 @@ class TestReportOutput:
def test_streaming_flush(self, tmp_path, capsys):
"""Report file is flushed after each record (readable before close)."""
from sglang.srt.debug_utils.comparator.output_types import report_sink
from sglang.srt.debug_utils.comparator.report_sink import report_sink
report_file: Path = tmp_path / "stream_report.jsonl"
report_sink.configure(
@@ -4362,5 +4363,102 @@ class TestEntrypointAutoDescend:
run(parse_args(argv))
class TestEntrypointDpAttentionMissingAlias:
"""Regression: dp-attention without ``# dp:=attn_dp`` → shape mismatch failure.
In dp-attention mode (tp_size=2, attn_dp_size=2), layer_input is dumped
after prepare_attn which DP-distributes tokens. One rank gets 0 tokens
(shape [0, H]), the other gets all tokens (shape [T, H]).
Without ``# dp:=attn_dp`` in dims, the comparator has no dp_rank/dp_size
to filter on, so it picks one rank via TP pick — potentially the empty
one — causing a shape mismatch with the baseline.
"""
@staticmethod
def _sglang_dp_attn_parallel_info(*, tp_rank: int) -> dict:
return {
"tp_rank": tp_rank,
"tp_size": 2,
"pp_rank": 0,
"pp_size": 1,
"moe_ep_rank": 0,
"moe_ep_size": 1,
"moe_tp_rank": tp_rank,
"moe_tp_size": 2,
"moe_dp_rank": 0,
"moe_dp_size": 1,
"enable_dp_attention": True,
"attn_tp_rank": 0,
"attn_tp_size": 1,
"attn_dp_rank": tp_rank,
"attn_dp_size": 2,
"local_attn_dp_rank": tp_rank,
"local_attn_dp_size": 2,
"attn_cp_rank": 0,
"attn_cp_size": 1,
}
def test_missing_dp_alias_causes_shape_mismatch(
self, tmp_path: Path, capsys
) -> None:
"""dims='t h' (no dp:=attn_dp) → comparator picks empty rank → shape_mismatch failure."""
torch.manual_seed(42)
tensor_data: torch.Tensor = torch.randn(5, 8)
target_data: torch.Tensor = tensor_data + torch.randn(5, 8) * 0.001
for side_name, data in [("baseline", tensor_data), ("target", target_data)]:
side_dir: Path = tmp_path / side_name
side_dir.mkdir()
# Baseline: single rank, no DP attention
if side_name == "baseline":
_create_rank_dump(
side_dir,
rank=0,
name="layer_input",
tensor=data,
dims="t h",
parallel_info={"tp_rank": 0, "tp_size": 1},
framework="sglang",
)
else:
# Target: dp-attention, tp_rank=0 gets 0 tokens, tp_rank=1 gets all
_create_rank_dump(
side_dir,
rank=0,
name="layer_input",
tensor=torch.empty(0, 8),
dims="t h",
parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=0),
framework="sglang",
)
_create_rank_dump(
side_dir,
rank=1,
name="layer_input",
tensor=data,
dims="t h",
parallel_info=self._sglang_dp_attn_parallel_info(tp_rank=1),
framework="sglang",
)
argv: list[str] = _make_argv(
tmp_path / "baseline" / _FIXED_EXP_NAME,
tmp_path / "target" / _FIXED_EXP_NAME,
diff_threshold=1e-3,
)
records, exit_code = _run_and_parse(argv, capsys)
assert exit_code == 1
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
assert len(comparisons) == 1
comparison: TensorComparisonRecord = comparisons[0]
assert comparison.shape_mismatch is True
assert comparison.diff is None
assert comparison.category == "failed"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

2
test/registered/debug_utils/comparator/test_log_sink.py
View File

@@ -7,8 +7,8 @@ from sglang.srt.debug_utils.comparator.log_sink import LogSink
from sglang.srt.debug_utils.comparator.output_types import (
ErrorLog,
InfoLog,
report_sink,
)
from sglang.srt.debug_utils.comparator.report_sink import report_sink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)

567
test/registered/debug_utils/comparator/test_output_types.py
View File

@@ -1,18 +1,67 @@
import sys
from io import StringIO
import pytest
from registered.debug_utils.comparator.testing_helpers import make_diff as _make_diff
from registered.debug_utils.comparator.testing_helpers import (
make_tensor_info as _make_tensor_info,
)
from rich.console import Console
from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSidePlan,
TracedStepPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
AlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
ReordererPlan,
ZigzagToNaturalParams,
)
from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import (
TokenAlignerPlan,
TokenLocator,
)
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
ConcatParams,
UnsharderPlan,
)
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
RecordLocation,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_format_aligner_plan,
_split_logs,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _render_rich(renderable: object) -> str:
buf: StringIO = StringIO()
Console(file=buf, force_terminal=False, width=120).print(renderable)
return buf.getvalue().rstrip("\n")
# ---------------------------------------------------------------------------
# Existing tests (preserved)
# ---------------------------------------------------------------------------
def test_split_logs_mixed_list() -> None:
"""_split_logs correctly partitions a mixed list of ErrorLog and InfoLog."""
errors, infos = _split_logs(
@@ -40,5 +89,523 @@ def test_log_record_to_text_format() -> None:
assert "ℹ fyi" in text
class TestLogRecord:
def test_format_body_returns_empty(self) -> None:
record: LogRecord = LogRecord()
assert record._format_body() == ""
def test_format_rich_body_returns_empty(self) -> None:
record: LogRecord = LogRecord()
assert record._format_rich_body() == ""
def test_to_text_empty_no_logs(self) -> None:
record: LogRecord = LogRecord()
assert record.to_text() == ""
def test_to_text_with_errors_and_infos(self) -> None:
record: LogRecord = LogRecord(
errors=[ErrorLog(category="a", message="bad thing")],
infos=[InfoLog(category="b", message="fyi")],
)
text: str = record.to_text()
assert text == "\n ✗ bad thing\n ℹ fyi"
# ---------------------------------------------------------------------------
# ConfigRecord
# ---------------------------------------------------------------------------
class TestConfigRecord:
def test_format_body(self) -> None:
record: ConfigRecord = ConfigRecord(config={"a": 1, "b": "two"})
assert record._format_body() == "Config: {'a': 1, 'b': 'two'}"
def test_to_text_with_errors(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"x": 1},
errors=[ErrorLog(category="cfg", message="bad config")],
)
text: str = record.to_text()
assert text.startswith("Config: {'x': 1}")
assert "✗ bad config" in text
# ---------------------------------------------------------------------------
# SkipComparisonRecord
# ---------------------------------------------------------------------------
class TestSkipComparisonRecord:
def test_format_body_no_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="layer.weight",
reason="zero-dim tensor",
)
assert record._format_body() == "Skip: layer.weight (zero-dim tensor)"
def test_format_body_with_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="layer.weight",
reason="scalar",
location=RecordLocation(step=3),
)
assert record._format_body() == "Skip: layer.weight (step=3) (scalar)"
def test_category_skipped(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="x",
reason="r",
)
assert record.category == "skipped"
def test_category_failed(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
name="x",
reason="r",
errors=[ErrorLog(category="e", message="boom")],
)
assert record.category == "failed"
# ---------------------------------------------------------------------------
# NonTensorComparisonRecord
# ---------------------------------------------------------------------------
class TestNonTensorComparisonRecord:
def test_format_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.001",
baseline_type="float",
target_type="float",
values_equal=True,
)
assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]"
def test_format_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.01",
baseline_type="float",
target_type="float",
values_equal=False,
)
assert record._format_body() == (
"NonTensor: config.lr\n"
" baseline = 0.001 (float)\n"
" target = 0.01 (float)"
)
def test_with_step(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="bias",
baseline_value="True",
target_value="True",
baseline_type="bool",
target_type="bool",
values_equal=True,
location=RecordLocation(step=5),
)
assert "(step=5)" in record._format_body()
def test_category(self) -> None:
passed: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="x",
baseline_value="1",
target_value="1",
baseline_type="int",
target_type="int",
values_equal=True,
)
failed: NonTensorComparisonRecord = NonTensorComparisonRecord(
name="x",
baseline_value="1",
target_value="2",
baseline_type="int",
target_type="int",
values_equal=False,
)
assert passed.category == "passed"
assert failed.category == "failed"
# ---------------------------------------------------------------------------
# SummaryRecord
# ---------------------------------------------------------------------------
class TestSummaryRecord:
def test_format_body(self) -> None:
record: SummaryRecord = SummaryRecord(
total=10,
passed=7,
failed=2,
skipped=1,
)
assert record._format_body() == (
"Summary: 7 passed, 2 failed, 1 skipped (total 10)"
)
def test_validation_error(self) -> None:
with pytest.raises(ValueError, match="total=5 !="):
SummaryRecord(total=5, passed=1, failed=1, skipped=1)
# ---------------------------------------------------------------------------
# TensorComparisonRecord._format_body
# ---------------------------------------------------------------------------
class TestTensorComparisonRecordFormatBody:
def test_basic(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_with_replicated_checks(self) -> None:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
replicated_checks=[
ReplicatedCheckResult(
axis="tp",
group_index=0,
compared_index=1,
baseline_index=0,
passed=True,
atol=1e-3,
diff=_make_diff(
rel_diff=1e-6, max_abs_diff=1e-5, mean_abs_diff=1e-6
),
),
],
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
"Replicated checks:\n"
" ✅ axis=tp group=0 idx=1 vs 0: "
"rel_diff=1.000000e-06 max_abs_diff=1.000000e-05 mean_abs_diff=1.000000e-06"
)
def test_with_aligner_plan(self) -> None:
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
)
traced: TracedAlignerPlan = TracedAlignerPlan(
plan=plan,
per_side=Pair(
x=TracedSidePlan(step_plans=[]),
y=TracedSidePlan(step_plans=[]),
),
)
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
traced_plan=traced,
)
body: str = record._format_body()
assert body == (
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005\n"
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)"
)
def test_with_step(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
location=RecordLocation(step=2),
)
body: str = record._format_body()
assert body == (
"[step=2] "
"Raw [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p50] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
# ---------------------------------------------------------------------------
# _format_aligner_plan
# ---------------------------------------------------------------------------
def _wrap_plan(plan: AlignerPlan) -> TracedAlignerPlan:
"""Wrap an AlignerPlan into a TracedAlignerPlan with no snapshots."""
baseline_traced_steps: list[TracedStepPlan] = [
TracedStepPlan(
step=sp.step,
input_object_indices=sp.input_object_indices,
sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans],
)
for sp in plan.per_step_plans.x
]
target_traced_steps: list[TracedStepPlan] = [
TracedStepPlan(
step=sp.step,
input_object_indices=sp.input_object_indices,
sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans],
)
for sp in plan.per_step_plans.y
]
return TracedAlignerPlan(
plan=plan,
per_side=Pair(
x=TracedSidePlan(step_plans=baseline_traced_steps),
y=TracedSidePlan(step_plans=target_traced_steps),
),
)
class TestFormatAlignerPlan:
def test_passthrough(self) -> None:
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: (no steps)"
)
def test_unsharder(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
)
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: unsharder]"
)
def test_reorderer(self) -> None:
reorderer: ReordererPlan = ReordererPlan(
params=ZigzagToNaturalParams(dim_name="s", cp_size=2),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0], sub_plans=[reorderer]
)
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: reorderer]"
)
def test_multi_step(self) -> None:
unsharder: UnsharderPlan = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
reorderer: ReordererPlan = ReordererPlan(
params=ZigzagToNaturalParams(dim_name="s", cp_size=2),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(
x=[],
y=[
AlignerPerStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
),
AlignerPerStepPlan(
step=1, input_object_indices=[0], sub_plans=[reorderer]
),
],
),
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder; step=1: reorderer]"
)
def test_with_token_aligner(self) -> None:
ta_plan: TokenAlignerPlan = TokenAlignerPlan(
locators=Pair(
x=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]),
y=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]),
),
layouts=Pair(x=TokenLayout.T, y=TokenLayout.T),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
token_aligner_plan=ta_plan,
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)\n"
" token_aligner: 3 tokens aligned"
)
def test_with_axis_aligner(self) -> None:
aa_plan: AxisAlignerPlan = AxisAlignerPlan(
pattern=Pair(x="b s d -> s b d", y=None),
)
plan: AlignerPlan = AlignerPlan(
per_step_plans=Pair(x=[], y=[]),
axis_aligner_plan=aa_plan,
)
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: (no steps)\n"
" axis_aligner: x: b s d -> s b d"
)
# ---------------------------------------------------------------------------
# _OutputRecord log attachment (to_text)
# ---------------------------------------------------------------------------
class TestOutputRecordLogAttachment:
def test_to_text_no_logs(self) -> None:
record: ConfigRecord = ConfigRecord(config={"a": 1})
text: str = record.to_text()
assert text == "Config: {'a': 1}"
def test_to_text_errors_only(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
errors=[ErrorLog(category="x", message="err1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n ✗ err1"
def test_to_text_infos_only(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
infos=[InfoLog(category="x", message="note1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n ℹ note1"
def test_to_text_mixed(self) -> None:
record: ConfigRecord = ConfigRecord(
config={"a": 1},
errors=[ErrorLog(category="x", message="err1")],
infos=[InfoLog(category="y", message="note1")],
)
text: str = record.to_text()
assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

84
test/registered/debug_utils/comparator/testing_helpers.py Normal file
View File

@@ -0,0 +1,84 @@
"""Shared test helpers for comparator tests."""
from __future__ import annotations
from typing import Optional
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorInfo,
TensorStats,
)
DEFAULT_PERCENTILES: dict[int, float] = {
1: -1.8,
5: -1.5,
50: 0.0,
95: 1.5,
99: 1.8,
}
DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = {
1: 0.0001,
5: 0.0001,
50: 0.0002,
95: 0.0004,
99: 0.0005,
}
def make_stats(
mean: float = 0.0,
abs_mean: float = 0.8,
std: float = 1.0,
min: float = -2.0,
max: float = 2.0,
percentiles: Optional[dict[int, float]] = None,
) -> TensorStats:
return TensorStats(
mean=mean,
abs_mean=abs_mean,
std=std,
min=min,
max=max,
percentiles=percentiles if percentiles is not None else DEFAULT_PERCENTILES,
)
def make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
abs_diff_percentiles: Optional[dict[int, float]] = None,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
abs_diff_percentiles=(
abs_diff_percentiles
if abs_diff_percentiles is not None
else DEFAULT_ABS_DIFF_PERCENTILES
),
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
diff_threshold=diff_threshold,
passed=passed,
)
def make_tensor_info(
shape: Optional[list[int]] = None,
dtype: str = "torch.float32",
stats: Optional[TensorStats] = None,
sample: Optional[str] = None,
) -> TensorInfo:
return TensorInfo(
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else make_stats(),
sample=sample,
)