Visualize comparison detailed results in dump comparator (#19565)

python/sglang/srt/debug_utils/comparator/bundle_comparator.py

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
 from sglang.srt.debug_utils.comparator.dims import apply_dim_names, parse_dim_names
 from sglang.srt.debug_utils.comparator.output_types import (
     ComparisonRecord,
+    GeneralWarning,
     NonTensorRecord,
     SkipRecord,
 )
@@ -45,6 +46,7 @@ def compare_bundle_pair(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
+    viz_output_dir: Optional[Path] = None,
 ) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
     with warning_sink.context() as collected_warnings:
         result = _compare_bundle_pair_inner(
@@ -55,6 +57,7 @@ def compare_bundle_pair(
             token_aligner_plan=token_aligner_plan,
             diff_threshold=diff_threshold,
             thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+            viz_output_dir=viz_output_dir,
         )
 
     return result.model_copy(update={"warnings": collected_warnings})
@@ -71,6 +74,7 @@ def _compare_bundle_pair_inner(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
+    viz_output_dir: Optional[Path] = None,
 ) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
     # 1. Load all successfully loaded values
     all_pair: Pair[list[ValueWithMeta]] = Pair(
@@ -96,6 +100,7 @@ def _compare_bundle_pair_inner(
         token_aligner_plan=token_aligner_plan,
         diff_threshold=diff_threshold,
         thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+        viz_output_dir=viz_output_dir,
     )
 
 
@@ -108,6 +113,7 @@ def _compare_bundle_pair_tensor_type(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]] = Pair(
         x=None, y=None
     ),
+    viz_output_dir: Optional[Path] = None,
 ) -> 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"
@@ -145,13 +151,59 @@ def _compare_bundle_pair_tensor_type(
         return SkipRecord(name=name, reason=reason)
 
     # Compare
+    aligned_baseline: torch.Tensor = aligner_result.tensors.x.rename(None)
+    aligned_target: torch.Tensor = aligner_result.tensors.y.rename(None)
+
     info = compare_tensor_pair(
-        x_baseline=aligner_result.tensors.x.rename(None),
-        x_target=aligner_result.tensors.y.rename(None),
+        x_baseline=aligned_baseline,
+        x_target=aligned_target,
         name=name,
         diff_threshold=diff_threshold,
     )
-    return ComparisonRecord(**info.model_dump(), aligner_plan=plan)
+    record = ComparisonRecord(**info.model_dump(), aligner_plan=plan)
+
+    if viz_output_dir is not None:
+        _try_generate_viz(
+            baseline=aligned_baseline,
+            target=aligned_target,
+            name=name,
+            viz_output_dir=viz_output_dir,
+        )
+
+    return record
+
+
+def _try_generate_viz(
+    *,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    name: str,
+    viz_output_dir: Path,
+) -> None:
+    from sglang.srt.debug_utils.comparator.visualizer import (
+        generate_comparison_figure,
+    )
+    from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+        _sanitize_filename,
+    )
+
+    filename: str = _sanitize_filename(name) + ".png"
+    output_path: Path = viz_output_dir / filename
+
+    try:
+        generate_comparison_figure(
+            baseline=baseline,
+            target=target,
+            name=name,
+            output_path=output_path,
+        )
+    except Exception as exc:
+        warning_sink.add(
+            GeneralWarning(
+                category="visualizer",
+                message=f"Visualization failed for {name}: {exc}",
+            )
+        )
 
 
 def _compare_bundle_pair_non_tensor_type(

python/sglang/srt/debug_utils/comparator/entrypoint.py

@@ -74,6 +74,10 @@ def run(args: argparse.Namespace) -> None:
         ),
     )
 
+    viz_output_dir: Optional[Path] = (
+        Path(args.viz_output_dir) if args.viz_bundle_details else None
+    )
+
     comparison_records = _compare_bundle_pairs(
         bundle_info_pairs=bundle_info_pairs,
         baseline_path=Path(args.baseline_path),
@@ -81,6 +85,7 @@ def run(args: argparse.Namespace) -> None:
         token_aligner_plan=ta_result.plan,
         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,
     )
     _consume_comparison_records(
         comparison_records=comparison_records, output_format=args.output_format
@@ -138,6 +143,7 @@ def _compare_bundle_pairs(
     token_aligner_plan: Optional[TokenAlignerPlan],
     diff_threshold: float,
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
+    viz_output_dir: Optional[Path] = None,
 ) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]:
     for bundle_info_pair in bundle_info_pairs:
         if not bundle_info_pair.y:
@@ -155,6 +161,7 @@ def _compare_bundle_pairs(
             token_aligner_plan=token_aligner_plan,
             diff_threshold=diff_threshold,
             thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
+            viz_output_dir=viz_output_dir,
         )
 
 
@@ -205,4 +212,16 @@ def _parse_args() -> argparse.Namespace:
         default=None,
         help="Tokenizer path for decoding input_ids (auto-discovered from dump metadata if not set)",
     )
+    parser.add_argument(
+        "--viz-bundle-details",
+        action="store_true",
+        default=False,
+        help="Generate comparison heatmap/histogram PNG for each compared tensor",
+    )
+    parser.add_argument(
+        "--viz-output-dir",
+        type=str,
+        default="/tmp/comparator_viz/",
+        help="Output directory for visualization PNGs (default: /tmp/comparator_viz/)",
+    )
     return parser.parse_args()

python/sglang/srt/debug_utils/comparator/visualizer/__init__.py

@@ -0,0 +1,3 @@
+from sglang.srt.debug_utils.comparator.visualizer.figure import (  # noqa: F401
+    generate_comparison_figure,
+)

python/sglang/srt/debug_utils/comparator/visualizer/figure.py

@@ -0,0 +1,116 @@
+"""Main orchestration logic for comparison figure generation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+    _preprocess_tensor,
+)
+
+
+@dataclass(frozen=True)
+class _PanelContext:
+    baseline_2d: torch.Tensor
+    target_2d: torch.Tensor
+    diff: Optional[torch.Tensor]  # None when shapes differ
+    name: str
+
+
+@dataclass(frozen=True)
+class _Panel:
+    label: str
+    requires_diff: bool
+    draw: Callable[[np.ndarray, int, _PanelContext], Optional[str]]
+
+
+def _build_panels() -> list[_Panel]:
+    from sglang.srt.debug_utils.comparator.visualizer.panels import (
+        _draw_baseline_heatmap,
+        _draw_diff_heatmap,
+        _draw_diff_histogram,
+        _draw_hist2d,
+        _draw_sampled,
+        _draw_target_heatmap,
+    )
+
+    return [
+        _Panel(
+            label="Baseline Heatmap", requires_diff=False, draw=_draw_baseline_heatmap
+        ),
+        _Panel(label="Target Heatmap", requires_diff=False, draw=_draw_target_heatmap),
+        _Panel(label="Abs Diff Heatmap", requires_diff=True, draw=_draw_diff_heatmap),
+        _Panel(label="Abs Diff Hist", requires_diff=True, draw=_draw_diff_histogram),
+        _Panel(label="Hist2D", requires_diff=True, draw=_draw_hist2d),
+        _Panel(label="Sampled", requires_diff=True, draw=_draw_sampled),
+    ]
+
+
+def generate_comparison_figure(
+    *,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    name: str,
+    output_path: Path,
+) -> None:
+    """Generate a multi-panel comparison PNG for a baseline/target tensor pair.
+
+    Panels (6 rows x 2 cols, left=normal, right=log10):
+      Row 0: Baseline heatmap
+      Row 1: Target heatmap
+      Row 2: Abs Diff heatmap
+      Row 3: Abs Diff histogram
+      Row 4: Hist2D scatter (baseline vs target density)
+      Row 5: Sampled scatter (10k sampled mini-heatmap)
+    """
+    import matplotlib.pyplot as plt
+
+    baseline_f: torch.Tensor = baseline.detach().cpu().float()
+    target_f: torch.Tensor = target.detach().cpu().float()
+
+    can_diff: bool = baseline_f.shape == target_f.shape
+
+    baseline_2d: torch.Tensor = _preprocess_tensor(baseline_f)
+    target_2d: torch.Tensor = _preprocess_tensor(target_f)
+
+    diff: Optional[torch.Tensor] = (baseline_2d - target_2d).abs() if can_diff else None
+
+    ctx = _PanelContext(
+        baseline_2d=baseline_2d,
+        target_2d=target_2d,
+        diff=diff,
+        name=name,
+    )
+
+    panels: list[_Panel] = _build_panels()
+    active: list[_Panel] = [p for p in panels if not p.requires_diff or can_diff]
+
+    nrows: int = len(active)
+    ncols: int = 2
+    fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows))
+    if nrows == 1:
+        axes = axes.reshape(1, -1)
+
+    stats_lines: list[str] = []
+    for i, panel in enumerate(active):
+        stats_line: Optional[str] = panel.draw(axes, i, ctx)
+        if stats_line is not None:
+            stats_lines.append(stats_line)
+
+    num_stats: int = len(stats_lines)
+    title_height: float = 0.015 * num_stats + 0.015
+    fig.suptitle(
+        "\n".join(stats_lines),
+        fontsize=9,
+        family="monospace",
+        y=1 - title_height / 2,
+    )
+    plt.tight_layout(rect=[0, 0, 1, 1 - title_height])
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(str(output_path), dpi=150, bbox_inches="tight")
+    plt.close(fig)

python/sglang/srt/debug_utils/comparator/visualizer/panels.py

@@ -0,0 +1,226 @@
+"""Panel draw functions for tensor comparison visualization."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.debug_utils.comparator.visualizer.figure import _PanelContext
+from sglang.srt.debug_utils.comparator.visualizer.preprocessing import (
+    _SCATTER_SAMPLE_SIZE,
+    _format_log_ticks,
+    _format_stats,
+    _maybe_downsample_numpy,
+    _safe_hist,
+    _to_log10,
+)
+
+
+def _draw_baseline_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    _draw_heatmap_pair(
+        axes, row_idx=row_idx, t=ctx.baseline_2d, title=f"{ctx.name} Baseline"
+    )
+    return _format_stats("Baseline", ctx.baseline_2d)
+
+
+def _draw_target_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    _draw_heatmap_pair(
+        axes, row_idx=row_idx, t=ctx.target_2d, title=f"{ctx.name} Target"
+    )
+    return _format_stats("Target", ctx.target_2d)
+
+
+def _draw_diff_heatmap(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    assert ctx.diff is not None
+    _draw_heatmap_pair(axes, row_idx=row_idx, t=ctx.diff, title=f"{ctx.name} Abs Diff")
+    return _format_stats("Abs Diff", ctx.diff)
+
+
+def _draw_diff_histogram(
+    axes: np.ndarray, row_idx: int, ctx: _PanelContext
+) -> Optional[str]:
+    assert ctx.diff is not None
+    _draw_histogram_pair(
+        axes, row_idx=row_idx, diff=ctx.diff, label=f"{ctx.name} Abs Diff"
+    )
+    return None
+
+
+def _draw_hist2d(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
+    _draw_scatter_hist2d(
+        axes,
+        row_idx=row_idx,
+        baseline=ctx.baseline_2d,
+        target=ctx.target_2d,
+        label=ctx.name,
+    )
+    return None
+
+
+def _draw_sampled(axes: np.ndarray, row_idx: int, ctx: _PanelContext) -> Optional[str]:
+    _draw_scatter_sampled(
+        axes,
+        row_idx=row_idx,
+        baseline=ctx.baseline_2d,
+        target=ctx.target_2d,
+        label=ctx.name,
+    )
+    return None
+
+
+# ────────────────────── internal drawing helpers ──────────────────────
+
+
+def _draw_heatmap_pair(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    t: torch.Tensor,
+    title: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    im = ax_normal.imshow(t.numpy(), aspect="auto", cmap="viridis")
+    ax_normal.set_title(title)
+    plt.colorbar(im, ax=ax_normal)
+
+    im_log = ax_log.imshow(_to_log10(t).numpy(), aspect="auto", cmap="viridis")
+    ax_log.set_title(f"{title} (Log10)")
+    cbar = plt.colorbar(im_log, ax=ax_log)
+    _format_log_ticks(cbar.ax, axis="y")
+
+
+def _draw_histogram_pair(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    diff: torch.Tensor,
+    label: str,
+) -> None:
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    diff_flat: np.ndarray = _maybe_downsample_numpy(diff.flatten())
+
+    _safe_hist(ax_normal, diff_flat, bins=100, edgecolor="none")
+    ax_normal.set_title(f"{label} Histogram")
+    ax_normal.set_xlabel("Abs Diff")
+    ax_normal.set_ylabel("Count")
+
+    log_flat: np.ndarray = np.log10(np.abs(diff_flat) + 1e-10)
+    _safe_hist(ax_log, log_flat, bins=100, edgecolor="none")
+    ax_log.set_title(f"{label} Histogram (Log10)")
+    ax_log.set_xlabel("Abs Diff")
+    ax_log.set_ylabel("Count")
+    _format_log_ticks(ax_log, axis="x")
+
+
+def _draw_scatter_hist2d(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    label: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_normal = axes[row_idx, 0]
+    ax_log = axes[row_idx, 1]
+
+    b_flat: np.ndarray = _maybe_downsample_numpy(baseline.flatten())
+    t_flat: np.ndarray = _maybe_downsample_numpy(target.flatten())
+    min_len: int = min(len(b_flat), len(t_flat))
+    b_flat = b_flat[:min_len]
+    t_flat = t_flat[:min_len]
+
+    # Normal scale
+    lim: float = float(max(np.abs(b_flat).max(), np.abs(t_flat).max())) * 1.05
+    if lim == 0:
+        lim = 1.0
+    _h, _xe, _ye, im = ax_normal.hist2d(
+        b_flat,
+        t_flat,
+        bins=200,
+        range=[[-lim, lim], [-lim, lim]],
+        cmap="viridis",
+        norm="log",
+    )
+    ax_normal.plot([-lim, lim], [-lim, lim], "r--", linewidth=0.5)
+    ax_normal.set_title(f"{label} Hist2D")
+    ax_normal.set_xlabel("Baseline")
+    ax_normal.set_ylabel("Target")
+    ax_normal.set_aspect("equal")
+    plt.colorbar(im, ax=ax_normal)
+
+    # Log scale
+    b_log: np.ndarray = np.log10(np.abs(b_flat) + 1e-10)
+    t_log: np.ndarray = np.log10(np.abs(t_flat) + 1e-10)
+    vmin: float = float(min(b_log.min(), t_log.min())) - 0.5
+    vmax: float = float(max(b_log.max(), t_log.max())) + 0.5
+    _h2, _xe2, _ye2, im2 = ax_log.hist2d(
+        b_log,
+        t_log,
+        bins=200,
+        range=[[vmin, vmax], [vmin, vmax]],
+        cmap="viridis",
+        norm="log",
+    )
+    ax_log.plot([vmin, vmax], [vmin, vmax], "r--", linewidth=0.5)
+    ax_log.set_title(f"{label} Hist2D (Log10 Abs)")
+    ax_log.set_xlabel("Baseline")
+    ax_log.set_ylabel("Target")
+    ax_log.set_aspect("equal")
+    plt.colorbar(im2, ax=ax_log)
+    _format_log_ticks(ax_log, axis="both")
+
+
+def _draw_scatter_sampled(
+    axes: np.ndarray,
+    *,
+    row_idx: int,
+    baseline: torch.Tensor,
+    target: torch.Tensor,
+    label: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    ax_baseline = axes[row_idx, 0]
+    ax_target = axes[row_idx, 1]
+
+    b_flat: np.ndarray = baseline.flatten().numpy()
+    t_flat: np.ndarray = target.flatten().numpy()
+
+    n_samples: int = min(_SCATTER_SAMPLE_SIZE, len(b_flat))
+    rng: np.random.Generator = np.random.default_rng(seed=42)
+    indices: np.ndarray = np.sort(rng.choice(len(b_flat), n_samples, replace=False))
+    b_sampled: np.ndarray = b_flat[indices]
+    t_sampled: np.ndarray = t_flat[indices]
+
+    side: int = int(np.sqrt(n_samples))
+    n_use: int = side * side
+    b_2d: np.ndarray = b_sampled[:n_use].reshape(side, side)
+    t_2d: np.ndarray = t_sampled[:n_use].reshape(side, side)
+
+    vmin: float = float(min(b_2d.min(), t_2d.min()))
+    vmax: float = float(max(b_2d.max(), t_2d.max()))
+
+    im_b = ax_baseline.imshow(b_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
+    ax_baseline.set_title(f"{label} Baseline (10k sampled)")
+    plt.colorbar(im_b, ax=ax_baseline)
+
+    im_t = ax_target.imshow(t_2d, aspect="auto", cmap="viridis", vmin=vmin, vmax=vmax)
+    ax_target.set_title(f"{label} Target (10k sampled)")
+    plt.colorbar(im_t, ax=ax_target)

python/sglang/srt/debug_utils/comparator/visualizer/preprocessing.py

@@ -0,0 +1,101 @@
+"""Tensor preprocessing and utility functions for visualization."""
+
+from __future__ import annotations
+
+import math
+import re
+
+import numpy as np
+import torch
+
+_DOWNSAMPLE_THRESHOLD: int = 10_000_000
+_SCATTER_SAMPLE_SIZE: int = 10_000
+
+
+def _preprocess_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    t: torch.Tensor = tensor.squeeze()
+
+    while t.ndim < 2:
+        t = t.unsqueeze(0)
+    if t.ndim > 2:
+        t = t.reshape(-1, t.shape[-1])
+
+    t = _reshape_to_balanced_aspect(t)
+    return t
+
+
+def _reshape_to_balanced_aspect(
+    t: torch.Tensor, max_ratio: float = 5.0
+) -> torch.Tensor:
+    assert t.ndim == 2
+
+    h, w = t.shape
+    ratio: float = h / w if w > 0 else float("inf")
+
+    if 1 / max_ratio <= ratio <= max_ratio:
+        return t
+
+    total: int = h * w
+    target_side: int = int(math.sqrt(total))
+
+    for new_h in range(target_side, 0, -1):
+        if total % new_h == 0:
+            new_w: int = total // new_h
+            new_ratio: float = new_h / new_w
+            if 1 / max_ratio <= new_ratio <= max_ratio:
+                return t.reshape(new_h, new_w)
+
+    return t.reshape(1, -1)
+
+
+# ────────────────────── utility ──────────────────────
+
+
+def _to_log10(t: torch.Tensor) -> torch.Tensor:
+    return t.abs().clamp(min=1e-10).log10()
+
+
+def _format_log_ticks(ax: object, axis: str = "both") -> None:
+    from matplotlib.ticker import FuncFormatter
+
+    formatter = FuncFormatter(
+        lambda x, _: f"1e{int(x)}" if x == int(x) else f"1e{x:.1f}"
+    )
+    if axis in ("x", "both"):
+        ax.xaxis.set_major_formatter(formatter)
+    if axis in ("y", "both"):
+        ax.yaxis.set_major_formatter(formatter)
+
+
+def _format_stats(name: str, t: torch.Tensor) -> str:
+    return (
+        f"{name}: shape={tuple(t.shape)}, "
+        f"min={t.min().item():.4g}, max={t.max().item():.4g}, "
+        f"mean={t.mean().item():.4g}, std={t.std().item():.4g}"
+    )
+
+
+def _safe_hist(
+    ax: object, data: np.ndarray, *, bins: int = 100, **kwargs: object
+) -> None:
+    data_f64: np.ndarray = data.astype(np.float64)
+    try:
+        ax.hist(data_f64, bins=bins, **kwargs)
+    except ValueError:
+        ax.hist(data_f64, bins=max(1, len(np.unique(data_f64[:1000]))), **kwargs)
+
+
+def _maybe_downsample_numpy(
+    t: torch.Tensor,
+    max_elements: int = _DOWNSAMPLE_THRESHOLD,
+) -> np.ndarray:
+    if t.numel() <= max_elements:
+        return t.numpy()
+
+    rng: np.random.Generator = np.random.default_rng(seed=0)
+    indices: np.ndarray = rng.choice(t.numel(), max_elements, replace=False)
+    return t.numpy()[indices]
+
+
+def _sanitize_filename(name: str) -> str:
+    return re.sub(r"[/\.\s]+", "_", name).strip("_")

python/sglang/srt/debug_utils/dump_loader.py

@@ -9,8 +9,6 @@ import torch
 
 LOAD_FAILED: object = object()
 
-LOAD_FAILED: object = object()
-
 
 def parse_meta_from_filename(path: Path) -> Dict[str, Any]:
     stem = Path(path).stem