Update layer id extraction, diffing, empty handling and error sentinel in dump comparator (#19562)

python/sglang/srt/debug_utils/comparator/aligner/unsharder/parallel_info.py

@@ -6,12 +6,17 @@ from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 _PARALLEL_INFO_KEYS = ("sglang_parallel_info", "megatron_parallel_info")
 
 
+def _is_error_sentinel(value: dict) -> bool:
+    """Check if a parallel_info dict is an error sentinel (e.g. {'megatron_error': True})."""
+    return any(k.endswith("_error") for k in value)
+
+
 def normalize_parallel_info(meta: dict) -> dict[ParallelAxis, AxisInfo]:
     """Extract unified parallel info from dump meta."""
     info: Optional[dict] = None
     for key in _PARALLEL_INFO_KEYS:
         value = meta.get(key)
-        if isinstance(value, dict) and value:
+        if isinstance(value, dict) and value and not _is_error_sentinel(value):
             if info is not None:
                 raise ValueError(
                     f"Meta contains multiple parallel_info keys among {_PARALLEL_INFO_KEYS}"

python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py

@@ -159,7 +159,11 @@ def _resolve_unshard_params(
             f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
         )
 
-    if spec.name == TOKEN_DIM_NAME and thd_global_seq_lens is not None:
+    if (
+        spec.name == TOKEN_DIM_NAME
+        and spec.parallel == ParallelAxis.CP
+        and thd_global_seq_lens is not None
+    ):
         if spec.parallel is None:
             raise ValueError(
                 f"THD unshard requires a parallel axis on dim '{spec.name}', but got None"

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

@@ -90,6 +90,16 @@ def compare_tensor_pair(
 
 
 def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
+    if x.numel() == 0:
+        return TensorStats(
+            mean=0.0,
+            abs_mean=0.0,
+            std=0.0,
+            min=0.0,
+            max=0.0,
+            percentiles={},
+        )
+
     include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD
     return TensorStats(
         mean=torch.mean(x).item(),
@@ -113,6 +123,19 @@ def _compute_diff(
     x_target: torch.Tensor,
     diff_threshold: float = 1e-3,
 ) -> DiffInfo:
+    if x_baseline.numel() == 0:
+        return DiffInfo(
+            rel_diff=0.0,
+            max_abs_diff=0.0,
+            mean_abs_diff=0.0,
+            abs_diff_percentiles={},
+            max_diff_coord=[],
+            baseline_at_max=0.0,
+            target_at_max=0.0,
+            diff_threshold=diff_threshold,
+            passed=True,
+        )
+
     raw_abs_diff = (x_target - x_baseline).abs()
     max_diff_coord = argmax_coord(raw_abs_diff)
 
@@ -133,9 +156,5 @@ def _compute_diff(
         baseline_at_max=x_baseline[max_diff_coord].item(),
         target_at_max=x_target[max_diff_coord].item(),
         diff_threshold=diff_threshold,
-        passed=(
-            rel_diff <= diff_threshold
-            and max_abs_diff <= diff_threshold
-            and mean_abs_diff <= diff_threshold
-        ),
+        passed=rel_diff <= diff_threshold,
     )

python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py

@@ -79,16 +79,12 @@ def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list
 
 
 def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
+    rel_diff_marker: str = "❌" if diff.rel_diff > diff.diff_threshold else "✅"
     lines: list[str] = [
         prefix_text
-        + "\t".join(
-            f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
-            for name, value in [
-                ("rel_diff", diff.rel_diff),
-                ("max_abs_diff", diff.max_abs_diff),
-                ("mean_abs_diff", diff.mean_abs_diff),
-            ]
-        ),
+        + f"{rel_diff_marker} rel_diff={diff.rel_diff}\t"
+        + f"max_abs_diff={diff.max_abs_diff}\t"
+        + f"mean_abs_diff={diff.mean_abs_diff}",
         f"max_abs_diff happens at coord={diff.max_diff_coord} with "
         f"baseline={diff.baseline_at_max} "
         f"target={diff.target_at_max}",

python/sglang/srt/debug_utils/dump_comparator.py

@@ -194,16 +194,12 @@ def _compute_and_print_diff(
     mean_abs_diff = raw_abs_diff.mean().item()
     rel_diff = _calc_rel_diff(x_target, x_baseline)
 
+    rel_diff_marker: str = "❌" if rel_diff > diff_threshold else "✅"
     print(
         prefix_text
-        + "\t".join(
-            f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
-            for name, value in [
-                ("rel_diff", rel_diff),
-                ("max_abs_diff", max_abs_diff),
-                ("mean_abs_diff", mean_abs_diff),
-            ]
-        )
+        + f"{rel_diff_marker} rel_diff={rel_diff}\t"
+        + f"max_abs_diff={max_abs_diff}\t"
+        + f"mean_abs_diff={mean_abs_diff}"
     )
 
     max_diff_coord = _argmax_coord(raw_abs_diff)