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

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

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.
 

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__]))

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__]))

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