Support overriding and post-hoc providing metadata in dump comparator (#19598)

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

@@ -25,6 +25,7 @@ from sglang.srt.debug_utils.comparator.dims import (
     resolve_dim_names,
 )
 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,
@@ -54,6 +55,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]:
     with warning_sink.context() as collected_warnings:
         result = _compare_bundle_pair_inner(
@@ -66,6 +68,7 @@ def compare_bundle_pair(
             thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
             viz_output_dir=viz_output_dir,
             compute_per_token=compute_per_token,
+            meta_overrider=meta_overrider,
         )
 
     return result.model_copy(update={"warnings": collected_warnings})
@@ -84,6 +87,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]:
     # 1. Load all successfully loaded values
     all_pair: Pair[list[ValueWithMeta]] = Pair(
@@ -101,6 +105,24 @@ def _compare_bundle_pair_inner(
         y=filter_to_non_empty_dp_rank(all_pair.y),
     )
 
+    # 1c. Dims override: patch meta["dims"] before downstream reads it
+    if meta_overrider is not None and not meta_overrider.is_empty:
+        _apply = meta_overrider.apply_to_meta
+        all_pair = Pair(
+            x=[
+                ValueWithMeta(
+                    value=v.value, meta=_apply(name=name, meta=v.meta, side="baseline")
+                )
+                for v in all_pair.x
+            ],
+            y=[
+                ValueWithMeta(
+                    value=v.value, meta=_apply(name=name, meta=v.meta, side="target")
+                )
+                for v in all_pair.y
+            ],
+        )
+
     # 2. Check if any side has non-tensor values → non-tensor display path
     has_non_tensor: bool = any(
         not isinstance(it.value, torch.Tensor) for it in [*all_pair.x, *all_pair.y]

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

@@ -22,6 +22,7 @@ from sglang.srt.debug_utils.comparator.bundle_matcher import (
     match_bundles,
 )
 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,
@@ -85,6 +86,13 @@ def run(args: argparse.Namespace) -> None:
         Path(args.visualize_per_token) if args.visualize_per_token else None
     )
 
+    meta_overrider: MetaOverrider = MetaOverrider.from_args_and_config(
+        override_dims=args.override_dims,
+        override_baseline_dims=args.override_baseline_dims,
+        override_target_dims=args.override_target_dims,
+        override_config=Path(args.override_config) if args.override_config else None,
+    )
+
     comparison_records = _compare_bundle_pairs(
         bundle_info_pairs=bundle_info_pairs,
         baseline_path=Path(args.baseline_path),
@@ -94,6 +102,7 @@ def run(args: argparse.Namespace) -> None:
         thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair,
         viz_output_dir=viz_output_dir,
         compute_per_token=visualize_per_token is not None,
+        meta_overrider=meta_overrider,
     )
     _consume_comparison_records(
         comparison_records=comparison_records,
@@ -155,6 +164,7 @@ def _compare_bundle_pairs(
     thd_seq_lens_by_step_pair: Pair[Optional[dict[int, list[int]]]],
     viz_output_dir: Optional[Path] = None,
     compute_per_token: bool = False,
+    meta_overrider: Optional[MetaOverrider] = None,
 ) -> Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]]:
     for bundle_info_pair in bundle_info_pairs:
         if not bundle_info_pair.y:
@@ -174,6 +184,7 @@ def _compare_bundle_pairs(
             thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
             viz_output_dir=viz_output_dir,
             compute_per_token=compute_per_token,
+            meta_overrider=meta_overrider,
         )
 
 
@@ -252,4 +263,31 @@ def _parse_args() -> argparse.Namespace:
         default=None,
         help="Output path for per-token relative difference heatmap PNG",
     )
+
+    # Dims override
+    parser.add_argument(
+        "--override-dims",
+        action="append",
+        default=[],
+        help="Override dims for both sides: 'name:dims_string' (repeatable)",
+    )
+    parser.add_argument(
+        "--override-baseline-dims",
+        action="append",
+        default=[],
+        help="Override dims for baseline only: 'name:dims_string' (repeatable)",
+    )
+    parser.add_argument(
+        "--override-target-dims",
+        action="append",
+        default=[],
+        help="Override dims for target only: 'name:dims_string' (repeatable)",
+    )
+    parser.add_argument(
+        "--override-config",
+        type=str,
+        default=None,
+        help="Path to YAML override config file (dims overrides, etc.)",
+    )
+
     return parser.parse_args()

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

@@ -0,0 +1,107 @@
+"""Meta overrider: replace metadata fields without re-running dumps.
+
+Currently only overrides 'dims', but the design supports overriding
+additional meta fields (e.g. parallel_info) in the future.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+from typing import Any, Literal, Optional
+
+import yaml
+
+from sglang.srt.debug_utils.comparator.utils import _StrictBase
+
+
+class MetaOverrideRule(_StrictBase):
+    """Single override rule: regex match on tensor name → replacement meta field(s).
+
+    Currently only 'dims' is supported; more fields may be added in the future.
+    """
+
+    match: str
+    dims: str
+    side: Literal["both", "baseline", "target"] = "both"
+
+
+class MetaOverrideConfig(_StrictBase):
+    """YAML top-level config for overriding comparator behavior."""
+
+    overrides: list[MetaOverrideRule] = []
+
+
+class MetaOverrider:
+    """Holds override rules and applies first-match-wins replacement."""
+
+    def __init__(self, rules: list[MetaOverrideRule]) -> None:
+        self._rules: list[MetaOverrideRule] = rules
+
+    @property
+    def is_empty(self) -> bool:
+        return len(self._rules) == 0
+
+    @classmethod
+    def from_args_and_config(
+        cls,
+        *,
+        override_dims: list[str],
+        override_baseline_dims: list[str],
+        override_target_dims: list[str],
+        override_config: Optional[Path],
+    ) -> "MetaOverrider":
+        per_side_args: list[tuple[list[str], Literal["both", "baseline", "target"]]] = [
+            (override_dims, "both"),
+            (override_baseline_dims, "baseline"),
+            (override_target_dims, "target"),
+        ]
+        cli_rules: list[MetaOverrideRule] = [
+            MetaOverrideRule(match=name, dims=dims_str, side=side)
+            for raw_args, side in per_side_args
+            for name, dims_str in [_parse_cli_override_arg(raw) for raw in raw_args]
+        ]
+
+        yaml_rules: list[MetaOverrideRule] = (
+            _load_yaml_rules(override_config) if override_config is not None else []
+        )
+
+        return cls(rules=cli_rules + yaml_rules)
+
+    def apply_to_meta(
+        self,
+        *,
+        name: str,
+        meta: dict[str, Any],
+        side: Literal["baseline", "target"],
+    ) -> dict[str, Any]:
+        """First-match-wins: return meta with dims replaced by the first matching rule for this side."""
+        for rule in self._rules:
+            if rule.side not in ("both", side):
+                continue
+            if re.search(rule.match, name):
+                return {**meta, "dims": rule.dims}
+
+        return meta
+
+
+def _parse_cli_override_arg(raw: str) -> tuple[str, str]:
+    """Parse 'name:dims_string' from a CLI --override-* argument."""
+    parts: list[str] = raw.split(":", maxsplit=1)
+    if len(parts) != 2 or not parts[0].strip() or not parts[1].strip():
+        raise ValueError(
+            f"Invalid override format: {raw!r}; expected 'name:dims_string'"
+        )
+    return parts[0].strip(), parts[1].strip()
+
+
+def _load_yaml_rules(path: Path) -> list[MetaOverrideRule]:
+    """Load override rules from a YAML config file."""
+    with open(path) as f:
+        raw_data: Any = yaml.safe_load(f)
+
+    if raw_data is None:
+        return []
+
+    config: MetaOverrideConfig = MetaOverrideConfig.model_validate(raw_data)
+    return config.overrides