Support presets and arbitrary skipping keys in dump comparator (#19676)

test/registered/debug_utils/comparator/test_manually_verify.py

@@ -207,7 +207,9 @@ class TestPerTokenHeatmapManualVerify:
         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.output_types import (
+            TensorComparisonRecord,
+        )
         from sglang.srt.debug_utils.comparator.per_token_visualizer import (
             generate_per_token_heatmap,
         )
@@ -220,7 +222,7 @@ class TestPerTokenHeatmapManualVerify:
         hidden_dim: int = 128
         num_tensors: int = 5
 
-        records: list[ComparisonRecord] = []
+        records: list[TensorComparisonRecord] = []
         for i in range(num_tensors):
             baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
             noise_scale: torch.Tensor = torch.linspace(
@@ -235,7 +237,7 @@ class TestPerTokenHeatmapManualVerify:
                 diff_threshold=1e-3,
                 seq_dim=0,
             )
-            records.append(ComparisonRecord(**info.model_dump()))
+            records.append(TensorComparisonRecord(**info.model_dump()))
 
         output_path: Path = tmp_path / "per_token_increasing_diff.png"
         result = generate_per_token_heatmap(records=records, output_path=output_path)
@@ -250,7 +252,9 @@ class TestPerTokenHeatmapManualVerify:
         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.output_types import (
+            TensorComparisonRecord,
+        )
         from sglang.srt.debug_utils.comparator.per_token_visualizer import (
             generate_per_token_heatmap,
         )
@@ -264,7 +268,7 @@ class TestPerTokenHeatmapManualVerify:
         spike_pos: int = 32
         num_tensors: int = 4
 
-        records: list[ComparisonRecord] = []
+        records: list[TensorComparisonRecord] = []
         for i in range(num_tensors):
             baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
             target: torch.Tensor = baseline.clone()
@@ -277,7 +281,7 @@ class TestPerTokenHeatmapManualVerify:
                 diff_threshold=1e-3,
                 seq_dim=0,
             )
-            records.append(ComparisonRecord(**info.model_dump()))
+            records.append(TensorComparisonRecord(**info.model_dump()))
 
         output_path: Path = tmp_path / "per_token_single_spike.png"
         result = generate_per_token_heatmap(records=records, output_path=output_path)