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

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