Visualize per-token information in dump comparator (#19594)

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,

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()

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)

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,
     )

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):

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()