From 67810828cf6014ce0586272633931aeba53eb980 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 1 Mar 2026 10:32:59 +0800 Subject: [PATCH] Visualize per-token information in dump comparator (#19594) --- .../comparator/bundle_comparator.py | 26 +++ .../srt/debug_utils/comparator/entrypoint.py | 30 +++- .../comparator/per_token_visualizer.py | 83 +++++++++ .../tensor_comparator/comparator.py | 11 ++ .../comparator/tensor_comparator/types.py | 1 + .../srt/debug_utils/comparator/utils.py | 20 +++ .../tensor_comparator/test_comparator.py | 42 +++++ .../debug_utils/comparator/test_entrypoint.py | 68 ++++++++ .../comparator/test_manually_verify.py | 90 +++++++++- .../comparator/test_per_token_visualizer.py | 159 ++++++++++++++++++ .../debug_utils/comparator/test_utils.py | 47 ++++++ 11 files changed, 575 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/per_token_visualizer.py create mode 100644 test/registered/debug_utils/comparator/test_per_token_visualizer.py diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index a14936c1f..940003e7b 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -19,6 +19,8 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import ( TokenAlignerPlan, ) from sglang.srt.debug_utils.comparator.dims import ( + SEQ_DIM_NAME, + TOKEN_DIM_NAME, apply_dim_names, resolve_dim_names, ) @@ -50,6 +52,7 @@ def compare_bundle_pair( x=None, y=None ), viz_output_dir: Optional[Path] = None, + compute_per_token: bool = False, ) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]: with warning_sink.context() as collected_warnings: result = _compare_bundle_pair_inner( @@ -61,6 +64,7 @@ def compare_bundle_pair( diff_threshold=diff_threshold, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, + compute_per_token=compute_per_token, ) return result.model_copy(update={"warnings": collected_warnings}) @@ -78,6 +82,7 @@ def _compare_bundle_pair_inner( x=None, y=None ), viz_output_dir: Optional[Path] = None, + compute_per_token: bool = False, ) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]: # 1. Load all successfully loaded values all_pair: Pair[list[ValueWithMeta]] = Pair( @@ -104,6 +109,7 @@ def _compare_bundle_pair_inner( diff_threshold=diff_threshold, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, + compute_per_token=compute_per_token, ) @@ -117,6 +123,7 @@ def _compare_bundle_pair_tensor_type( x=None, y=None ), viz_output_dir: Optional[Path] = None, + compute_per_token: bool = False, ) -> Union[ComparisonRecord, SkipRecord]: if not valid_pair.x or not valid_pair.y: reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed" @@ -153,6 +160,11 @@ def _compare_bundle_pair_tensor_type( reason: str = f"{side_name}_load_failed" return SkipRecord(name=name, reason=reason) + # Resolve seq_dim for per-token computation + seq_dim: Optional[int] = ( + _resolve_seq_dim(aligner_result.tensors.y) if compute_per_token else None + ) + # Compare aligned_baseline: torch.Tensor = aligner_result.tensors.x.rename(None) aligned_target: torch.Tensor = aligner_result.tensors.y.rename(None) @@ -162,6 +174,7 @@ def _compare_bundle_pair_tensor_type( x_target=aligned_target, name=name, diff_threshold=diff_threshold, + seq_dim=seq_dim, ) record = ComparisonRecord(**info.model_dump(), aligner_plan=plan) @@ -209,6 +222,19 @@ def _try_generate_viz( ) +def _resolve_seq_dim(tensor: torch.Tensor) -> Optional[int]: + """Find the token/seq dimension index from the tensor's named dims.""" + if tensor.names[0] is None: + return None + + names: tuple[Optional[str], ...] = tensor.names + for target_name in (TOKEN_DIM_NAME, SEQ_DIM_NAME): + if target_name in names: + return list(names).index(target_name) + + return None + + def _compare_bundle_pair_non_tensor_type( *, name: str, diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 2587a9566..840f79da6 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -30,6 +30,9 @@ from sglang.srt.debug_utils.comparator.output_types import ( SummaryRecord, print_record, ) +from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, +) from sglang.srt.debug_utils.comparator.utils import Pair from sglang.srt.debug_utils.comparator.warning_sink import warning_sink from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path @@ -78,6 +81,10 @@ def run(args: argparse.Namespace) -> None: Path(args.viz_output_dir) if args.viz_bundle_details else None ) + visualize_per_token: Optional[Path] = ( + Path(args.visualize_per_token) if args.visualize_per_token else None + ) + comparison_records = _compare_bundle_pairs( bundle_info_pairs=bundle_info_pairs, baseline_path=Path(args.baseline_path), @@ -86,9 +93,12 @@ def run(args: argparse.Namespace) -> None: diff_threshold=args.diff_threshold, thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, + compute_per_token=visualize_per_token is not None, ) _consume_comparison_records( - comparison_records=comparison_records, output_format=args.output_format + comparison_records=comparison_records, + output_format=args.output_format, + visualize_per_token=visualize_per_token, ) @@ -144,6 +154,7 @@ def _compare_bundle_pairs( diff_threshold: float, thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]], viz_output_dir: Optional[Path] = None, + compute_per_token: bool = False, ) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]: for bundle_info_pair in bundle_info_pairs: if not bundle_info_pair.y: @@ -162,6 +173,7 @@ def _compare_bundle_pairs( diff_threshold=diff_threshold, thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, viz_output_dir=viz_output_dir, + compute_per_token=compute_per_token, ) @@ -169,18 +181,28 @@ def _consume_comparison_records( *, comparison_records: Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]], output_format: str, + visualize_per_token: Optional[Path] = None, ) -> None: counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} + collected_comparisons: list[ComparisonRecord] = [] for record in comparison_records: counts[record.category] += 1 print_record(record, output_format=output_format) + if visualize_per_token is not None and isinstance(record, ComparisonRecord): + collected_comparisons.append(record) print_record( SummaryRecord(total=sum(counts.values()), **counts), output_format=output_format, ) + if visualize_per_token is not None and collected_comparisons: + generate_per_token_heatmap( + records=collected_comparisons, + output_path=visualize_per_token, + ) + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() @@ -224,4 +246,10 @@ def _parse_args() -> argparse.Namespace: default="/tmp/comparator_viz/", help="Output directory for visualization PNGs (default: /tmp/comparator_viz/)", ) + parser.add_argument( + "--visualize-per-token", + type=str, + default=None, + help="Output path for per-token relative difference heatmap PNG", + ) return parser.parse_args() diff --git a/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py new file mode 100644 index 000000000..899b50777 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py @@ -0,0 +1,83 @@ +"""Per-token relative difference heatmap generator. + +Produces a single PNG with rows = tensor names, columns = token positions, +color = log10(rel_diff). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Optional + +from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord + + +def generate_per_token_heatmap( + *, + records: list[ComparisonRecord], + output_path: Path, +) -> Optional[Path]: + """Generate a per-token relative difference heatmap PNG. + + Returns the output path if a file was written, or None if no data was available. + """ + rows_data: list[tuple[str, list[float]]] = _collect_per_token_data(records=records) + if not rows_data: + return None + + _render_heatmap(rows_data=rows_data, output_path=output_path) + return output_path + + +def _collect_per_token_data( + *, + records: list[ComparisonRecord], +) -> list[tuple[str, list[float]]]: + rows: list[tuple[str, list[float]]] = [] + for record in records: + if record.diff is None or record.diff.per_token_rel_diff is None: + continue + rows.append((record.name, record.diff.per_token_rel_diff)) + return rows + + +def _render_heatmap( + *, + rows_data: list[tuple[str, list[float]]], + output_path: Path, +) -> None: + import matplotlib + import numpy as np + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + max_len: int = max(len(vals) for _, vals in rows_data) + labels: list[str] = [label for label, _ in rows_data] + + matrix: np.ndarray = np.full((len(rows_data), max_len), np.nan, dtype=np.float64) + for i, (_, vals) in enumerate(rows_data): + matrix[i, : len(vals)] = vals + + fig_width: float = max(12.0, max_len * 0.15) + fig_height: float = max(6.0, len(rows_data) * 0.3) + fig, ax = plt.subplots(figsize=(fig_width, fig_height)) + + im = ax.imshow( + np.log10(matrix + 1e-10), aspect="auto", cmap="hot", interpolation="nearest" + ) + + ax.set_xlabel("Token Position") + ax.set_ylabel("Tensor") + ax.set_yticks(range(len(labels))) + ax.set_yticklabels(labels, fontsize=8) + + colorbar = fig.colorbar(im, ax=ax) + colorbar.set_label("log10(rel_diff)") + + ax.set_title("Per-Token Relative Difference Heatmap") + fig.tight_layout() + + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(str(output_path), dpi=150) + plt.close(fig) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py index f6cff930b..46016bfbd 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -12,6 +12,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( from sglang.srt.debug_utils.comparator.utils import ( Pair, argmax_coord, + calc_per_token_rel_diff, calc_rel_diff, compute_smaller_dtype, try_unify_shape, @@ -27,6 +28,7 @@ def compare_tensor_pair( x_target: torch.Tensor, name: str = "", diff_threshold: float = 1e-3, + seq_dim: Optional[int] = None, ) -> TensorComparisonInfo: baseline_info = TensorInfo( shape=list(x_baseline.shape), @@ -59,6 +61,7 @@ def compare_tensor_pair( x_baseline=x_baseline_f, x_target=x_target_f, diff_threshold=diff_threshold, + seq_dim=seq_dim, ) needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD @@ -122,6 +125,7 @@ def _compute_diff( x_baseline: torch.Tensor, x_target: torch.Tensor, diff_threshold: float = 1e-3, + seq_dim: Optional[int] = None, ) -> DiffInfo: if x_baseline.numel() == 0: return DiffInfo( @@ -145,6 +149,12 @@ def _compute_diff( include_quantiles: bool = raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD + per_token_rel_diff: Optional[list[float]] = None + if seq_dim is not None and x_baseline.dim() > seq_dim: + per_token_rel_diff = calc_per_token_rel_diff( + x_baseline, x_target, seq_dim=seq_dim + ).tolist() + return DiffInfo( rel_diff=rel_diff, max_abs_diff=max_abs_diff, @@ -157,4 +167,5 @@ def _compute_diff( target_at_max=x_target[max_diff_coord].item(), diff_threshold=diff_threshold, passed=rel_diff <= diff_threshold, + per_token_rel_diff=per_token_rel_diff, ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py index 89f7fd2ba..e505d022e 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py @@ -31,6 +31,7 @@ class DiffInfo(_StrictBase): target_at_max: float diff_threshold: float passed: bool + per_token_rel_diff: Optional[list[float]] = None class TensorComparisonInfo(_StrictBase): diff --git a/python/sglang/srt/debug_utils/comparator/utils.py b/python/sglang/srt/debug_utils/comparator/utils.py index c899d1c95..a6f30fe49 100644 --- a/python/sglang/srt/debug_utils/comparator/utils.py +++ b/python/sglang/srt/debug_utils/comparator/utils.py @@ -66,3 +66,23 @@ def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: denominator = (x * x + y * y).sum() sim = 2 * (x * y).sum() / denominator return 1 - sim + + +def calc_per_token_rel_diff( + x: torch.Tensor, y: torch.Tensor, *, seq_dim: int +) -> torch.Tensor: + """Cosine-distance-like metric per token position. + + Sums over all dims except seq_dim. + """ + x, y = x.double(), y.double() + other_dims: list[int] = [d for d in range(x.dim()) if d != seq_dim] + + if other_dims: + denominator: torch.Tensor = (x * x + y * y).sum(dim=other_dims) + sim: torch.Tensor = 2 * (x * y).sum(dim=other_dims) / (denominator + 1e-10) + else: + denominator = x * x + y * y + sim = 2 * (x * y) / (denominator + 1e-10) + + return (1 - sim).float() diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py index 0943eaf98..9851b064d 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py @@ -10,6 +10,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( _compute_tensor_stats, compare_tensor_pair, ) +from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=20, suite="default", nightly=True) @@ -96,6 +97,47 @@ class TestComputeDiff: assert diff.rel_diff == pytest.approx(1.0, abs=1e-5) assert diff.passed is False + def test_per_token_with_seq_dim(self) -> None: + """seq_dim provided → per_token_rel_diff is list[float].""" + torch.manual_seed(42) + x: torch.Tensor = torch.randn(8, 16) + y: torch.Tensor = x + torch.randn_like(x) * 0.01 + + diff: DiffInfo = _compute_diff( + x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0 + ) + + assert diff.per_token_rel_diff is not None + assert isinstance(diff.per_token_rel_diff, list) + assert len(diff.per_token_rel_diff) == 8 + assert all(isinstance(v, float) for v in diff.per_token_rel_diff) + + def test_per_token_without_seq_dim(self) -> None: + """No seq_dim → per_token_rel_diff is None.""" + x: torch.Tensor = torch.randn(8, 16) + y: torch.Tensor = x + torch.randn_like(x) * 0.01 + + diff: DiffInfo = _compute_diff(x_baseline=x, x_target=y, diff_threshold=1e-3) + + assert diff.per_token_rel_diff is None + + def test_per_token_json_roundtrip(self) -> None: + """DiffInfo with per_token_rel_diff survives JSON serialization.""" + torch.manual_seed(42) + x: torch.Tensor = torch.randn(4, 8) + y: torch.Tensor = x + torch.randn_like(x) * 0.01 + + diff: DiffInfo = _compute_diff( + x_baseline=x, x_target=y, diff_threshold=1e-3, seq_dim=0 + ) + + json_str: str = diff.model_dump_json() + assert "per_token_rel_diff" in json_str + + roundtripped: DiffInfo = DiffInfo.model_validate_json(json_str) + assert roundtripped.per_token_rel_diff is not None + assert len(roundtripped.per_token_rel_diff) == 4 + class TestCompareTensors: def test_normal(self): diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 0be688b21..028bb5a40 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -1785,6 +1785,7 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace grouping="logical", viz_bundle_details=False, viz_output_dir="/tmp/comparator_viz/", + visualize_per_token=None, ) defaults.update(overrides) return Namespace(**defaults) @@ -2180,6 +2181,73 @@ def _create_thd_cp_zigzag_dumps( return directory / _FIXED_EXP_NAME +class TestEntrypointPerTokenVisualization: + """Test --visualize-per-token CLI flag integration.""" + + def test_visualize_per_token_creates_png(self, tmp_path: Path, capsys) -> None: + """--visualize-per-token with dims metadata produces per-token data in records.""" + pytest.importorskip("matplotlib") + + torch.manual_seed(42) + baseline_dir: Path = tmp_path / "baseline" + target_dir: Path = tmp_path / "target" + baseline_dir.mkdir() + target_dir.mkdir() + + baseline_tensor: torch.Tensor = torch.randn(10, 10) + target_tensor: torch.Tensor = baseline_tensor + torch.randn(10, 10) * 0.01 + + for name in ["tensor_a", "tensor_b"]: + _create_rank_dump( + baseline_dir, + rank=0, + name=name, + tensor=baseline_tensor, + dims="t h", + ) + _create_rank_dump( + target_dir, + rank=0, + name=name, + tensor=target_tensor, + dims="t h", + ) + + baseline_path: Path = baseline_dir / _FIXED_EXP_NAME + target_path: Path = target_dir / _FIXED_EXP_NAME + + output_png: Path = tmp_path / "per_token.png" + args = _make_args( + baseline_path, + target_path, + grouping="raw", + visualize_per_token=str(output_png), + ) + records = _run_and_parse(args, capsys) + + comparisons = _get_comparisons(records) + assert len(comparisons) == 2 + + # per_token_rel_diff should be populated + for comp in comparisons: + assert comp.diff is not None + assert comp.diff.per_token_rel_diff is not None + assert isinstance(comp.diff.per_token_rel_diff, list) + assert len(comp.diff.per_token_rel_diff) == 10 + + def test_no_visualize_no_per_token(self, tmp_path: Path, capsys) -> None: + """Without --visualize-per-token, per_token_rel_diff is None.""" + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) + args = _make_args(baseline_path, target_path, grouping="raw") + + records = _run_and_parse(args, capsys) + + comparisons = _get_comparisons(records) + assert len(comparisons) == 1 + assert comparisons[0].diff is not None + assert comparisons[0].diff.per_token_rel_diff is None + + class TestEntrypointThdCpZigzag: """E2E entrypoint tests for THD CP zigzag format. diff --git a/test/registered/debug_utils/comparator/test_manually_verify.py b/test/registered/debug_utils/comparator/test_manually_verify.py index 0c2b822a1..51a9ef0f8 100644 --- a/test/registered/debug_utils/comparator/test_manually_verify.py +++ b/test/registered/debug_utils/comparator/test_manually_verify.py @@ -76,7 +76,7 @@ def _skip_if_no_matplotlib() -> None: pytest.importorskip("matplotlib") -class TestManuallyVerify: +class TestBundleDetailsManualVerify: def test_normal_small_diff(self, tmp_path: Path, publish_dir: Path) -> None: """Two nearly-identical tensors (randn + 0.01 noise). @@ -199,5 +199,93 @@ class TestManuallyVerify: ) +class TestPerTokenHeatmapManualVerify: + def test_increasing_diff(self, tmp_path: Path, publish_dir: Path) -> None: + """Per-token heatmap with linearly increasing diff across token positions. + + Expected: Heatmap shows a clear left-to-right gradient — dark/cold on + the left (small diff), bright/hot on the right (large diff). Multiple + rows for different tensor names. Colorbar shows log10 scale. + """ + from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( + compare_tensor_pair, + ) + + torch.manual_seed(42) + seq_len: int = 64 + hidden_dim: int = 128 + num_tensors: int = 5 + + records: list[ComparisonRecord] = [] + for i in range(num_tensors): + baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) + noise_scale: torch.Tensor = torch.linspace( + 1e-6, 0.5, steps=seq_len + ).unsqueeze(1) + target: torch.Tensor = baseline + torch.randn_like(baseline) * noise_scale + + info = compare_tensor_pair( + x_baseline=baseline, + x_target=target, + name=f"layer_{i}_hidden_states", + diff_threshold=1e-3, + seq_dim=0, + ) + records.append(ComparisonRecord(**info.model_dump())) + + output_path: Path = tmp_path / "per_token_increasing_diff.png" + result = generate_per_token_heatmap(records=records, output_path=output_path) + + assert result is not None + _assert_valid_png(output_path) + shutil.copy2(src=output_path, dst=publish_dir / output_path.name) + + def test_single_spike(self, tmp_path: Path, publish_dir: Path) -> None: + """Per-token heatmap where only one token position has large diff. + + Expected: Heatmap shows one bright vertical stripe at the spike position, + rest is dark/cold. + """ + from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( + compare_tensor_pair, + ) + + torch.manual_seed(42) + seq_len: int = 64 + hidden_dim: int = 128 + spike_pos: int = 32 + num_tensors: int = 4 + + records: list[ComparisonRecord] = [] + for i in range(num_tensors): + baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) + target: torch.Tensor = baseline.clone() + target[spike_pos, :] += torch.randn(hidden_dim) * 5.0 + + info = compare_tensor_pair( + x_baseline=baseline, + x_target=target, + name=f"layer_{i}_attn_output", + diff_threshold=1e-3, + seq_dim=0, + ) + records.append(ComparisonRecord(**info.model_dump())) + + output_path: Path = tmp_path / "per_token_single_spike.png" + result = generate_per_token_heatmap(records=records, output_path=output_path) + + assert result is not None + _assert_valid_png(output_path) + shutil.copy2(src=output_path, dst=publish_dir / output_path.name) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_per_token_visualizer.py b/test/registered/debug_utils/comparator/test_per_token_visualizer.py new file mode 100644 index 000000000..3f0e0e0ad --- /dev/null +++ b/test/registered/debug_utils/comparator/test_per_token_visualizer.py @@ -0,0 +1,159 @@ +"""Layer 2: PNG generation tests for per-token heatmap visualizer. + +Requires matplotlib — uses pytest.importorskip to gracefully skip if absent. +""" + +import sys +from pathlib import Path + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord +from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( + compare_tensor_pair, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="default", nightly=True) + +_PNG_MAGIC: bytes = b"\x89PNG" + + +@pytest.fixture(autouse=True) +def _skip_if_no_matplotlib() -> None: + pytest.importorskip("matplotlib") + + +def _make_comparison_record( + *, + name: str, + baseline: torch.Tensor, + target: torch.Tensor, + seq_dim: int = 0, +) -> ComparisonRecord: + """Build a ComparisonRecord with per-token data from raw tensors.""" + info = compare_tensor_pair( + x_baseline=baseline, + x_target=target, + name=name, + diff_threshold=1e-3, + seq_dim=seq_dim, + ) + return ComparisonRecord(**info.model_dump()) + + +class TestPerTokenVisualizer: + def test_no_data_returns_none(self, tmp_path: Path) -> None: + """Empty records list → None returned, no file created.""" + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + + output_path: Path = tmp_path / "empty.png" + result = generate_per_token_heatmap(records=[], output_path=output_path) + + assert result is None + assert not output_path.exists() + + def test_no_per_token_data_returns_none(self, tmp_path: Path) -> None: + """Records without per_token_rel_diff → None.""" + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + + info = compare_tensor_pair( + x_baseline=torch.randn(4, 8), + x_target=torch.randn(4, 8), + name="no_per_token", + diff_threshold=1e-3, + ) + record = ComparisonRecord(**info.model_dump()) + + output_path: Path = tmp_path / "no_data.png" + result = generate_per_token_heatmap(records=[record], output_path=output_path) + + assert result is None + + def test_generates_valid_png(self, tmp_path: Path) -> None: + """Records with per-token data → valid PNG file.""" + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + + torch.manual_seed(42) + records: list[ComparisonRecord] = [ + _make_comparison_record( + name=f"tensor_{i}", + baseline=torch.randn(16, 32), + target=torch.randn(16, 32), + ) + for i in range(3) + ] + + output_path: Path = tmp_path / "heatmap.png" + result = generate_per_token_heatmap(records=records, output_path=output_path) + + assert result == output_path + assert output_path.exists() + assert output_path.stat().st_size > 0 + with open(output_path, "rb") as f: + magic: bytes = f.read(4) + assert magic == _PNG_MAGIC + + def test_variable_length_sequences(self, tmp_path: Path) -> None: + """Records with different token lengths → NaN padding, no crash.""" + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + + torch.manual_seed(42) + records: list[ComparisonRecord] = [ + _make_comparison_record( + name="short", + baseline=torch.randn(4, 8), + target=torch.randn(4, 8), + ), + _make_comparison_record( + name="medium", + baseline=torch.randn(16, 8), + target=torch.randn(16, 8), + ), + _make_comparison_record( + name="long", + baseline=torch.randn(64, 8), + target=torch.randn(64, 8), + ), + ] + + output_path: Path = tmp_path / "variable.png" + result = generate_per_token_heatmap(records=records, output_path=output_path) + + assert result == output_path + assert output_path.exists() + with open(output_path, "rb") as f: + magic: bytes = f.read(4) + assert magic == _PNG_MAGIC + + def test_creates_parent_dirs(self, tmp_path: Path) -> None: + """Output path with non-existent parent dirs → dirs created automatically.""" + from sglang.srt.debug_utils.comparator.per_token_visualizer import ( + generate_per_token_heatmap, + ) + + torch.manual_seed(42) + record = _make_comparison_record( + name="test", + baseline=torch.randn(8, 16), + target=torch.randn(8, 16), + ) + + output_path: Path = tmp_path / "nested" / "deep" / "heatmap.png" + result = generate_per_token_heatmap(records=[record], output_path=output_path) + + assert result == output_path + assert output_path.exists() + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_utils.py b/test/registered/debug_utils/comparator/test_utils.py index 89a2b9856..81ea7c2fc 100644 --- a/test/registered/debug_utils/comparator/test_utils.py +++ b/test/registered/debug_utils/comparator/test_utils.py @@ -6,6 +6,7 @@ import torch from sglang.srt.debug_utils.comparator.utils import ( Pair, argmax_coord, + calc_per_token_rel_diff, calc_rel_diff, compute_smaller_dtype, try_unify_shape, @@ -38,6 +39,52 @@ class TestCalcRelDiff: assert result == pytest.approx(2.0, abs=1e-5) +class TestCalcPerTokenRelDiff: + def test_identical_tensors(self) -> None: + """Identical tensors → per-token diff all zero.""" + x: torch.Tensor = torch.randn(8, 16) + result: torch.Tensor = calc_per_token_rel_diff(x, x, seq_dim=0) + + assert result.shape == (8,) + assert torch.allclose(result, torch.zeros(8), atol=1e-6) + + def test_different_tensors(self) -> None: + """Single token position differs → that position has higher diff.""" + torch.manual_seed(42) + x: torch.Tensor = torch.randn(8, 16) + y: torch.Tensor = x.clone() + y[3, :] += 10.0 + + result: torch.Tensor = calc_per_token_rel_diff(x, y, seq_dim=0) + + assert result.shape == (8,) + assert result[3] > result[0] + assert result[3] > result[7] + for i in [0, 1, 2, 4, 5, 6, 7]: + assert result[i] < 1e-6 + + def test_seq_dim_selection(self) -> None: + """Different seq_dim values produce correct output shapes.""" + x: torch.Tensor = torch.randn(4, 8, 16) + y: torch.Tensor = x + torch.randn_like(x) * 0.01 + + assert calc_per_token_rel_diff(x, y, seq_dim=0).shape == (4,) + assert calc_per_token_rel_diff(x, y, seq_dim=1).shape == (8,) + assert calc_per_token_rel_diff(x, y, seq_dim=2).shape == (16,) + + def test_1d_tensor(self) -> None: + """1D tensor with seq_dim=0 returns per-element diff.""" + x: torch.Tensor = torch.tensor([1.0, 2.0, 3.0]) + y: torch.Tensor = torch.tensor([1.0, 2.0, 4.0]) + + result: torch.Tensor = calc_per_token_rel_diff(x, y, seq_dim=0) + + assert result.shape == (3,) + assert result[0] < 1e-6 + assert result[1] < 1e-6 + assert result[2] > 0.01 + + class TestArgmaxCoord: def test_1d_tensor(self): x = torch.tensor([0.0, 0.0, 5.0, 0.0])