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

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

@@ -1,6 +1,6 @@
 from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (  # noqa: F401
     AlignerPlan,
 )
-from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord
+from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord
 
-ComparisonRecord.model_rebuild()
+TensorComparisonRecord.model_rebuild()

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

@@ -48,13 +48,13 @@ def compute_maybe_token_aligner_result(
     args: argparse.Namespace,
     dfs: Pair[pl.DataFrame],
 ) -> TokenAlignerResult:
-    if args.grouping != "logical":
+    token_aligner_mode: Optional[TokenAlignerMode] = getattr(
-        return TokenAlignerResult(
+        args, "token_aligner", None
-            mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD
     )
 
-    token_aligner_mode: TokenAlignerMode = getattr(
+    if token_aligner_mode is None:
-        args, "token_aligner", "concat_steps"
+        return TokenAlignerResult(
+            mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD
         )
 
     if token_aligner_mode == "concat_steps":

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

@@ -28,10 +28,10 @@ from sglang.srt.debug_utils.comparator.dims import (
 from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
 from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
 from sglang.srt.debug_utils.comparator.output_types import (
-    ComparisonRecord,
     GeneralWarning,
-    NonTensorRecord,
+    NonTensorComparisonRecord,
-    SkipRecord,
+    SkipComparisonRecord,
+    TensorComparisonRecord,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
     compare_tensor_pair,
@@ -58,7 +58,7 @@ def compare_bundle_pair(
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
     meta_overrider: Optional[MetaOverrider] = None,
-) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
+) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]:
     with warning_sink.context() as collected_warnings:
         result = _compare_bundle_pair_inner(
             name=name,
@@ -92,7 +92,7 @@ def _compare_bundle_pair_inner(
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
     meta_overrider: Optional[MetaOverrider] = None,
-) -> Union[ComparisonRecord, SkipRecord, NonTensorRecord]:
+) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]:
     # 1. Load all successfully loaded values
     all_pair: Pair[list[ValueWithMeta]] = Pair(
         x=_load_all_values(filenames=filenames_pair.x, base_path=baseline_path),
@@ -101,7 +101,7 @@ def _compare_bundle_pair_inner(
 
     if not all_pair.x or not all_pair.y:
         reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
-        return SkipRecord(name=name, reason=reason)
+        return SkipComparisonRecord(name=name, reason=reason)
 
     # 1b. Dims override: patch meta["dims"] before DP filter reads it
     # (--override-dims may add ``# dp:=moe_dp``, so it must run first)
@@ -171,10 +171,10 @@ def _compare_bundle_pair_tensor_type(
     ),
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
-) -> Union[ComparisonRecord, SkipRecord]:
+) -> Union[TensorComparisonRecord, SkipComparisonRecord]:
     if not valid_pair.x or not valid_pair.y:
         reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
-        return SkipRecord(name=name, reason=reason)
+        return SkipComparisonRecord(name=name, reason=reason)
 
     # Plan (meta only, no tensor)
     metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
@@ -207,7 +207,7 @@ def _compare_bundle_pair_tensor_type(
         assert aligner_result.failed_side_xy is not None
         side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
         reason: str = f"{side_name}_load_failed"
-        return SkipRecord(name=name, reason=reason)
+        return SkipComparisonRecord(name=name, reason=reason)
 
     # Resolve seq_dim for per-token computation
     seq_dim: Optional[int] = (
@@ -225,7 +225,7 @@ def _compare_bundle_pair_tensor_type(
         diff_threshold=diff_threshold,
         seq_dim=seq_dim,
     )
-    record = ComparisonRecord(
+    record = TensorComparisonRecord(
         **info.model_dump(),
         aligner_plan=plan,
         replicated_checks=replicated_checks,
@@ -292,7 +292,7 @@ def _compare_bundle_pair_non_tensor_type(
     *,
     name: str,
     value_pair: Pair[list[ValueWithMeta]],
-) -> NonTensorRecord:
+) -> NonTensorComparisonRecord:
     baseline_value: Any = value_pair.x[0].value
     target_value: Any = value_pair.y[0].value
 
@@ -301,7 +301,7 @@ def _compare_bundle_pair_non_tensor_type(
     except Exception:
         values_equal = False
 
-    return NonTensorRecord(
+    return NonTensorComparisonRecord(
         name=name,
         baseline_value=repr(baseline_value),
         target_value=repr(target_value),

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

@@ -26,22 +26,26 @@ from sglang.srt.debug_utils.comparator.bundle_matcher import (
 from sglang.srt.debug_utils.comparator.display import emit_display_records
 from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
 from sglang.srt.debug_utils.comparator.output_types import (
-    ComparisonRecord,
     ConfigRecord,
-    NonTensorRecord,
+    NonTensorComparisonRecord,
-    SkipRecord,
+    RecordLocation,
+    SkipComparisonRecord,
     SummaryRecord,
+    TensorComparisonRecord,
     report_sink,
 )
 from sglang.srt.debug_utils.comparator.per_token_visualizer import (
     generate_per_token_heatmap,
 )
+from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
 from sglang.srt.debug_utils.comparator.utils import Pair
 from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
 
+_DEFAULT_SKIP_KEYS: set[str] = {"dump_index", "filename"}
+
 
 def main() -> None:
-    args = _parse_args()
+    args = parse_args(sys.argv[1:])
     sys.exit(run(args))
 
 
@@ -76,9 +80,7 @@ def run(args: argparse.Namespace) -> int:
 
         bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
             dfs=dfs,
-            skip_keys=_compute_skip_keys(
+            skip_keys=_compute_skip_keys(args),
-                args, has_token_aligner=ta_result.mode is not None
-            ),
         )
 
         viz_output_dir: Optional[Path] = (
@@ -202,13 +204,8 @@ def _read_df(args: argparse.Namespace) -> Pair[pl.DataFrame]:
     return Pair(x=df_baseline, y=df_target)
 
 
-def _compute_skip_keys(args, *, has_token_aligner: bool) -> set[str]:
+def _compute_skip_keys(args: argparse.Namespace) -> set[str]:
-    skip_keys: set[str] = {"dump_index", "filename"}
+    return _DEFAULT_SKIP_KEYS | set(args.grouping_skip_keys or [])
-    if args.grouping == "logical":
-        skip_keys |= {"rank", "recompute_status"}
-        if has_token_aligner:
-            skip_keys |= {"step"}
-    return skip_keys
 
 
 def _compare_bundle_pairs(
@@ -223,7 +220,9 @@ def _compare_bundle_pairs(
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
     meta_overrider: Optional[MetaOverrider] = None,
-) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]:
+) -> Iterator[
+    Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]
+]:
     for bundle_info_pair in bundle_info_pairs:
         if not bundle_info_pair.y:
             continue
@@ -232,7 +231,9 @@ def _compare_bundle_pairs(
         filenames_pair: Pair[list[str]] = bundle_info_pair.map(
             lambda infos: [info.filename for info in infos]
         )
-        yield compare_bundle_pair(
+        record: Union[
+            TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord
+        ] = compare_bundle_pair(
             name=name,
             filenames_pair=filenames_pair,
             baseline_path=baseline_path,
@@ -246,22 +247,33 @@ def _compare_bundle_pairs(
             meta_overrider=meta_overrider,
         )
 
+        target_steps: set[int] = {info.step for info in bundle_info_pair.y}
+        step: Optional[int] = target_steps.pop() if len(target_steps) == 1 else None
+        if step is not None:
+            record = record.model_copy(update={"location": RecordLocation(step=step)})
+
+        yield record
+
 
 def _consume_comparison_records(
     *,
-    comparison_records: Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]],
+    comparison_records: Iterator[
+        Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]
+    ],
     visualize_per_token: Optional[Path] = None,
 ) -> tuple[SummaryRecord, list[str]]:
     counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
-    collected_comparisons: list[ComparisonRecord] = []
+    collected_comparisons: list[TensorComparisonRecord] = []
     skipped_names: list[str] = []
 
     for record in comparison_records:
         counts[record.category] += 1
         report_sink.add(record)
-        if isinstance(record, SkipRecord) and record.category == "skipped":
+        if isinstance(record, SkipComparisonRecord) and record.category == "skipped":
             skipped_names.append(record.name)
-        if visualize_per_token is not None and isinstance(record, ComparisonRecord):
+        if visualize_per_token is not None and isinstance(
+            record, TensorComparisonRecord
+        ):
             collected_comparisons.append(record)
 
     summary: SummaryRecord = SummaryRecord(total=sum(counts.values()), **counts)
@@ -276,7 +288,10 @@ def _consume_comparison_records(
     return summary, skipped_names
 
 
-def _parse_args() -> argparse.Namespace:
+def parse_args(argv: list[str]) -> argparse.Namespace:
+    """Parse CLI arguments from an argv list. Applies preset expansion."""
+    argv = expand_preset(argv, presets=PRESETS)
+
     parser = argparse.ArgumentParser()
     parser.add_argument("--baseline-path", type=str)
     parser.add_argument("--target-path", type=str)
@@ -294,18 +309,28 @@ def _parse_args() -> argparse.Namespace:
         help="Output format: text (default) or json (JSONL, one JSON object per line)",
     )
     parser.add_argument(
-        "--grouping",
+        "--preset",
         type=str,
-        choices=["logical", "raw"],
+        choices=list(PRESETS.keys()),
-        default="logical",
+        default=None,
-        help="Grouping mode: logical (cross-rank unshard) or raw (rank-by-rank)",
+        help="Preset configuration (expanded before parsing). "
+        f"Available: {list(PRESETS.keys())}",
+    )
+    parser.add_argument(
+        "--grouping-skip-keys",
+        nargs="*",
+        default=None,
+        help="Metadata keys to skip when grouping bundles (additive on top of "
+        "always-skipped dump_index and filename). "
+        "E.g. '--grouping-skip-keys rank step' skips rank and step.",
     )
     parser.add_argument(
         "--token-aligner",
         type=str,
         choices=["smart", "concat_steps"],
-        default="concat_steps",
+        default=None,
-        help="Token aligner mode: concat_steps (BS=1, no aux needed) or smart (BS>1, sequence matching)",
+        help="Token aligner mode: concat_steps (BS=1, no aux needed) or smart (BS>1, sequence matching). "
+        "Default None (per-step comparison).",
     )
     parser.add_argument(
         "--tokenizer",
@@ -374,4 +399,4 @@ def _parse_args() -> argparse.Namespace:
         "Pass empty string '' to disable.",
     )
 
-    return parser.parse_args()
+    return parser.parse_args(argv)

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

@@ -61,6 +61,24 @@ class _OutputRecord(_StrictBase):
         return body
 
 
+class RecordLocation(_StrictBase):
+    step: Optional[int] = None
+
+
+class _BaseComparisonRecord(_OutputRecord):
+    location: RecordLocation = Field(default_factory=RecordLocation)
+
+    def _format_location_prefix(self) -> str:
+        if self.location.step is not None:
+            return f"[step={self.location.step}] "
+        return ""
+
+    def _format_location_suffix(self) -> str:
+        if self.location.step is not None:
+            return f" (step={self.location.step})"
+        return ""
+
+
 class ConfigRecord(_OutputRecord):
     type: Literal["config"] = "config"
     config: dict[str, Any]
@@ -74,7 +92,7 @@ class ConfigRecord(_OutputRecord):
         return f"Config: {self.config}"
 
 
-class SkipRecord(_OutputRecord):
+class SkipComparisonRecord(_BaseComparisonRecord):
     type: Literal["skip"] = "skip"
     name: str
     reason: str
@@ -86,7 +104,7 @@ class SkipRecord(_OutputRecord):
         return "skipped"
 
     def _format_body(self) -> str:
-        return f"Skip: {self.name} ({self.reason})"
+        return f"Skip: {self.name}{self._format_location_suffix()} ({self.reason})"
 
 
 class _TableRecord(_OutputRecord):
@@ -118,7 +136,7 @@ class InputIdsRecord(_TableRecord):
         return f"{self.label} input_ids & positions"
 
 
-class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
+class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
     model_config = ConfigDict(extra="forbid", defer_build=True)
 
     type: Literal["comparison"] = "comparison"
@@ -134,7 +152,7 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
         return "passed" if self.diff is not None and self.diff.passed else "failed"
 
     def _format_body(self) -> str:
-        body: str = format_comparison(self)
+        body: str = self._format_location_prefix() + format_comparison(self)
         if self.replicated_checks:
             body += "\n" + format_replicated_checks(self.replicated_checks)
         if self.aligner_plan is not None:
@@ -142,7 +160,7 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
         return body
 
 
-class NonTensorRecord(_OutputRecord):
+class NonTensorComparisonRecord(_BaseComparisonRecord):
     type: Literal["non_tensor"] = "non_tensor"
     name: str
     baseline_value: str
@@ -158,10 +176,11 @@ class NonTensorRecord(_OutputRecord):
         return "passed" if self.values_equal else "failed"
 
     def _format_body(self) -> str:
+        suffix: str = self._format_location_suffix()
         if self.values_equal:
-            return f"NonTensor: {self.name} = {self.baseline_value} ({self.baseline_type}) [equal]"
+            return f"NonTensor: {self.name}{suffix} = {self.baseline_value} ({self.baseline_type}) [equal]"
         return (
-            f"NonTensor: {self.name}\n"
+            f"NonTensor: {self.name}{suffix}\n"
             f"  baseline = {self.baseline_value} ({self.baseline_type})\n"
             f"  target   = {self.target_value} ({self.target_type})"
         )
@@ -237,9 +256,9 @@ AnyRecord = Annotated[
         ConfigRecord,
         RankInfoRecord,
         InputIdsRecord,
-        SkipRecord,
+        SkipComparisonRecord,
-        ComparisonRecord,
+        TensorComparisonRecord,
-        NonTensorRecord,
+        NonTensorComparisonRecord,
         SummaryRecord,
         WarningRecord,
     ],

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

@@ -9,12 +9,12 @@ from __future__ import annotations
 from pathlib import Path
 from typing import Optional
 
-from sglang.srt.debug_utils.comparator.output_types import ComparisonRecord
+from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord
 
 
 def generate_per_token_heatmap(
     *,
-    records: list[ComparisonRecord],
+    records: list[TensorComparisonRecord],
     output_path: Path,
 ) -> Optional[Path]:
     """Generate a per-token relative difference heatmap PNG.
@@ -31,7 +31,7 @@ def generate_per_token_heatmap(
 
 def _collect_per_token_data(
     *,
-    records: list[ComparisonRecord],
+    records: list[TensorComparisonRecord],
 ) -> list[tuple[str, list[float]]]:
     rows: list[tuple[str, list[float]]] = []
     for record in records:

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

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+PRESETS: dict[str, list[str]] = {
+    "raw": [
+        "--grouping-skip-keys",
+    ],
+    "sglang_dev": [
+        "--grouping-skip-keys",
+        "rank",
+    ],
+    "sglang_megatron": [
+        "--grouping-skip-keys",
+        "rank",
+        "step",
+        "--token-aligner",
+        "concat_steps",
+    ],
+}
+
+DEFAULT_PRESET: str = "sglang_dev"
+
+
+def expand_preset(argv: list[str], presets: dict[str, list[str]]) -> list[str]:
+    """Expand ``--preset <name>`` into the corresponding argv fragment.
+
+    If ``--preset`` is absent **and** ``--grouping-skip-keys`` is also absent,
+    the DEFAULT_PRESET is applied automatically.
+    """
+    if (expanded := _expand_flag(argv, "--preset", presets)) is not None:
+        return expanded
+
+    if "--grouping-skip-keys" not in argv:
+        return presets[DEFAULT_PRESET] + argv
+
+    return argv
+
+
+def _expand_flag(
+    argv: list[str], flag: str, mapping: dict[str, list[str]]
+) -> list[str] | None:
+    """Replace ``flag <name>`` in *argv* with the corresponding argv fragment from *mapping*."""
+    if flag not in argv:
+        return None
+
+    idx: int = argv.index(flag)
+    name: str = argv[idx + 1]
+    if name not in mapping:
+        raise ValueError(
+            f"Unknown value for {flag}: {name}. Available: {list(mapping.keys())}"
+        )
+
+    return argv[:idx] + mapping[name] + argv[idx + 2 :]

test/registered/debug_utils/comparator/tensor_comparator/test_types.py

@@ -4,12 +4,12 @@ import sys
 import pytest
 
 from sglang.srt.debug_utils.comparator.output_types import (
-    ComparisonRecord,
     ConfigRecord,
     GeneralWarning,
     ReplicatedCheckResult,
-    SkipRecord,
+    SkipComparisonRecord,
     SummaryRecord,
+    TensorComparisonRecord,
     WarningRecord,
     parse_record_json,
 )
@@ -83,7 +83,7 @@ class TestStrictBase:
 
 class TestRecordTypes:
     def test_comparison_record_inherits_tensor_fields(self):
-        record = ComparisonRecord(
+        record = TensorComparisonRecord(
             name="hidden_states",
             baseline=_make_tensor_info(),
             target=_make_tensor_info(),
@@ -108,8 +108,8 @@ class TestRecordTypes:
                     "end_step": 100,
                 }
             ),
-            SkipRecord(name="attn", reason="no_baseline"),
+            SkipComparisonRecord(name="attn", reason="no_baseline"),
-            ComparisonRecord(
+            TensorComparisonRecord(
                 name="mlp",
                 baseline=_make_tensor_info(),
                 target=_make_tensor_info(),
@@ -148,8 +148,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
 
 class TestWarnings:
     def test_comparison_record_failed_when_diff_passed_but_warnings(self):
-        """ComparisonRecord with diff.passed=True but warnings → category=='failed'."""
+        """TensorComparisonRecord with diff.passed=True but warnings → category=='failed'."""
-        record = ComparisonRecord(
+        record = TensorComparisonRecord(
             name="hidden",
             baseline=_make_tensor_info(),
             target=_make_tensor_info(),
@@ -161,8 +161,8 @@ class TestWarnings:
         assert record.category == "failed"
 
     def test_skip_record_failed_when_warnings(self):
-        """SkipRecord with warnings → category=='failed' instead of 'skipped'."""
+        """SkipComparisonRecord with warnings → category=='failed' instead of 'skipped'."""
-        record = SkipRecord(
+        record = SkipComparisonRecord(
             name="x",
             reason="no_baseline",
             warnings=[GeneralWarning(category="test", message="some warning")],
@@ -170,8 +170,8 @@ class TestWarnings:
         assert record.category == "failed"
 
     def test_replicated_checks_all_passed(self):
-        """ComparisonRecord with all replicated_checks passed → category=='passed'."""
+        """TensorComparisonRecord with all replicated_checks passed → category=='passed'."""
-        record = ComparisonRecord(
+        record = TensorComparisonRecord(
             name="hidden",
             baseline=_make_tensor_info(),
             target=_make_tensor_info(),
@@ -183,8 +183,8 @@ class TestWarnings:
         assert record.category == "passed"
 
     def test_replicated_checks_failed_means_record_failed(self):
-        """ComparisonRecord with any replicated_check.passed=False → category=='failed'."""
+        """TensorComparisonRecord with any replicated_check.passed=False → category=='failed'."""
-        record = ComparisonRecord(
+        record = TensorComparisonRecord(
             name="hidden",
             baseline=_make_tensor_info(),
             target=_make_tensor_info(),
@@ -196,7 +196,7 @@ class TestWarnings:
         assert record.category == "failed"
 
     def test_replicated_check_json_round_trip(self):
-        """ReplicatedCheckResult survives JSON round-trip via ComparisonRecord."""
+        """ReplicatedCheckResult survives JSON round-trip via TensorComparisonRecord."""
         check = _make_replicated_check(
             axis="cp",
             group_index=2,
@@ -204,7 +204,7 @@ class TestWarnings:
             baseline_index=0,
             passed=False,
         )
-        record = ComparisonRecord(
+        record = TensorComparisonRecord(
             name="mlp",
             baseline=_make_tensor_info(),
             target=_make_tensor_info(),
@@ -215,7 +215,7 @@ class TestWarnings:
         )
 
         restored = parse_record_json(record.model_dump_json())
-        assert isinstance(restored, ComparisonRecord)
+        assert isinstance(restored, TensorComparisonRecord)
         assert len(restored.replicated_checks) == 1
 
         restored_check: ReplicatedCheckResult = restored.replicated_checks[0]

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)

test/registered/debug_utils/comparator/test_model_validation.py

@@ -22,11 +22,11 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
 )
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import (
-    ComparisonRecord,
     GeneralWarning,
-    NonTensorRecord,
+    NonTensorComparisonRecord,
-    SkipRecord,
+    SkipComparisonRecord,
     SummaryRecord,
+    TensorComparisonRecord,
     parse_record_json,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -208,9 +208,9 @@ def _make_comparison_record(
     *,
     diff: DiffInfo | None,
     warnings: list | None = None,
-) -> ComparisonRecord:
+) -> TensorComparisonRecord:
     ti: TensorInfo = _make_tensor_info()
-    return ComparisonRecord(
+    return TensorComparisonRecord(
         name="t",
         baseline=ti,
         target=ti,
@@ -223,7 +223,7 @@ def _make_comparison_record(
 
 class TestOutputRecordCategories:
     def test_skip_record_with_warnings_is_failed(self) -> None:
-        record = SkipRecord(
+        record = SkipComparisonRecord(
             name="t",
             reason="test",
             warnings=[GeneralWarning(category="c", message="m")],
@@ -231,28 +231,28 @@ class TestOutputRecordCategories:
         assert record.category == "failed"
 
     def test_skip_record_no_warnings_is_skipped(self) -> None:
-        record = SkipRecord(name="t", reason="test")
+        record = SkipComparisonRecord(name="t", reason="test")
         assert record.category == "skipped"
 
     def test_comparison_record_diff_none_is_failed(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(diff=None)
+        record: TensorComparisonRecord = _make_comparison_record(diff=None)
         assert record.category == "failed"
 
     def test_comparison_record_passed_with_warnings_is_failed(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
             warnings=[GeneralWarning(category="c", message="m")],
         )
         assert record.category == "failed"
 
     def test_comparison_record_passed_no_warnings_is_passed(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
         )
         assert record.category == "passed"
 
     def test_non_tensor_record_equal_is_passed(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.125",
@@ -263,7 +263,7 @@ class TestOutputRecordCategories:
         assert record.category == "passed"
 
     def test_non_tensor_record_different_is_failed(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.25",
@@ -274,7 +274,7 @@ class TestOutputRecordCategories:
         assert record.category == "failed"
 
     def test_non_tensor_record_with_warnings_is_failed(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.125",
@@ -286,7 +286,7 @@ class TestOutputRecordCategories:
         assert record.category == "failed"
 
     def test_non_tensor_record_json_roundtrip(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.25",
@@ -296,14 +296,14 @@ class TestOutputRecordCategories:
         )
         json_str: str = record.model_dump_json()
         roundtripped = parse_record_json(json_str)
-        assert isinstance(roundtripped, NonTensorRecord)
+        assert isinstance(roundtripped, NonTensorComparisonRecord)
         assert roundtripped.name == "sm_scale"
         assert roundtripped.values_equal is False
         assert roundtripped.baseline_value == "0.125"
         assert roundtripped.target_value == "0.25"
 
     def test_non_tensor_record_text_format_equal(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.125",
@@ -316,7 +316,7 @@ class TestOutputRecordCategories:
         assert "[equal]" in text
 
     def test_non_tensor_record_text_format_different(self) -> None:
-        record = NonTensorRecord(
+        record = NonTensorComparisonRecord(
             name="sm_scale",
             baseline_value="0.125",
             target_value="0.25",
@@ -351,10 +351,10 @@ def _make_aligner_plan() -> AlignerPlan:
     )
 
 
-class TestAlignerPlanInComparisonRecord:
+class TestAlignerPlanInTensorComparisonRecord:
     def test_comparison_record_with_aligner_plan(self) -> None:
         plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
         )
         record_with_plan = record.model_copy(update={"aligner_plan": plan})
@@ -363,7 +363,7 @@ class TestAlignerPlanInComparisonRecord:
 
     def test_aligner_plan_json_roundtrip(self) -> None:
         plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
         )
         record_with_plan = record.model_copy(update={"aligner_plan": plan})
@@ -376,7 +376,7 @@ class TestAlignerPlanInComparisonRecord:
             == "unsharder"
         )
 
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        roundtripped: TensorComparisonRecord = parse_record_json(json_str)
         assert roundtripped.aligner_plan is not None
         assert (
             roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
@@ -384,16 +384,16 @@ class TestAlignerPlanInComparisonRecord:
         )
 
     def test_comparison_record_without_aligner_plan(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
         )
         json_str: str = record.model_dump_json()
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        roundtripped: TensorComparisonRecord = parse_record_json(json_str)
         assert roundtripped.aligner_plan is None
 
     def test_aligner_plan_text_format(self) -> None:
         plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
+        record: TensorComparisonRecord = _make_comparison_record(
             diff=_make_diff_info(passed=True),
         )
         record_with_plan = record.model_copy(update={"aligner_plan": plan})

test/registered/debug_utils/comparator/test_per_token_visualizer.py

@@ -9,7 +9,7 @@ from pathlib import Path
 import pytest
 import torch
 
-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.tensor_comparator.comparator import (
     compare_tensor_pair,
 )
@@ -31,8 +31,8 @@ def _make_comparison_record(
     baseline: torch.Tensor,
     target: torch.Tensor,
     seq_dim: int = 0,
-) -> ComparisonRecord:
+) -> TensorComparisonRecord:
-    """Build a ComparisonRecord with per-token data from raw tensors."""
+    """Build a TensorComparisonRecord with per-token data from raw tensors."""
     info = compare_tensor_pair(
         x_baseline=baseline,
         x_target=target,
@@ -40,7 +40,7 @@ def _make_comparison_record(
         diff_threshold=1e-3,
         seq_dim=seq_dim,
     )
-    return ComparisonRecord(**info.model_dump())
+    return TensorComparisonRecord(**info.model_dump())
 
 
 class TestPerTokenVisualizer:
@@ -68,7 +68,7 @@ class TestPerTokenVisualizer:
             name="no_per_token",
             diff_threshold=1e-3,
         )
-        record = ComparisonRecord(**info.model_dump())
+        record = TensorComparisonRecord(**info.model_dump())
 
         output_path: Path = tmp_path / "no_data.png"
         result = generate_per_token_heatmap(records=[record], output_path=output_path)
@@ -82,7 +82,7 @@ class TestPerTokenVisualizer:
         )
 
         torch.manual_seed(42)
-        records: list[ComparisonRecord] = [
+        records: list[TensorComparisonRecord] = [
             _make_comparison_record(
                 name=f"tensor_{i}",
                 baseline=torch.randn(16, 32),
@@ -108,7 +108,7 @@ class TestPerTokenVisualizer:
         )
 
         torch.manual_seed(42)
-        records: list[ComparisonRecord] = [
+        records: list[TensorComparisonRecord] = [
             _make_comparison_record(
                 name="short",
                 baseline=torch.randn(4, 8),

test/registered/debug_utils/comparator/test_preset.py

@@ -0,0 +1,44 @@
+import pytest
+
+from sglang.srt.debug_utils.comparator.preset import PRESETS, expand_preset
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default", nightly=True)
+
+
+class TestExpandPreset:
+    """Test preset expansion logic."""
+
+    def test_explicit_preset(self):
+        """--preset sglang_megatron expands into its argv."""
+        argv = [
+            "--baseline-path",
+            "/a",
+            "--preset",
+            "sglang_megatron",
+            "--diff-threshold",
+            "0.01",
+        ]
+        result = expand_preset(argv, presets=PRESETS)
+        assert "--preset" not in result
+        assert "--grouping-skip-keys" in result
+        assert "concat_steps" in result
+        assert "--baseline-path" in result
+        assert "--diff-threshold" in result
+
+    def test_default_preset_applied(self):
+        """No --preset and no --grouping-skip-keys triggers default preset."""
+        argv = ["--baseline-path", "/a"]
+        result = expand_preset(argv, presets=PRESETS)
+        assert "--grouping-skip-keys" in result
+
+    def test_explicit_skip_keys_prevents_default(self):
+        """Explicit --grouping-skip-keys prevents default preset injection."""
+        argv = ["--grouping-skip-keys", "rank", "--baseline-path", "/a"]
+        result = expand_preset(argv, presets=PRESETS)
+        assert result == argv
+
+    def test_unknown_preset_raises(self):
+        """Unknown preset name raises ValueError."""
+        with pytest.raises(ValueError, match="Unknown value for --preset"):
+            expand_preset(["--preset", "nonexistent"], presets=PRESETS)