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

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