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,15 +48,15 @@ def compute_maybe_token_aligner_result(
     args: argparse.Namespace,
     dfs: Pair[pl.DataFrame],
 ) -> TokenAlignerResult:
-    if args.grouping != "logical":
+    token_aligner_mode: Optional[TokenAlignerMode] = getattr(
+        args, "token_aligner", None
+    )
+
+    if token_aligner_mode is None:
         return TokenAlignerResult(
             mode=None, plan=None, thd_seq_lens_by_step_pair=_NONE_THD
         )
 
-    token_aligner_mode: TokenAlignerMode = getattr(
-        args, "token_aligner", "concat_steps"
-    )
-
     if token_aligner_mode == "concat_steps":
         thd_pair: Pair[Optional[dict[int, list[int]]]] = _load_thd_seq_lens_pair(
             args=args, dfs=dfs

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,
-    SkipRecord,
+    NonTensorComparisonRecord,
+    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,
-    SkipRecord,
+    NonTensorComparisonRecord,
+    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(
-                args, has_token_aligner=ta_result.mode is not None
-            ),
+            skip_keys=_compute_skip_keys(args),
         )
 
         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]:
-    skip_keys: set[str] = {"dump_index", "filename"}
-    if args.grouping == "logical":
-        skip_keys |= {"rank", "recompute_status"}
-        if has_token_aligner:
-            skip_keys |= {"step"}
-    return skip_keys
+def _compute_skip_keys(args: argparse.Namespace) -> set[str]:
+    return _DEFAULT_SKIP_KEYS | set(args.grouping_skip_keys or [])
 
 
 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"],
-        default="logical",
-        help="Grouping mode: logical (cross-rank unshard) or raw (rank-by-rank)",
+        choices=list(PRESETS.keys()),
+        default=None,
+        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",
-        help="Token aligner mode: concat_steps (BS=1, no aux needed) or smart (BS>1, sequence matching)",
+        default=None,
+        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,
-        ComparisonRecord,
-        NonTensorRecord,
+        SkipComparisonRecord,
+        TensorComparisonRecord,
+        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 :]